Laplace Approximation in PyTorch

laplace_approximation.py

"""
Laplace approximation of a Beta distribution.
"""
import matplotlib.pyplot as plt
import torch

x = torch.linspace(0, 1, 200)
p = torch.distributions.Beta(2, 5)


def pdf(index_points, distribution):
    return torch.exp(distribution.log_prob(index_points))


mode = x[torch.argmax(pdf(x, p))]
plt.plot(x, pdf(x, p))
plt.scatter(mode, pdf(mode, p))

# noinspection PyTypeChecker
hessian = -torch.autograd.functional.hessian(p.log_prob, mode)
q = torch.distributions.Normal(mode, hessian ** -1)
plt.plot(x, pdf(x, q))
plt.title("Laplace Approximation")
plt.show()