inverse CDF for Gamma/Beta distributions

[alexxthiery]

Hi!

I was wondering whether there is a way to get an equivalent of the Gamma/Beta (and other standard distributions) quantile function: scipy.stats.gamma.ppf
https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.gamma.html
https://docs.scipy.org/doc/scipy/reference/generated/scipy.special.gammaincinv.html

The inverse cumulative function is already implemented for the Gaussian distribution, but not for the Gamma / Beta distributions.
Although this may not be the best way to do so [1], this could be useful for using the "reparametrization trick" for Gamma / Beta distributions.

Thanks!
Alex.

[1]: Figurnov, Mikhail, Shakir Mohamed, and Andriy Mnih. "Implicit reparameterization gradients." Advances in neural information processing systems 31 (2018).

[jakevdp]

Hi - thanks for the request! These kinds of features are generally added on an as-needed basis. If you're interested in putting together a pull request, we'd be happy to review it! But we don't have anyone actively working on expanding the jax.scipy.stats and jax.scipy.special API coverage at the moment.

[VARUN3WARE]

was trying to work on this issue for few days
I'd really appreciate your thoughts on it..

@jakevdp

[jakevdp]

Hi @VARUN3WARE – implementing this would require two steps:

1. Implement jax.scipy.special.gammaincinv. This will be tricky, as there's no HLO primitive for this function: scipy implements this via cephes, which typically relies on a fixed series expansion. In similar cases it's often possible to port that implementation to JAX (search "cephes" in the JAX source for some examples).
2. Once that's done, scipy.stats.gamma.ppf can be trivially implemented in terms of this new function.

[VARUN3WARE]

finally got time for fixing this issue

waiting for your comments on it

also i added some tests for checking if implementation was correct or not..

@jakevdp