improve sinc jvp at zero, fixes #5054

jax/_src/numpy/lax_numpy.py

@@ -1005,10 +1005,23 @@ 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_deriv(0, pi_x),
+ lax.div(lax.sin(safe_pi_x), safe_pi_x))
+
+@partial(custom_jvp, nondiff_argnums=(0,))
+def _sinc_deriv(k, x):
+ # compute the kth derivative of x -> sin(x)/x evaluated at zero
+ if k % 2:
+ return lax._const(x, 0)
+ else:
+ return lax._const(x, (-1) ** (k // 2) / (k + 1))
+
+@_sinc_deriv.defjvp
+def _sinc_maclaurin_jvp(k, primals, tangents):
+ (x,), (t,) = primals, tangents
+ return _sinc_deriv(k, x), _sinc_deriv(k + 1, x) * t
 
 
 @_wraps(np.transpose)

tests/lax_numpy_test.py

@@ -4531,6 +4531,33 @@ 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 testTakeAlongAxisIssue1521(self):
  # https://github.com/google/jax/issues/1521
  idx = jnp.repeat(jnp.arange(3), 10).reshape((30, 1))