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!


Discussion:
It may also be of interest (apologies if not!) that autograd doesn't have this problem:

```python
import autograd.numpy as np
from autograd.scipy.stats import norm
from autograd import elementwise_grad

inputs = np.linspace(-50, 20, 100)

fun = norm.logcdf
grad_fun = elementwise_grad(fun)
grad_grad_fun = elementwise_grad(grad_fun)
grad_grad_grad_fun = elementwise_grad(grad_grad_fun)

f = fun(inputs)
fp = grad_fun(inputs)
fpp = grad_grad_fun(inputs)
fppp = grad_grad_grad_fun(inputs)

[np.any(np.isnan(x)) for x in [f, fp, fpp, fppp]]
```

Gives:

```
[False, False, False, False]
```

And the plots actually look nicer than what my ratio trick gives, so it seems to have a nice solution there. Looks like there, the vjp is defined using exp(logcdf - logpdf):

https://github.com/HIPS/autograd/blob/master/autograd/scipy/stats/norm.py

Sorry if this was all obvious already, but I thought I'd add this information in case it's helpful!

EDIT: I could try to implement this using:
https://github.com/google/jax/pull/636
If that would be useful.
Our guess is that differentiating the (rather complicated!) implementation of `log_ndtr` isn't working that well, and that problem can be fixed by adding a derivative that treats `log_ndtr` as a black box, much like autograd does.

Resolving patch:
diff --git a/jax/scipy/special.py b/jax/scipy/special.py
--- 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 _create_polynomial(var, coeffs):
   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))))


Test patch:
diff --git a/tests/lax_scipy_test.py b/tests/lax_scipy_test.py
--- a/tests/lax_scipy_test.py
+++ b/tests/lax_scipy_test.py
@@ -65,8 +65,8 @@ def op_record(name, nargs, dtypes, rng, test_grad, test_name=None):
     # 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 @@ def testScipySpecialFun(self, scipy_op, lax_op, rng, shapes, dtypes,
     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__":