code: def test_init(self):
      class Foo(nnx.Module):
          def __init__(self):
              self.w = nnx.Param(
                  nnx.with_partitioning(
                      lambda: jnp.ones((8, 2)),
                      sharding=("model", "data"),
                  )()
              )
  
          def __call__(self, x):
              return x @ self.w
  
      @jax.jit
      def create_module():
          return Foo().split()
  
      mesh = Mesh(
          mesh_utils.create_device_mesh((2, 2)), ("model", "data")
      )
  
      with mesh:
          m: Foo = nnx.merge(create_module())
  
      assert m.w.shape == (8, 2)
      assert m.w.sharding.shard_shape(m.w.shape) == (4, 1)
  
file_path: flax/flax/experimental/nnx/tests/test_spmd.py
imports: import jax
  import jax.numpy
  import optax
  from jax._src import test_util
  from jax.experimental import mesh_utils
  from jax.sharding import Mesh, PartitionSpec
  from flax.experimental import nnx
label_desc: Use image node_id 1 for calling the TestSPMD obj's underlying member method code with example usage: obj.test_init() and returns: Foo, unknown