Merge pull request #18706 from jakevdp:i0-grad

PiperOrigin-RevId: 586079527

jax/_src/numpy/lax_numpy.py

@@ -49,6 +49,7 @@
 from jax._src import api_util
 from jax._src import config
 from jax._src import core
+from jax._src.custom_derivatives import custom_jvp
 from jax._src import dispatch
 from jax._src import dtypes
 from jax._src.api_util import _ensure_index_tuple
@@ -2644,6 +2645,7 @@ def meshgrid(*xi: ArrayLike, copy: bool = True, sparse: bool = False,
   return output
 
 
+@custom_jvp
 @util._wraps(np.i0)
 @jit
 def i0(x: ArrayLike) -> Array:
@@ -2653,6 +2655,11 @@ def i0(x: ArrayLike) -> Array:
   x_arr = lax.abs(x_arr)
   return lax.mul(lax.exp(x_arr), lax.bessel_i0e(x_arr))
 
+@i0.defjvp
+def _i0_jvp(primals, tangents):
+  primal_out, tangent_out = jax.jvp(i0.fun, primals, tangents)
+  return primal_out, where(primals[0] == 0, 0.0, tangent_out)
+
 
 @util._wraps(np.ix_)
 def ix_(*args: ArrayLike) -> tuple[Array, ...]:

tests/lax_numpy_operators_test.py

@@ -641,6 +641,11 @@ def __rmul__(self, other):
     self.assertIsInstance(b * a, MyArray)
     self.assertIsInstance(jax.jit(operator.mul)(b, a), MyArray)
 
+  def testI0Grad(self):
+    # Regression test for https://github.com/google/jax/issues/11479
+    dx = jax.grad(jax.numpy.i0)(0.0)
+    self.assertArraysEqual(dx, 0.0)
+
 
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())