Numpyro - enumeration of discrete variables

Just getting into pyro/numpyro - looks great!

I’d like to do NUTS with automatic enumeration of discrete variables. I know that pyro supports this, but I haven’t been able to find documentation on whether numpyro does? My guess is not, and that I should be using pyro, but wanted to check.

Thanks!

1 Like

Hi @jpm, we are working on supporting models with discrete latent variables in NumPyro. Currently, it is not supported.