Probabilistic programming interface

Is there a plan of adding probabilistic programming interface, such as rlax.distributions or tfp.

Right now, numpyro is the only JAX-powered library for this. I could add an example of how to combine dm-haiku and numpyro if you want.

[trevorcai]

To be honest, I haven’t yet explored probabilistic programming (past using tfp.distributions for other projects in the past) in sufficient depth to give a strong answer here. I like declaring things out of scope where sensible; if Haiku composes with numpyro or other pure-functional libraries, we’ll just declare this out of scope!

My limited experience with Haiku + RLAX and Sonnet + TFP has given me no indicator that Haiku won’t compose with probabilistic programming libraries; I’ve previously written VAEs in Haiku with no issues.

That said, I’d love to see an example combining Haiku & numpyro, though it could be tricky for us to merge internally.

Any thoughts, @ibab ?

[ibab]

Here’s an example of how to replace the stax networks in the numpyro VAE example with Haiku: https://github.com/ibab/numpyro/commit/076bb2a15026734a41abf4eb39ce384ab68c297f We could add a Haiku to stax converter function to Haiku to streamline use cases like this.