Am I using HMC/NUTS/MCMC correctly?

16

As an alternate approach, I think you could also define a distribution class with log_prob that gives you unnormalized density and use it in your model. Sample code:

import torch
from torch.distributions import constraints

import pyro
from pyro.distributions import TorchDistribution
from pyro.infer.mcmc import MCMC, NUTS


class UnnormalizedDist(TorchDistribution):
    support = constraints.real
    has_rsample = False

    def __init__(self):
        self.mu = torch.tensor(2.)
        self.std = torch.tensor(3.)
        super(TorchDistribution, self).__init__()

    # HACK: Only used for model initialization.
    def sample(self, sample_shape=torch.Size()):
        return torch.tensor(1.)

    def log_prob(self, value):
        return -((value - self.mu) ** 2) / (2 * self.std ** 2)


def model():
    pyro.sample("y", UnnormalizedDist())


if __name__ == "__main__":
    kernel = NUTS(model, jit_compile=True, ignore_jit_warnings=True)
    mcmc = MCMC(kernel=kernel, num_samples=10000, warmup_steps=200).run()
    marginal = mcmc.marginal(["y"]).empirical["y"]
    print(marginal.mean)
    print(marginal.variance)

As a note for us, we can facilitate such use cases by not sampling from the model for the initial setup in case an initial_trace is provided to the kernel, thereby avoiding the constant sampling hack above.

2 Likes