code: def test_scan_over_layers(self):
      def f(scope, x):
          def body(scope, x):
              kernel_init = meta.with_partitioning(
                  nn.initializers.ones_init(), ("in", "out")
              )
              y = nn.dense(scope, x, 3, kernel_init=kernel_init)
              return y, ()
  
          c, _ = lift.scan(
              body,
              variable_axes={"params": 0},
              split_rngs={"params": True},
              length=8,
              metadata_params={meta.PARTITION_NAME: "layers"},
          )(scope, x)
          return c
  
      _, variables = init(f)(random.key(0), jnp.zeros((8, 3)))
      boxed_shapes = jax.tree_map(jnp.shape, variables["params"])
      self.assertEqual(
          boxed_shapes,
          {
              "kernel": meta.Partitioned(
                  (8, 3, 3), ("layers", "in", "out")
              ),
              "bias": (8, 3),
          },
      )
  
file_path: flax/tests/core/core_meta_test.py
imports: import jax
  from absl.testing import absltest
  from jax import numpy
  from jax import random, sharding
  from jax.experimental import mesh_utils
  from flax import errors
  from flax.core import init, lift, meta, nn
label_desc: Use image node_id 5 for calling the MetaTest obj's underlying member method code with example usage: obj.test_scan_over_layers() and returns: c, y,  