Repository: pyro-ppl/numpyro

Issue:
inf's with TruncatedNormal
I've seen the discussion in #1184 and #1185, but I'm still seeing this issue with numpyro v0.10.1. Here's a MWE, comparing to scipy's `scipy.stats.truncnorm` implementation:
```python
import numpyro
numpyro.enable_x64()
import numpyro.distributions as dist
import jax.numpy as jnp
from scipy.stats import truncnorm

loc = 1.35
scale = jnp.geomspace(0.01, 1, 10)
low, high = (-20, -1.0)
a, b = (low - loc) / scale, (high - loc) / scale
x = -15.

scipy_val = truncnorm.logpdf(x, loc=loc, scale=scale, a=a, b=b)
numpyro_val = dist.TruncatedNormal(loc, scale, low=low, high=high).log_prob(x)
```
(arbitrary values chosen to get into the tail)

Comparing the output values:
```python
>>> scipy_val
array([-1.30898994e+06, -4.70421167e+05, -1.69055833e+05, -6.07514028e+04,
       -2.18294637e+04, -7.84229794e+03, -2.81622225e+03, -1.01058785e+03,
       -3.62302368e+02, -1.29911728e+02])
>>> numpyro_val
DeviceArray([            inf,             inf,             inf,
                         inf, -21829.46367826,  -7842.29793866,
              -2816.22224529,  -1010.58784742,   -362.30236837,
               -129.91172764], dtype=float64)
```

It's possible to avoid this by special-casing the truncated normal distribution, as I [recently implemented in Jax](https://github.com/google/jax/pull/12646) -- it would be great to have this in numpyro as well.
```python
from jax.scipy.stats import truncnorm as jax_truncnorm
jax_val = jax_truncnorm.logpdf(x, loc=loc, scale=scale, a=a, b=b)
print(jax_val)

DeviceArray([-1.30898994e+06, -4.70421167e+05, -1.69055833e+05,
             -6.07514028e+04, -2.18294637e+04, -7.84229794e+03,
             -2.81622225e+03, -1.01058785e+03, -3.62302368e+02,
             -1.29911728e+02], dtype=float64)
```

Would you consider a PR to special-case `TruncatedNormal`? I'm not familiar with the numpyro codebase but have just started using it and am loving it - thanks for the work and maintenance on this project!


Discussion:
> consider a PR to special-case TruncatedNormal?

