Hi! I have this stochstic function:

`def prior(self, batch_size, device): concentration = torch.tensor(3.4, device=device) translation_x = (pyro.sample('translation_x', Beta(concentration, concentration).expand(batch_size)) - 0.5) * 6.0 translation_y = (pyro.sample('translation_y', Beta(concentration, concentration).expand(batch_size)) - 0.5) * 6.0 return torch.stack((translation_x, translation_y))`

With poutine.trace(prior).get_trace([10], ‘cpu’).log_prob_sum() I can calculate the pdf of the whole batch. However, I want to calculate the pdf of the individual samples in the batch and not all of them combined. What is the best way to do that?