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.