That would be awesome, @adrn!!! Regarding numpyro codebase, I think you can add a method `log_cdf` in `Normal` and use it to compute sort of [_log_diff_tail_probs](https://github.com/pyro-ppl/numpyro/blob/e9af0c6facc38ba8421c27b27abedbcbae5ff937/numpyro/distributions/truncated.py#L238-L250) which computes log of `self._tail_prob_at_high - self._tail_prob_at_low`. You can add a method `_log_diff_tail_probs` to `TwoSidedTruncatedDistribution` class that uses more numerical stable formula if `log_cdf` method is available in `base_dist`. Please let me know if this is not clear.

By the way, do you have any idea to address similar issue for [sample](https://github.com/pyro-ppl/numpyro/blob/e9af0c6facc38ba8421c27b27abedbcbae5ff937/numpyro/distributions/truncated.py#L269) method? Maybe it is enough to clip values there?
Hi,
I'm wondering if there are any updates on this? 
I'm currently running into inf's when trying to sample from `dist.TruncatedNormal()`.

Thanks in advance!
I also had similar issues with inf's and [implemented ](https://gist.github.com/nikmich1/353f6ddc27cbceb7d3e9cf0fe2dee096)the steps as suggested by fehiepsi.

@fehiepsi is this potentially interesting for a PR? In case it is, would it be ok to simply call jax.scipy.stats.norm.logcdf for the new log_cdf method in Normal, or does it need to be implemented from scratch?
> would it be ok to simply call jax.scipy.stats.norm.logcdf

Yes, that would be nice. Thanks for working on this issue, your solution looks great to me. Could you also address the issue at `sample` method by clip [the prob value](https://github.com/pyro-ppl/numpyro/blob/e9af0c6facc38ba8421c27b27abedbcbae5ff937/numpyro/distributions/truncated.py#L269) to the open (0, 1) interval. You can use [clamp_probs](https://github.com/pyro-ppl/numpyro/blob/e9af0c6facc38ba8421c27b27abedbcbae5ff937/numpyro/distributions/util.py#L412) utility for it.

Resolving patch:
diff --git a/numpyro/distributions/continuous.py b/numpyro/distributions/continuous.py
--- a/numpyro/distributions/continuous.py
+++ b/numpyro/distributions/continuous.py
@@ -47,6 +47,7 @@
     xlog1py,
     xlogy,
 )
+from jax.scipy.stats import norm as jax_norm
 
 from numpyro.distributions import constraints
 from numpyro.distributions.discrete import _to_logits_bernoulli
@@ -2077,6 +2078,9 @@ def cdf(self, value):
         scaled = (value - self.loc) / self.scale
         return ndtr(scaled)
 
+    def log_cdf(self, value):
+        return jax_norm.logcdf(value, loc=self.loc, scale=self.scale)
+
     def icdf(self, q):
         return self.loc + self.scale * ndtri(q)
 
diff --git a/numpyro/distributions/truncated.py b/numpyro/distributions/truncated.py
--- a/numpyro/distributions/truncated.py
+++ b/numpyro/distributions/truncated.py
@@ -18,6 +18,7 @@
 )
 from numpyro.distributions.distribution import Distribution
 from numpyro.distributions.util import (
+    clamp_probs,
     is_prng_key,
     lazy_property,
     promote_shapes,
@@ -249,6 +250,23 @@ def _tail_prob_at_high(self):
         sign = jnp.where(loc >= self.low, 1.0, -1.0)
         return self.base_dist.cdf(loc - sign * (loc - self.high))
 
+    @lazy_property
+    def _log_diff_tail_probs(self):
+        # use log_cdf method, if available, to avoid inf's in log_prob
+        # fall back to cdf, if log_cdf not available
+        log_cdf = getattr(self.base_dist, "log_cdf", None)
+        if callable(log_cdf):
+            return logsumexp(
+                a=jnp.stack([log_cdf(self.high), log_cdf(self.low)], axis=-1),
+                axis=-1,
+                b=jnp.array([1, -1]),  # subtract low from high
+            )
+
+        else:
+            loc = self.base_dist.loc
+            sign = jnp.where(loc >= self.low, 1.0, -1.0)
+            return jnp.log(sign * (self._tail_prob_at_high - self._tail_prob_at_low))
+
     def sample(self, key, sample_shape=()):
         assert is_prng_key(key)
         dtype = jnp.result_type(float)
@@ -266,7 +284,7 @@ def sample(self, key, sample_shape=()):
         loc = self.base_dist.loc
         sign = jnp.where(loc >= self.low, 1.0, -1.0)
         return (1 - sign) * loc + sign * self.base_dist.icdf(
-            (1 - u) * self._tail_prob_at_low + u * self._tail_prob_at_high
+            clamp_probs((1 - u) * self._tail_prob_at_low + u * self._tail_prob_at_high)
         )
 
     @validate_sample
@@ -276,10 +294,7 @@ def log_prob(self, value):
         #   cdf(high) - cdf(low) = as-is
         # if low > loc
         #   cdf(high) - cdf(low) = cdf(2 * loc - low) - cdf(2 * loc - high)
-        sign = jnp.where(self.base_dist.loc >= self.low, 1.0, -1.0)
-        return self.base_dist.log_prob(value) - jnp.log(
-            sign * (self._tail_prob_at_high - self._tail_prob_at_low)
-        )
+        return self.base_dist.log_prob(value) - self._log_diff_tail_probs
 
     def tree_flatten(self):
         base_flatten, base_aux = self.base_dist.tree_flatten()


Test patch:
diff --git a/test/test_distributions.py b/test/test_distributions.py
--- a/test/test_distributions.py
+++ b/test/test_distributions.py
@@ -19,6 +19,7 @@
 import jax.numpy as jnp
 import jax.random as random
 from jax.scipy.special import expit, logsumexp
+from jax.scipy.stats import norm as jax_norm, truncnorm as jax_truncnorm
 from jax.tree_util import tree_map
 
 import numpyro.distributions as dist
@@ -2758,3 +2759,46 @@ def f(x):
     x = dist.Multinomial(10, probs).sample(key)
     y = jax.jit(f)(x)
     assert_allclose(x, y, rtol=1e-6)
+
+
+def test_normal_log_cdf():
+    # test if log_cdf method agrees with jax.scipy.stats.norm.logcdf
+    # and if exp(log_cdf) agrees with cdf
+    loc = jnp.array([[0.0, -10.0, 20.0]])
+    scale = jnp.array([[1, 5, 7]])
+    values = jnp.linspace(-5, 5, 100).reshape(-1, 1)
+    numpyro_log_cdf = dist.Normal(loc=loc, scale=scale).log_cdf(values)
+    numpyro_cdf = dist.Normal(loc=loc, scale=scale).cdf(values)
+    jax_log_cdf = jax_norm.logcdf(loc=loc, scale=scale, x=values)
+    assert_allclose(numpyro_log_cdf, jax_log_cdf)
+    assert_allclose(jnp.exp(numpyro_log_cdf), numpyro_cdf, rtol=1e-6)
+
+
+@pytest.mark.parametrize(
+    "value",
+    [
+        -15.0,
+        jnp.array([[-15.0], [-10.0], [-5.0]]),
+        jnp.array([[[-15.0], [-10.0], [-5.0]], [[-14.0], [-9.0], [-4.0]]]),
+    ],
+)
+def test_truncated_normal_log_prob_in_tail(value):
+    # define set of distributions truncated in tail of distribution
+    loc = 1.35
+    scale = jnp.geomspace(0.01, 1, 10)
+    low, high = (-20, -1.0)
+    a, b = (low - loc) / scale, (high - loc) / scale  # rescale for jax input
+
+    numpyro_log_prob = dist.TruncatedNormal(loc, scale, low=low, high=high).log_prob(
+        value
+    )
+    jax_log_prob = jax_truncnorm.logpdf(value, loc=loc, scale=scale, a=a, b=b)
+    assert_allclose(numpyro_log_prob, jax_log_prob, rtol=1e-06)
+
+
+def test_sample_truncated_normal_in_tail():
+    # test, if samples from distributions truncated in
+    # tail of distribution returns any inf's
+    tail_dist = dist.TruncatedNormal(loc=0, scale=1, low=-16, high=-15)
+    samples = tail_dist.sample(random.PRNGKey(0), sample_shape=(10_000,))
+    assert ~jnp.isinf(samples).any()