code: def test_get_partition_spec(self):
      class Foo(nnx.Module):
          def __init__(self):
              self.w = nnx.Param(
                  nnx.with_partitioning(
                      lambda: jnp.ones((8, 2)),
                      sharding=("row", "col"),
                  )()
              )
  
          def __call__(self, x):
              return x @ self.w
  
      params, graphdef = Foo().split()
      state = nnx.TrainState(
          graphdef,
          params=params,
          tx=optax.adam(1e-3),
      )
      state_spec = nnx.get_partition_spec(state)
  
      assert state_spec.params["w"] == PartitionSpec("row", "col")
      assert state_spec.opt_state[0].mu["w"] == PartitionSpec(
          "row", "col"
      )
      assert state_spec.opt_state[0].nu["w"] == PartitionSpec(
          "row", "col"
      )
  
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 2 for calling the TestSPMD obj's underlying member method code with example usage: obj.test_get_partition_spec() and returns: unknown