[nnx] improve nnx.scan in_axes/out_axes

[cgarciae]

What does this PR do?

Makes nnx.scan's in_axes and out_axes API is more flexible in the following ways:

• The function is no longer limited to returning carry, outputs, out_axes can use Carry to control its structure, by default its now (Carry, 0) to match current behavior.
• Carry is no longer required to be present in in_axes/out_axes, however if Carry is present in in_axes it must also be present in out_axes and vice versa.
• in_axes is not limited to a Sequence, as with vmap a single top-level integer or None can now be used to configure all arguments. A top-level Carry is also accepted by in_axes but its interpreted as (Carry,), it means that the function accepts a single argument that is the Carry.

Example: Scan over layers

state_axes = nnx.StateAxes({(nnx.Param, nnx.RngState): 0, ...: None})

class MLP(nnx.Module):
  @nnx.split_rngs(splits=5)
  @nnx.vmap(in_axes=(state_axes, 0))
  def __init__(self, rngs: nnx.Rngs):
    self.linear = nnx.Linear(3, 3, rngs=rngs)

  @nnx.scan(in_axes=(state_axes, nnx.Carry), out_axes=nnx.Carry)
  def __call__(self, x: jax.Array):
    return nnx.gelu(self.linear(x))

module = MLP(nnx.Rngs(0))

x = jnp.ones((1, 3))
y = module(x)

assert y.shape == (1, 3)

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB