MCMC inference for Poisson GPLVM

I would like to code up a GPLVM with Poisson emissions and perform MCMC inference (VI is also fine, if that is easier). The model is

x_n    ~ N(0, I)
f_j    ~ GP(0, K)
Y_{nj} ~ Poisson(exp(f_j(x_n)))

In other words, Y is an N x J count matrix, and its J columns are Gaussian process-distributed maps from a latent variable X, an N x D matrix.

Following the docstring example here, I have tried to code this up:

import pyro
import pyro.contrib.gp as gp
from   pyro.distributions import Normal, Poisson
from   pyro.infer.mcmc import MCMC, NUTS
import torch


N = 500  # Samples.
J = 50   # Features.
D = 2    # Latent dimension.

Y      = torch.distributions.Poisson(100).sample((N, J))
kernel = gp.kernels.RBF(input_dim=D)
mu_x   = torch.zeros(D)
cov_x  = torch.eye(D)
mu_f   = torch.zeros(N)


def model(data):
    X     = pyro.sample('X', Normal(mu_x, cov_x))
    K_x   = kernel(X)
    F     = pyro.sample('F', Normal(mu_f, K_x))
    theta = torch.exp(F)
    return pyro.sample('Y', Poisson(theta), obs=data)


nuts_kernel = NUTS(model)
mcmc = MCMC(nuts_kernel, num_samples=500, warmup_steps=100)
mcmc.run(data=Y)
mcmc.get_samples()['X'].mean(0)

However, I get the error:

RuntimeError: The size of tensor a (500) must match the size of tensor b (2) at non-singleton dimension 1
Trace Shapes:
 Param Sites:
  lengthscale
     variance
Sample Sites:
       X dist 2 2 |
        value 2 2 |

It seems clear that Y and X do not have the same first dimension, but I don’t know how to sample N random variables x_n using Pyro.

Thanks for any help.

please refer to the tensor shapes tutorial

1 Like

Thank you. This is helpful. While my code no longer crashes, it appears to do nothing:

...
_SMALL = 0.00001

def model(Y):
    # `sample` should produce i.i.d. samples X = {x_n}.
    X = MultivariateNormal(mu_x, cov_x).sample(torch.Size([N]))
    # `+ _SMALL` to ensure PSD.
    K_x = kernel(X) + _SMALL * torch.eye(N)
    # `expand` specifies that `F` are conditionally independent given `K_x`.
    F = MultivariateNormal(mu_f, K_x).expand(torch.Size([J])).sample().T
    theta = torch.exp(F)
    assert(theta.shape == Y.shape)
    return pyro.sample('obs', Poisson(theta), obs=Y)

nuts_kernel = NUTS(model, adapt_step_size=True)
mcmc = MCMC(nuts_kernel, num_samples=500, warmup_steps=100)
mcmc.run(Y=Y)
print(mcmc.get_samples())  # Empty `dict`.

you need to use pyro.sample statements and the like if you want pyro to know about the randomness in your model. (pyro doesn’t “see” the randomness you’re injecting with your MVN distribution, for example). please refer to the intro tutorials.