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

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

[patrickvonplaten]

Do you think you can do a patch release for this one as we have a couple of Flax/JAX models in Transformers failing to due this on master at the moment and are also looking into doing another release soon.

[patrickvonplaten]

Thanks a lot for helping so quickly!