Thanks to these two fantastic resources (I hope they help someone else), I understand how to instantiate batched Normal and MultivariateNormal objects:
https://ericmjl.github.io/blog/2019/5/29/reasoning-about-shapes-and-probability-distributions/
I also understand the output shapes of sampling the batched distributions. However, I am somewhat lost as to the input shape requirements of the log_prob method.
Here’s an example implementation that I just got right through random permutation until it worked… I still have no idea how I did it:
# 1D mesh. Each point on the mesh will be the mean of a different distribution
mesh = torch.arange(0, 10)
scale = torch.ones_like(mesh, dtype=torch.float32)
scale += torch.normal(mean=0, std=torch.arange(0.1, 0.2, 0.01))
x_range = torch.arange(-5, 15, 0.1)
results = torch.exp(gaussians.log_prob(x_range))
for i, res in enumerate(results.T):
sns.lineplot(x=x_range, y=res)
Okay, so I need to unsqueeze the input and it works fine…
results = torch.exp(gaussians.log_prob(x_range.unsqueeze(1)))
The 0th dimension is the number of samples. But what is this new dimension?
And what if I have a batch of samples:
x_range = torch.arange(-5, 15, 0.1)
x_range = torch.stack([x_range, x_range + 1]) # <--- Two batches
results = torch.exp(gaussians.log_prob(x_range))
for i, res in enumerate(results.T):
sns.lineplot(x=x_range, y=res)
Again I get
So I unsqueeze the last dimension again. Now I’m dealing with three-dimensional input and I have no idea what two of the dimension are, nor the dimensions of the output.
Here is what the docs say:
Nothing…