code: def test_in_axes(self):
      class Block(nnx.Module):
          def __init__(self, *, rngs: nnx.Rngs):
              self.linear = nnx.Linear(3, 3, rngs=rngs)
              self.node = nnx.Variable(jnp.ones((2,)))
  
          def __call__(
              self, x: jax.Array, a: jax.Array
          ) -> tp.Tuple[jax.Array, None]:
              assert x.shape == a.shape
              x = x + a
              x = self.linear(x)
              x = nnx.gelu(x)
              return x, None
  
      MLP = nnx.Scan(
          Block,
          variable_axes={nnx.Param: 0},
          length=5,
      )
  
      module = MLP(rngs=nnx.Rngs(0))
  
      assert module.scan_module.linear.kernel.shape == (5, 3, 3)
      assert module.scan_module.linear.bias.shape == (5, 3)
      assert module.scan_module.node.shape == (2,)
  
      x = jnp.ones((1, 3))
      a = jnp.ones((5, 1, 3))
      y, out = module(x, a)
  
      assert y.shape == (1, 3)
      assert out is None
  
imports: import typing
  from functools import partial
  import jax
  import jax.numpy
  import pytest
  from flax.experimental import nnx
label_desc: Use image node_id 4 for calling the TestScan obj's underlying member method code with example usage: obj.test_in_axes() and returns: x,  None