jax.nn.glu: fix static argname issue

google · Mar 17, 2022 · c762e07 · c762e07
1 parent c3a4a6e
commit c762e07
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/jax/_src/nn/functions.py b/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.
 

diff --git a/tests/nn_test.py b/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)