Discrete uniform from categorical leads to error during inference

Misc.
#1

I am experiencing some difficulties with a simple model from Bayesian Methods for Hackers:

def model(obs_data):
    n = obs_data.shape[0]

    # hyperparams
    alpha = 1 / obs_data.mean()

    # 0 ...(λ1)...τ...(λ2)... n #

    with pyro.plate('lam_t', 2):
        lam_t = pyro.sample("lam", dist.Exponential(alpha))

    tau = 1 + pyro.sample("tau", dist.Categorical(T.ones(n-1) / (n-1)))
    
    lam = T.cat([
        lam_t[0].expand(tau),
        lam_t[1].expand(n - tau)
    ])

    with pyro.plate("data", len(obs_data)) as idx:
        obs = pyro.sample("obs", dist.Poisson(lam), obs=obs_data)

# Inference
kernel = NUTS(
    model, jit_compile=True, ignore_jit_warnings=True, max_tree_depth=5
)
posterior = MCMC(
    kernel, num_samples=2000, warmup_steps=500
)
posterior.run(obs_data)

During inference, the following exception is raised:

     14     lam = T.cat([
---> 15         lam_t[0].expand(tau),
     16         lam_t[1].expand(n - tau)
     17     ])
TypeError: 
expand(): argument 'size' (position 1) must be tuple of ints, not Tensor

It appears that the problem is fixed if I sample tau from dist.Uniform(0, 1) and make the appropriate transformation.

Therefore, my question reduces to: is there a more subtle difference (at least regarding the size of the sampled data) between

tau = 1 + pyro.sample("tau", dist.Categorical(T.ones(n-1) / (n-1)))

and

tau = 1 + ((n-1) * pyro.sample("tau", dist.Uniform(0, 1))).long()

Thank you!

#2

Hi @AlexD, the difference is Categorical is a discrete distribution and Uniform is a continuous distribution. NUTS algorithm only works for continuous latent variables so if you use Categorical distribution, you need to either:

  • use NUTS within Gibbs (this is the default in PyMC3 I guess)
  • or marginalize the discrete latent variable (this is the default in Pyro, see more in enumerate tutorial on how to write code to support marginalization).

If you want to use NUTS within Gibbs, you can try DiscreteHMCGibbs in NumPyro.

1 Like