Merge pull request #9942 from jakevdp:glu-fix

PiperOrigin-RevId: 435453524

Author: jax authors

jax/_src/nn/functions.py

@@ -258,7 +258,7 @@ def gelu(x: Array, approximate: bool = True) -> Array:
   else:
     return jnp.array(x * (lax.erf(x / np.sqrt(2)) + 1) / 2, dtype=x.dtype)
 
-@partial(jax.jit, static_argnames=("glu",))
+@partial(jax.jit, static_argnames=("axis",))
 def glu(x: Array, axis: int = -1) -> Array:
   """Gated linear unit activation function.
 

tests/nn_test.py

@@ -81,7 +81,7 @@ def testEluValue(self):
     self.assertAllClose(val, 1e4, check_dtypes=False)
 
   def testGluValue(self):
-    val = nn.glu(jnp.array([1.0, 0.0]))
+    val = nn.glu(jnp.array([1.0, 0.0]), axis=0)
     self.assertAllClose(val, jnp.array([0.5]))
 
   @parameterized.parameters(False, True)