Numpyro: Wrap distributions in scipy.stats for use with JAX

[neerajprad]

To make the internals of the sampler and logpdf methods visible to the tracer and compatible with jax, we need to:

• Write the operations in terms of jax.numpy operations.
• We need to use jax's count based random number generator (via PRNGKey), and not the default used by scipy distributions which is numpy's globally mutable mtrand. Currently, the PRNGKey is being passed in through the random_state kwarg to the distribution's _rvs method.

Currently, only the normal distribution is wrapped. We can similarly wrap the following distributions (we may need to wrap over the samplers / rewrite the logic using jax operations):

Continuous

• [x] beta @fehiepsi
• [x] cauchy @fehiepsi
• [x] dirichlet @fehiepsi
• [x] exponential @fehiepsi
• [x] uniform
• [x] gamma @fehiepsi
• [x] lognormal @fehiepsi
• [x] normal
• [x] student-t @fehiepsi

Discrete

• [x] bernoulli @neerajprad
• [x] categorical @neerajprad
• [x] binomial
• [ ] poisson
• [x] multinomial

[fehiepsi]

@neerajprad Have you worked on this? I can't quickly add these distributions in case it is not overlapping with what you are doing. :)

[neerajprad]

@neerajprad Have you worked on this? I can't quickly add these distributions in case it is not overlapping with what you are doing. :)

Feel free to work on this. I only took a brief look. I think the blocker is that we need to add a gamma sampler to jax first before we can support distributions like dirichlet and beta. What do you think?

[neerajprad]

@fehiepsi - Just add your name to the distribution that you are working on so that we don't end up doing double work! :)

[fehiepsi]

we need to add a gamma sampler to jax first before we can support distributions like dirichlet and beta.

Agree! I am mostly interested in seeing performance on hmm so I would like to support dirichlet distribution. Will try to implement a gamma sampler first (mainly to learn). :)

[fehiepsi]

Hmm, looks like it is more complicated than I thought. Not because of understanding algorithms (Fritz already made it clear in his PR to pytorch), but because of my lacking knowledge on how jax/lax works (e.g. how to take gradient with shape parameters, how to implement if/else/while logics,...). I think that we'd need help from JAX devs on this.

[neerajprad]

I believe you'll need to define the sampler as a primitive operation and register its gradient separately, and not try to make the sampler itself differentiable. It is possible to do that but the documentation is a bit lacking - https://github.com/google/jax/issues/116. I think all we'll need is digamma to specify the gradient function which is already implemented in xla. What do you think? I can take a stab at it next week too.

[fehiepsi]

Yeah, dwelling into autograd documentation, then searching for similar pattern in jax.lax, I come up with the following template, which works for some simple functions. ;)

from jax.interpreters import ad
from jax.lax import standard_unop, _float

def standard_gamma(x):
    # implement forward
    return x

def _standard_gamma_jvp_rule(g, ans, x):
    # implement backward
    return g * x

standard_gamma_p = standard_unop(_float, 'standard_gamma')
ad.defjvp2(standard_gamma_p, _standard_gamma_jvp_rule)

Edit: the above is not true (it is just applicable for XLA-compatible primitives), it seems that we need to define jvp_rule (forward-mode), transpose_rule (reverse-mode), batch_rule (vmap),... for custom primitives. I have implemented the sampler but the rules are still tricky for me.

[neerajprad]

@fehiepsi - We can ask the jax folks if this is already on their roadmap. In the meantime, feel free to add a PR for the sampler. For HMC at least, we just need the logpdf methods which can be wrapped over easily, and even with just a non-reparametrized sampler we can experiment with other algorithms.

[fehiepsi]

Yes, I'll make a PR soon to not block us from adding more distributions.

[neerajprad]

No hurry, please take your time.