Constructing batched distributions from data in pytorch and sampling from them

batched_normal_distributions.py

import torch
from torch.distributions.normal import Normal
from torch.distributions.multivariate_normal import MultivariateNormal


"""
Example of computing c batched Normal and Multivariate distributions from data
and sampling batches from them
"""

c = 3
n = 10

x = torch.randn(n, c)

loc = x.mean(dim=0)
scale = x.std(dim=0)

dist = Normal(loc, scale)
samples = dist.sample((n,))
print(dist.mean)
print(dist.stddev)
print(samples.shape)

c = 4
n = 3
d = 2

x = torch.randn(n, c, d)
loc = x.mean(0)
delta = x - loc.unsqueeze(0)
covariance_matrix = torch.matmul(delta.permute(1, 2, 0), delta.permute(1, 0, 2)) / (n - 1)
dist = MultivariateNormal(loc, covariance_matrix)
samples = dist.sample((n, ))
print(samples.shape)