Strange behaviour when infering with NUTS

Hi,

I was asking why the model below is failing, I’ve been running it line by line and it works, but something strange is happening inside the traces. Im not an expert on Pyro, so maybe I forgot something. Thanks in advance.

Start by defining observed data:

import pyro
import pyro.distributions as dist
import torch 
from pyro.infer import HMC, MCMC, NUTS

observed = torch.Tensor([0, 2, 3, 1, 5, 1, 2, 0, 1, 2, 1, 0, 2, 6, 1, 0]).long()
hteam = torch.Tensor([0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3]).long()

Let’s define the model:

def bad_model(observed, hteam):
    mu = pyro.sample("mu", dist.Normal(0, 0.1))
    tau = pyro.sample("tau", dist.Gamma(0.1, 0.1))

    home = pyro.sample("home", dist.Normal(0, 0.1))

    att = pyro.sample("att_t", dist.Normal(mu, tau), [4])  # <-- From the 2nd iteration onwards, does not return a `torch.Tensor([a, b, c, d])` with 4 elements, just `torch.tensor(e)`
    theta = torch.exp(home + att[hteam])
    
    with pyro.plate("data", observed.size(0)):
        pyro.sample("obs", dist.Poisson(theta), obs=observed)

And then, when infering:

kernel = NUTS(bad_model, jit_compile=True, ignore_jit_warnings=True)
posterior = MCMC(kernel, num_samples=500, warmup_steps=500)
posterior.run(observed, hteam)

It raises:

Warmup:   0%|                                                                                                                                                          | 0/1000 [00:00, ?it/s]
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
/tmp/ipykernel_10074/1049478554.py in <module>
      3 kernel = NUTS(bad_model, jit_compile=True, ignore_jit_warnings=True)
      4 posterior = MCMC(kernel, num_samples=500, warmup_steps=500)
----> 5 posterior.run(observed, hteam)
.
.
.
/tmp/ipykernel_10074/1742051827.py in bad_model(observed, hteam)
      6 
      7     att = pyro.sample("att_t", dist.Normal(mu, tau), [4])  # <-- From the 2nd iteration, does not return a `torch.Tensor([a, b, c, d])` with 4 elements, just `torch.tensor(e)`
----> 8     theta = torch.exp(home + att[hteam])
      9 
     10     with pyro.plate("data", observed.size(0)):

IndexError: too many indices for tensor of dimension 0

Does anyone know what I am doing wrong?

what is the purpose of [4]?

The purpose of [4] is just sampling 4 elements from that distribution:

att = pyro.sample("att_t", dist.Normal(0, 0.1), [4])
att
# tensor([ 0.1478,  0.1022, -0.1932,  0.0375])

att[hteam]
# tensor([ 0.1478,  0.1478,  0.1478,  0.1478,  0.1022,  0.1022,  0.1022,  0.1022,
#         -0.1932, -0.1932, -0.1932, -0.1932,  0.0375,  0.0375,  0.0375,  0.0375])

ok but unfortunately that syntax is not supported.

i think you want something more like this:

def bad_model(observed, hteam):
    mu = pyro.sample("mu", dist.Normal(0, 0.1))
    tau = pyro.sample("tau", dist.Gamma(0.1, 0.1))

    home = pyro.sample("home", dist.Normal(0, 0.1))

    att = pyro.sample("att_t", dist.Normal(mu, tau).expand([4]).to_event(1))
    theta = torch.exp(home + att[hteam])
    
    with pyro.plate("data", observed.size(0)):
        pyro.sample("obs", dist.Poisson(theta), obs=observed)
1 Like

Wow, definetly it works after days of work! Could you explain what I was doing wrong?

well for one thing you were passing a random argument to sample that isn’t in the API. honestly i’m not sure where your [4] was being piped to

please refer to this tutorial for details about expand, tensor shapes, etc

Well, I see

I was using incorrectly that third argument (obs) from pyro.sample. From now, I’m using expand().to_event() to draw samples.

Thank you very much!!