code: def test_scan_with_axes(self):
      # MLP Hparams
      B, L, E = 8, 4, 32  # pylint: disable=invalid-name
      # fake inputs
      x = jnp.ones((B, E))
      k = random.key(0)
  
      class SinDot(nn.Module):
          depth: int
  
          @nn.compact
          def __call__(self, x):
              W1 = partitioning.param_with_axes(  # pylint: disable=invalid-name
                  "W1",
                  nn.initializers.xavier_normal(),
                  (x.shape[-1], self.depth),
                  axes=("emb", "mlp"),
              )
              W2 = partitioning.param_with_axes(  # pylint: disable=invalid-name
                  "W2",
                  nn.initializers.xavier_normal(),
                  (self.depth, x.shape[-1]),
                  axes=("mlp", "emb"),
              )
              y = jnp.dot(jnp.sin(jnp.dot(x, W1)), W2)
              _ = partitioning.variable_with_axes(
                  "stats", "y_st", lambda: y, axes=("batch", "emb")
              )
              # scan expects a (carry, out) return signature.
              return y, None
  
      class Scanned(nn.Module):
          num_layers: int
          depth: int
  
          @nn.compact
          def __call__(self, x):
              scanned_sindot = partitioning.scan_with_axes(
                  SinDot,
                  in_axes=(),
                  variable_axes={"params": 0, "stats": 1},
                  split_rngs={"params": True},
                  axis_name="layer",
                  axes_collections=("params", "stats"),
                  length=self.num_layers,
              )(self.depth, name="scanned_layer")
              y, _ = scanned_sindot(x)
              # test calling again to test metadata compatibility across calls
              _, _ = scanned_sindot(x)
              return y
  
      p_rules = (("emb", "data"), ("mlp", "model"), ("batch", "data"))
      with partitioning.axis_rules(p_rules):
          variables = Scanned(L, E).init(k, x)
  
          # Ensure that the module can be called when 'params_axes' is not mutable.
          Scanned(L, E).apply(variables, x)
      self.assertIn("params", variables)
      self.assertIn("params_axes", variables)
      self.assertIn("stats", variables)
      self.assertIn("stats_axes", variables)
      self.assertEqual(
          variables["params_axes"]["scanned_layer"]["W1_axes"],
          partitioning.AxisMetadata(names=("layer", "emb", "mlp")),
      )
      logical_axis_names = partitioning.get_axis_names(
          variables["params_axes"]
      )
      self.assertEqual(
          logical_axis_names,
          {
              "scanned_layer": {
                  "W1": jax.sharding.PartitionSpec(
                      "layer", "emb", "mlp"
                  ),
                  "W2": jax.sharding.PartitionSpec(
                      "layer", "mlp", "emb"
                  ),
              }
          },
      )
      logical_axis_names = partitioning.get_axis_names(
          variables["stats_axes"]
      )
      self.assertEqual(
          logical_axis_names,
          {
              "scanned_layer": {
                  "y_st": jax.sharding.PartitionSpec(
                      "batch", "layer", "emb"
                  )
              }
          },
      )
  
imports: import jax
  import jax.numpy
  from absl.testing import absltest, parameterized
  from jax import random, sharding
  from jax.experimental import mesh_utils
  from flax import linen
  from flax.core import freeze, unfreeze
  from flax.linen import partitioning
label_desc: Use image node_id 15 for calling the PartitioningTest obj's underlying member method code with example usage: obj.test_scan_with_axes() and returns: y,  None, y