code: def test_pjit_scan_over_layers(self):
      class MLP(nn.Module):
          hidden_size: int
  
          @nn.compact
          def __call__(self, x):
              ki = nn.linear.default_kernel_init
              h = nn.Dense(
                  self.hidden_size,
                  kernel_init=nn.with_partitioning(
                      ki, ("data", "model")
                  ),
              )(x)
              h = nn.relu(h)
              return nn.Dense(
                  x.shape[-1],
                  kernel_init=nn.with_partitioning(
                      ki, ("model", "data")
                  ),
              )(h)
  
      class Model(nn.Module):
          @nn.compact
          def __call__(self, x):
              def body(_, c):
                  c = MLP(512)(c)
                  return c, ()
  
              c, _ = nn.scan(
                  body,
                  variable_axes={"params": 0},
                  split_rngs={"params": 0},
                  length=8,
                  metadata_params={nn.PARTITION_NAME: None},
              )(self, x)
              return c
  
      devs = mesh_utils.create_device_mesh((jax.device_count(), 1))
      mesh = Mesh(devs, ["data", "model"])
      model = Model()
      x = jnp.ones((8, 128))
      spec = nn.get_partition_spec(
          jax.eval_shape(model.init, random.key(0), x)
      )
      self.assertEqual(
          spec,
          {
              "params": {
                  "MLP_0": {
                      "Dense_0": {
                          "bias": PartitionSpec(),
                          "kernel": PartitionSpec(
                              None, "data", "model"
                          ),
                      },
                      "Dense_1": {
                          "bias": PartitionSpec(),
                          "kernel": PartitionSpec(
                              None, "model", "data"
                          ),
                      },
                  },
              },
          },
      )
      x_spec = PartitionSpec("data", "model")
      f = lambda x: jax.sharding.NamedSharding(mesh, x)
      key_spec = PartitionSpec()
      init_fn = jax.jit(
          model.init,
          in_shardings=jax.tree_map(f, (key_spec, x_spec)),
          out_shardings=jax.tree_map(f, spec),
      )
      variables = init_fn(random.key(0), x)
      apply_fn = jax.jit(
          model.apply,
          in_shardings=jax.tree_map(f, (spec, x_spec)),
          out_shardings=jax.tree_map(f, x_spec),
      )
      y = apply_fn(variables, x)
      self.assertEqual(y.shape, (8, 128))
  
file_path: flax/tests/linen/linen_meta_test.py
imports: import jax
  from absl.testing import absltest
  from jax import numpy
  from jax import random
  from jax.experimental import mesh_utils
  from jax.sharding import Mesh, PartitionSpec
  from flax import linen
label_desc: Use image node_id 3 for calling the LinenMetaTest obj's underlying member method code with example usage: obj.test_pjit_scan_over_layers() and returns: nn, c, c,  