Simplify implementation of nn.relu.

jax/_src/nn/functions.py

@@ -25,7 +25,6 @@
 
 import jax
 import jax.numpy as jnp
-from jax import custom_jvp
 from jax import lax
 from jax._src import config
 from jax._src import core
@@ -50,7 +49,6 @@ def __repr__(self):
 
 # activations
 
-@custom_jvp
 @jax.jit
 def relu(x: ArrayLike) -> Array:
   r"""Rectified linear unit activation function.
@@ -83,9 +81,7 @@ def relu(x: ArrayLike) -> Array:
     :func:`relu6`
 
   """
-  return jnp.maximum(x, 0)
-# 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)))
+  return jnp.where(jnp.greater(x, 0), x, lax.zeros_like_array(x))
 
 @jax.jit
 def squareplus(x: ArrayLike, b: ArrayLike = 4) -> Array: