Repository: pyro-ppl/numpyro

Issue:
icdf of beta distribution is not implemented
```python
>>> import numpyro.distributions as dist
>>> d = dist.Beta(1, 5)
>>> d.icdf(0.975)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/ayaka/venv/lib/python3.10/site-packages/numpyro/distributions/distribution.py", line 463, in icdf
    raise NotImplementedError
NotImplementedError
```

NumPyro version: 0.9.1


Discussion:
Hi @ayaka14732, please subscribe this upstream issue https://github.com/google/jax/issues/2399 Just curious, what is your usage case? If you don't need differentiation, then we can use scipy implementation for it.
> Just curious, what is your usage case? If you don't need differentiation, then we can use scipy implementation for it.

Just for learning, so I can use the scipy implementation. But for the NumPyro library, I think we should wait for it to be added in JAX, so users won't experience unexpected performance degradation.
I'd like to use the beta icdf in a `pyro.deterministic` statement in a model guide. Do I understand correctly that that is not possible by simply using the scipy function because gradient information is required? (I tried doing so and got errors that seemed to indicate this.) 

Any recommendations for what would currently be the best / easiest way to circumvent this problem? Can I somehow implement and provide the gradient information by myself?
I guess if beta icdf is only used at deterministic sites for prediction, then you can use `host_callback` as in [TruncatedDistribution](https://num.pyro.ai/en/stable/tutorials/truncated_distributions.html) tutorial. If it is used at some downstream observation sites, then we need grad rules for it. Please see [this tutorial](https://jax.readthedocs.io/en/latest/notebooks/Custom_derivative_rules_for_Python_code.html) for more information.
(In case somebody else finds this and looks for a solution, I opted to use the Kumaraswamy distribution instead, as recommended in the jax issue linked to above. It's very similar / in some cases identical to the beta distribution and has the icdf implemented.)
As suggested in https://github.com/google/jax/issues/2399, we can use `tfp.math.betaincinv` in the implementation of this cdf.
This can be added in the same way as #1475 .

Resolving patch:
diff --git a/numpyro/contrib/tfp/distributions.py b/numpyro/contrib/tfp/distributions.py
--- a/numpyro/contrib/tfp/distributions.py
+++ b/numpyro/contrib/tfp/distributions.py
@@ -177,7 +177,8 @@ def init(self, *args, **kwargs):
                 "scale": constraints.positive,
             }
         else:
-            if hasattr(numpyro_dist, tfd_class_name):
+            # Mixture distributions in tfp and numpyro are different.
+            if hasattr(numpyro_dist, tfd_class_name) and tfd_class_name != "Mixture":
                 numpyro_dist_class = getattr(numpyro_dist, tfd_class_name)
                 # resolve FooProbs/FooLogits namespaces
                 numpyro_dist_class = getattr(
diff --git a/numpyro/distributions/continuous.py b/numpyro/distributions/continuous.py
--- a/numpyro/distributions/continuous.py
+++ b/numpyro/distributions/continuous.py
@@ -34,7 +34,6 @@
 import jax.random as random
 from jax.scipy.linalg import cho_solve, solve_triangular
 from jax.scipy.special import (
-    betainc,
     betaln,
     expit,
     gammainc,
@@ -58,7 +57,10 @@
     SigmoidTransform,
 )
 from numpyro.distributions.util import (
+    betainc,
+    betaincinv,
     cholesky_of_inverse,
+    gammaincinv,
     is_prng_key,
     lazy_property,
     matrix_to_tril_vec,
@@ -187,6 +189,9 @@ def variance(self):
     def cdf(self, value):
         return betainc(self.concentration1, self.concentration0, value)
 
+    def icdf(self, q):
+        return betaincinv(self.concentration1, self.concentration0, q)
+
 
 class Cauchy(Distribution):
     arg_constraints = {"loc": constraints.real, "scale": constraints.positive}
@@ -363,6 +368,9 @@ def variance(self):
     def cdf(self, x):
         return gammainc(self.concentration, self.rate * x)
 
+    def icdf(self, q):
+        return gammaincinv(self.concentration, q) / self.rate
+
 
 class Chi2(Gamma):
     arg_constraints = {"df": constraints.positive}
@@ -1909,11 +1917,6 @@ def variance(self):
         return jnp.broadcast_to(var, self.batch_shape)
 
     def cdf(self, value):
-        try:
-            from tensorflow_probability.substrates.jax.math import betainc as betainc_fn
-        except ImportError:
-            from jax.scipy.special import betainc as betainc_fn
-
         # Ref: https://en.wikipedia.org/wiki/Student's_t-distribution#Related_distributions
         # X^2 ~ F(1, df) -> df / (df + X^2) ~ Beta(df/2, 0.5)
         scaled = (value - self.loc) / self.scale
@@ -1922,19 +1925,17 @@ def cdf(self, value):
 
         # when scaled < 0, returns 0.5 * Beta(df/2, 0.5).cdf(beta_value)
         # when scaled > 0, returns 1 - 0.5 * Beta(df/2, 0.5).cdf(beta_value)
-        scaled_sign_half = 0.5 * jnp.sign(scaled)
-        return (
-            0.5
-            + scaled_sign_half
-            - 0.5
-            * jnp.sign(scaled)
-            * betainc_fn(0.5 * jnp.asarray(self.df), 0.5, jnp.asarray(beta_value))
+        return 0.5 * (
+            1
+            + jnp.sign(scaled)
+            - jnp.sign(scaled) * betainc(0.5 * self.df, 0.5, beta_value)
         )
 
     def icdf(self, q):
-        # scipy.special.betaincinv is not avaiable yet in JAX
-        # upstream issue: https://github.com/google/jax/issues/2399
-        raise NotImplementedError
+        beta_value = betaincinv(0.5 * self.df, 0.5, 1 - jnp.abs(1 - 2 * q))
+        scaled_squared = self.df * (1 / beta_value - 1)
+        scaled = jnp.sign(q - 0.5) * jnp.sqrt(scaled_squared)
+        return scaled * self.scale + self.loc
 
 
 class Uniform(Distribution):
diff --git a/numpyro/distributions/util.py b/numpyro/distributions/util.py
--- a/numpyro/distributions/util.py
+++ b/numpyro/distributions/util.py
@@ -414,6 +414,37 @@ def clamp_probs(probs):
     return jnp.clip(probs, a_min=finfo.tiny, a_max=1.0 - finfo.eps)
 
 
+def betainc(a, b, x):
+    try:
+        from tensorflow_probability.substrates.jax.math import betainc as betainc_fn
+    except ImportError:
+        from jax.scipy.special import betainc as betainc_fn
+
+    return betainc_fn(jnp.asarray(a), jnp.asarray(b), jnp.asarray(x))
+
+
+def betaincinv(a, b, y):
+    try:
+        from tensorflow_probability.substrates.jax.math import special as tfp_special
+
+        return tfp_special.betaincinv(jnp.asarray(a), jnp.asarray(b), jnp.asarray(y))
+    except ImportError as e:
+        raise ImportError(
+            "Please install `tensorflow_probability>=0.18` for betaincinv."
+        ) from e
+
+
+def gammaincinv(a, y):
+    try:
+        from tensorflow_probability.substrates.jax import math as tfp_math
+
+        return tfp_math.igammainv(jnp.asarray(a), jnp.asarray(y))
+    except ImportError as e:
+        raise ImportError(
+            "Please install `tensorflow_probability>=0.18` for gammaincinv."
+        ) from e
+
+
 def is_identically_zero(x):
     """
     Check if argument is exactly the number zero. True for the number zero;


Test patch:
diff --git a/test/test_distributions.py b/test/test_distributions.py
--- a/test/test_distributions.py
+++ b/test/test_distributions.py
@@ -1187,6 +1187,7 @@ def test_log_prob(jax_dist, sp_dist, params, prepend_shape, jit):
     # TODO: add more complete pattern for Discrete.cdf
     CONTINUOUS + [T(dist.Poisson, 2.0), T(dist.Poisson, np.array([2.0, 3.0, 5.0]))],
 )
+@pytest.mark.filterwarnings("ignore:overflow encountered:RuntimeWarning")
 def test_cdf_and_icdf(jax_dist, sp_dist, params):
     d = jax_dist(*params)
     if d.event_dim > 0:
@@ -1194,8 +1195,8 @@ def test_cdf_and_icdf(jax_dist, sp_dist, params):
     samples = d.sample(key=random.PRNGKey(0), sample_shape=(100,))
     quantiles = random.uniform(random.PRNGKey(1), (100,) + d.shape())
     try:
+        rtol = 2e-3 if jax_dist in (dist.Gamma, dist.StudentT) else 1e-5
         if d.shape() == () and not d.is_discrete:
-            rtol = 1e-3 if jax_dist is dist.StudentT else 1e-5
             assert_allclose(
                 jax.vmap(jax.grad(d.cdf))(samples),
                 jnp.exp(d.log_prob(samples)),
@@ -1209,7 +1210,7 @@ def test_cdf_and_icdf(jax_dist, sp_dist, params):
                 rtol=rtol,
             )
         assert_allclose(d.cdf(d.icdf(quantiles)), quantiles, atol=1e-5, rtol=1e-5)
-        assert_allclose(d.icdf(d.cdf(samples)), samples, atol=1e-5, rtol=1e-5)
+        assert_allclose(d.icdf(d.cdf(samples)), samples, atol=1e-5, rtol=rtol)
     except NotImplementedError:
         pass
 
@@ -1223,7 +1224,7 @@ def test_cdf_and_icdf(jax_dist, sp_dist, params):
         assert_allclose(actual_cdf, expected_cdf, atol=1e-5, rtol=1e-5)
         actual_icdf = d.icdf(quantiles)
         expected_icdf = sp_dist.ppf(quantiles)
-        assert_allclose(actual_icdf, expected_icdf, atol=1e-5, rtol=1e-4)
+        assert_allclose(actual_icdf, expected_icdf, atol=1e-4, rtol=1e-4)
     except NotImplementedError:
         pass