Repository: google/jax

Issue:
Add custom JVP rule for jnp.sinc
The current implementation has incorrect even-ordered derivatives at `x=0.0`:
```python
import matplotlib.pyplot as plt
import jax.numpy as jnp
from jax import grad, vmap

x = jnp.linspace(-5, 5, 101)
y = vmap(grad(grad(jnp.sinc)))(x)

plt.plot(x, y)
```
![download](https://user-images.githubusercontent.com/781659/100658044-4de98480-3303-11eb-9b6c-b08ab1e28f5f.png)

Related to #5039


Discussion:

Resolving patch:
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -1005,10 +1005,24 @@ def sinc(x):
   _check_arraylike("sinc", x)
   x, = _promote_dtypes_inexact(x)
   eq_zero = lax.eq(x, lax._const(x, 0))
-  safe_x = where(eq_zero, lax._const(x, 0), x)
-  pi_x = lax.mul(lax._const(x, pi), safe_x)
-  return where(eq_zero,
-               lax._const(x, 1), lax.div(lax.sin(pi_x), pi_x))
+  pi_x = lax.mul(lax._const(x, pi), x)
+  safe_pi_x = where(eq_zero, lax._const(x, 0), pi_x)
+  return where(eq_zero, _sinc_maclaurin(0, pi_x),
+               lax.div(lax.sin(safe_pi_x), safe_pi_x))
+
+@partial(custom_jvp, nondiff_argnums=(0,))
+def _sinc_maclaurin(k, x):
+  # compute the kth derivative of x -> sin(x)/x evaluated at zero (since we
+  # compute the monomial term in the jvp rule)
+  if k % 2:
+    return lax.full_like(x, 0)
+  else:
+    return lax.full_like(x, (-1) ** (k // 2) / (k + 1))
+
+@_sinc_maclaurin.defjvp
+def _sinc_maclaurin_jvp(k, primals, tangents):
+  (x,), (t,) = primals, tangents
+  return _sinc_maclaurin(k, x), _sinc_maclaurin(k + 1, x) * t
 
 
 @_wraps(np.transpose)


Test patch:
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -4531,6 +4531,37 @@ def testOpGradSpecialValue(self, op, special_value, order):
     check_grads(op, (special_value,), order, ["fwd", "rev"],
                 atol={np.float32: 3e-3})
 
+  def testSincAtZero(self):
+    # Some manual tests for sinc at zero, since it doesn't have well-behaved
+    # numerical derivatives at zero
+    def deriv(f):
+      return lambda x: api.jvp(f, (x,), (1.,))[1]
+
+    def apply_all(fns, x):
+      for f in fns:
+        x = f(x)
+      return x
+
+    d1 = 0.
+    for ops in itertools.combinations_with_replacement([deriv, api.grad], 1):
+      self.assertAllClose(apply_all(ops, jnp.sinc)(0.), d1)
+
+    d2 = -np.pi ** 2 / 3
+    for ops in itertools.combinations_with_replacement([deriv, api.grad], 2):
+      self.assertAllClose(apply_all(ops, jnp.sinc)(0.), d2)
+
+    d3 = 0.
+    for ops in itertools.combinations_with_replacement([deriv, api.grad], 3):
+      self.assertAllClose(apply_all(ops, jnp.sinc)(0.), d3)
+
+    d4 = np.pi ** 4 / 5
+    for ops in itertools.combinations_with_replacement([deriv, api.grad], 4):
+      self.assertAllClose(apply_all(ops, jnp.sinc)(0.), d4)
+
+  def testSincGradArrayInput(self):
+    # tests for a bug almost introduced in #5077
+    jax.grad(lambda x: jnp.sinc(x).sum())(jnp.arange(10.))  # doesn't crash
+
   def testTakeAlongAxisIssue1521(self):
     # https://github.com/google/jax/issues/1521
     idx = jnp.repeat(jnp.arange(3), 10).reshape((30, 1))