Example of plotting MV Norm distribution in Pytorch

plot_multivariate_normal.py

import torch
from torch.distributions.normal import Normal
from torch.distributions.multivariate_normal import MultivariateNormal
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

d = Normal(0.0, 1.0)
x = torch.linspace(-4, 4.0, 50)
y = torch.exp(d.log_prob(x))
plt.plot(x, y)
plt.show()

d = MultivariateNormal(torch.zeros(2), torch.eye(2))
x_, y_, = torch.linspace(-3, 3, 100), torch.linspace(-3, 3, 100)
x, y = torch.meshgrid([x_, y_])
z = d.log_prob(torch.stack((x, y)).T)
z = torch.exp(z)

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.plot_surface(x.cpu().numpy(), y.cpu().numpy(), z.cpu().numpy())
plt.show()