bug in jax.nn.glu
There is an issue with `jax.nn.glu`  in the latest release

https://github.com/google/jax/blob/728e4fd3fad334e551dcd1a13b47b8d51aae5a9f/jax/_src/nn/functions.py#L261

Here the `static_argnames` should be `axis`. Right now  calling `glu`  results in a `TracerIntegerConversionError`

```python
import jax
import jax.numpy as jnp

jax.nn.glu(jnp.ones((1, 4)), axis=1)
```

full error messages/tracebacks.
```
TracerIntegerConversionError              Traceback (most recent call last)
Input In [8], in <module>
----> 1 jax.nn.glu(jnp.ones((1, 4)), 1)

    [... skipping hidden 14 frame]

File ~/transformers-env/lib/python3.8/site-packages/jax/_src/nn/functions.py:269, in glu(x, axis)
    261 @partial(jax.jit, static_argnames=("glu",))
    262 def glu(x: Array, axis: int = -1) -> Array:
    263   """Gated linear unit activation function.
    264 
    265   Args:
    266     x : input array
    267     axis: the axis along which the split should be computed (default: -1)
    268   """
--> 269   size = x.shape[axis]
    270   assert size % 2 == 0, "axis size must be divisible by 2"
    271   x1, x2 = jnp.split(x, 2, axis)

File ~/transformers-env/lib/python3.8/site-packages/jax/core.py:473, in Tracer.__index__(self)
    472 def __index__(self):
--> 473   raise TracerIntegerConversionError(self)

TracerIntegerConversionError: The __index__() method was called on the JAX Tracer object Traced<ShapedArray(int32[], weak_type=True)>with<DynamicJaxprTrace(level=0/1)>
See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerIntegerConversionError
```

diff --git a/jax/_src/nn/functions.py b/jax/_src/nn/functions.py
index 3d568983b..bc386bb7c 100644
--- 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
index 1659f7d40..60b0764a8 100644
--- a/tests/nn_test.py
+++ b/tests/nn_test.py
@@ -81,7 +81,7 @@ class NNFunctionsTest(jtu.JaxTestCase):
     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)
['jax/_src/nn/functions.py', 'tests/nn_test.py']