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.

Hi @fehiepsi, has this been implemented in numpyro yet, specifically for SVI? I found Support enumeration in SVI · Issue #741 · pyro-ppl/numpyro · GitHub but don’t see a conclusion. I am working on a problem with discrete latent variables, but the size of the data and complexity of the model make MC solutions impractical. Thanks in advance!

Hi @tjm, it is not available yet. If you need this, I think you can port pyro funsor’s TraceEnum_ELBO over, which is mostly independent of the backend.