Wrap distributions in scipy.stats for use with JAX about numpyro

[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 - Just add your name to the distribution that you are working on so that we don't end up doing double work! :)

[fehiepsi]

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

[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?

[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 - google/jax#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.

[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.

[neerajprad]

No hurry, please take your time.