Repository: google/jax
Issue: Higher order derivatives of norm.logcdf seem to have numerical problems at low input values

Hi there,

I think I might have found a problem while investigating https://github.com/google/jax/issues/475 (though still checking!).

It looks like second derivatives of the logcdf and higher have problems for small input values (<-10) -- see reproducible code below:

```python
import os
os.environ['XLA_FLAGS'] = '--xla_cpu_enable_fast_math=false'
import jax.numpy as np
from jax.scipy.stats import norm
from jax import grad, vmap

x_vals = np.linspace(-30, 20, 100)
first_grad = vmap(grad(norm.logcdf))(x_vals)
second_grad = vmap(grad(grad(norm.logcdf)))(x_vals)
third_grad = vmap(grad(grad(grad(norm.logcdf))))(x_vals)

count_nans = lambda x: int(np.sum(np.isnan(x)))

print(f'Found {count_nans(first_grad)} NaNs in first derivative.')
print(f'Found {count_nans(second_grad)} NaNs in second derivative.')
print(f'Found {count_nans(third_grad)} NaNs in third derivative.')
```

Output is:

```
Found 0 NaNs in first derivative.
Found 40 NaNs in second derivative.
Found 42 NaNs in third derivative.
```

I need derivatives up to third order and have been able to get around the issue by defining the ratio:

```python
def pdf_cdf_ratio(x):
    return np.exp(norm.logpdf(x) - norm.logcdf(x))
```

Which is the gradient of the logcdf (subtracting and then exponentiating is much more stable than computing the ratio directly). You can express the second and third derivatives in terms of the ratio. Here's the full code:

```python
def pdf_cdf_ratio(x):
    return np.exp(norm.logpdf(x) - norm.logcdf(x))

def pdf_cdf_ratio_grad(x):

    g = pdf_cdf_ratio(x)
    return -g * (x + g)

def pdf_cdf_ratio_grad_grad(x):

    g = pdf_cdf_ratio(x)
    gp = pdf_cdf_ratio_grad(x)

    return -gp * (x + g) - g * (1 + gp)
```

With these alternatives, I get:

```python
first_grad_alt = pdf_cdf_ratio(x_vals)
second_grad_alt = pdf_cdf_ratio_grad(x_vals)
third_grad_alt = pdf_cdf_ratio_grad_grad(x_vals)

count_nans = lambda x: int(np.sum(np.isnan(x)))

print(f'Found {count_nans(first_grad_alt)} NaNs in first derivative.')
print(f'Found {count_nans(second_grad_alt)} NaNs in second derivative.')
print(f'Found {count_nans(third_grad_alt)} NaNs in third derivative.')
```

```
Found 0 NaNs in first derivative.
Found 0 NaNs in second derivative.
Found 0 NaNs in third derivative.
```

Mind you, plotting the functions, although it has no `nan`s, it looks like the third gradient gets pretty wonky below values of -10, so maybe there's a better alternative.

Anyway, I'm not sure what to do to fix this with autodiff, and whether that's even possible, but thought I'd raise this as a possible problem!

Changed files:
tests/lax_scipy_test.py, jax/scipy/special.py

diff --git a/jax/scipy/special.py b/jax/scipy/special.py
index 4a425e043..c2e0ef692 100644
--- a/jax/scipy/special.py
+++ b/jax/scipy/special.py
@@ -216,7 +216,7 @@ def _ndtr(x):
                       lax.select(lax.gt(w, dtype(0.)),
                                       dtype(2.) - lax.erfc(z),
                                       lax.erfc(z)))
-  return 0.5 * y
+  return dtype(0.5) * y
 
 
 def ndtri(p):
@@ -355,6 +355,7 @@ def _ndtri(p):
   return x_nan_replaced
 
 
+@custom_transforms
 def log_ndtr(x, series_order=3):
   r"""Log Normal distribution function.
 
@@ -488,3 +489,14 @@ def _log_ndtr_asymptotic_series(x, series_order):
 def _double_factorial(n):
   """The double factorial function for small Python integer `n`."""
   return onp.prod(onp.arange(n, 1, -2))
+
+
+_norm_logpdf_constant = onp.log(onp.sqrt(2 * onp.pi))
+
+def _norm_logpdf(x):
+  neg_half = _constant_like(x, -0.5)
+  log_normalizer = _constant_like(x, _norm_logpdf_constant)
+  return lax.sub(lax.mul(neg_half, lax.square(x)), log_normalizer)
+
+defjvp(log_ndtr,
+       lambda g, ans, x: lax.mul(g, lax.exp(lax.sub(_norm_logpdf(x), ans))))
diff --git a/tests/lax_scipy_test.py b/tests/lax_scipy_test.py
index f1be13310..daf1001a5 100644
--- a/tests/lax_scipy_test.py
+++ b/tests/lax_scipy_test.py
@@ -65,8 +65,8 @@ JAX_SPECIAL_FUNCTION_RECORDS = [
     # TODO: gammaln has slightly high error.
     op_record("gammaln", 1, float_dtypes, jtu.rand_positive(), False),
     op_record("logit", 1, float_dtypes, jtu.rand_small_positive(), False),
-    op_record("log_ndtr", 1, float_dtypes, jtu.rand_small(), True),
-    op_record("ndtri", 1, float_dtypes, jtu.rand_uniform(0., 1.), True),
+    op_record("log_ndtr", 1, float_dtypes, jtu.rand_default(), True),
+    op_record("ndtri", 1, float_dtypes, jtu.rand_uniform(0.05, 0.95), True),
     op_record("ndtr", 1, float_dtypes, jtu.rand_default(), True),
 ]
 
@@ -119,7 +119,7 @@ class LaxBackedScipyTests(jtu.JaxTestCase):
     self._CompileAndCheck(lax_op, args_maker, check_dtypes=True)
 
     if test_autodiff:
-      jtu.check_grads(lax_op, args, order=1, atol=1e-3, rtol=3e-3)
+      jtu.check_grads(lax_op, args, order=1, atol=1e-3, rtol=3e-3, eps=1e-3)
 
 
 if __name__ == "__main__":