right, but the original tensor has a shape and I need to preserve the shape which would exist if I were to do…
x = torch.ones((2, 3, 4))
x.mean(dim=0)

which would be a (3,4) tensor of the means along dimension 0. I have found that I could do this by summing the elements along the dimension and then somehow dividing by the count of nonzero elements along the same dimension…
But then I realized my problem is a little bit harder than that because I am dealing with one tensor of values and then one probability distribution tensor with the same shape, which gives the probabilities of those values…
x = torch.rand((2, 3, 4))
y = torch.distributions.Normal(torch.randn((2, 3, 4), torch.randn((2, 3, 4)))
pr = y.log_prob(x)
...

at this point I need to get the mean of dimension 0 of the values in y which are 0 in x but they are not 0 in y because of the log_prob call… very confusing about how to go about it