code: def test_scan_with_sharding(self):
      class Block(nnx.Module):
          def __init__(self, *, rngs: nnx.Rngs):
              self.linear = nnx.Linear(
                  3,
                  3,
                  kernel_init=nnx.with_metadata(
                      nnx.initializers.lecun_normal(),
                      sharding=("din", "dout"),
                  ),
                  bias_init=nnx.with_metadata(
                      nnx.initializers.zeros_init(),
                      sharding=("dout",),
                  ),
                  rngs=rngs,
              )
  
          def __call__(
              self, x: jax.Array, _
          ) -> tp.Tuple[jax.Array, None]:
              x = self.linear(x)
  
              # test sharding layer axes is not present inside scan
              state = self.linear.get_state()
              assert state.kernel.shape == (3, 3)
              assert state.variables.kernel.sharding == ("din", "dout")
              assert state.bias.shape == (3,)
              assert state.variables.bias.sharding == ("dout",)
  
              return x, None
  
      MLP = nnx.Scan(
          Block,
          variable_axes={nnx.Param: 0},
          length=5,
          scan_metadata={nnx.PARTITION_NAME: "layers"},
      )
  
      m = MLP(rngs=nnx.Rngs(0))
  
      # test sharding layers axes is set
      state = m.get_state()
      assert state.scan_module.linear.variables.kernel.value.shape == (
          5,
          3,
          3,
      )
      assert state.scan_module.linear.variables.kernel.sharding == (
          "layers",
          "din",
          "dout",
      )
      assert state.scan_module.linear.variables.bias.value.shape == (
          5,
          3,
      )
      assert state.scan_module.linear.variables.bias.sharding == (
          "layers",
          "dout",
      )
  
      x = jnp.ones((1, 3))
      y, out = m(x, None)
  
      # test sharding axes is preserved
      state = m.get_state()
      assert state.scan_module.linear.kernel.shape == (5, 3, 3)
      assert state.scan_module.linear.variables.kernel.sharding == (
          "layers",
          "din",
          "dout",
      )
      assert state.scan_module.linear.bias.shape == (5, 3)
      assert state.scan_module.linear.variables.bias.sharding == (
          "layers",
          "dout",
      )
  
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 9 for calling the TestScan obj's underlying member method code with example usage: obj.test_scan_with_sharding() and returns: x,  None