Merge pull request jax-ml#25392 from carlosgmartin:add_nn_relu_grad_at_zero_test_update_paper_link

PiperOrigin-RevId: 704947960

Author: Google-ML-Automation

jax/_src/nn/functions.py

@@ -67,7 +67,7 @@ def relu(x: ArrayLike) -> Array:
 
   For more information see
   `Numerical influence of ReLU’(0) on backpropagation
-  <https://openreview.net/forum?id=urrcVI-_jRm>`_.
+  <https://dl.acm.org/doi/10.5555/3540261.3540297>`_.
 
   Args:
     x : input array
@@ -84,7 +84,7 @@ def relu(x: ArrayLike) -> Array:
 
   """
   return jnp.maximum(x, 0)
-# For behavior at 0, see https://openreview.net/forum?id=urrcVI-_jRm
+# For behavior at 0, see https://dl.acm.org/doi/10.5555/3540261.3540297
 relu.defjvps(lambda g, ans, x: lax.select(x > 0, g, lax.full_like(g, 0)))
 
 @jax.jit

tests/nn_test.py

@@ -317,6 +317,11 @@ def testReluGrad(self):
     jaxpr = jax.make_jaxpr(jax.grad(nn.relu))(0.)
     self.assertGreaterEqual(len(jaxpr.jaxpr.eqns), 2)
 
+  def testReluGradAtZero(self):
+    # https://dl.acm.org/doi/10.5555/3540261.3540297
+    grad = jax.grad(nn.relu)(0.)
+    self.assertEqual(grad, 0.)
+
   def testRelu6Grad(self):
     rtol = 1e-2 if jtu.test_device_matches(["tpu"]) else None
     check_grads(nn.relu6, (1.,), order=3, rtol=rtol)