code: def test_get_paritition(self):
      m = nnx.Dict(
          a=nnx.Sequence([nnx.Param(10.0), nnx.Param(20.0)]),
          b=nnx.Param(10.0),
          c=7,
          d=5.0,
      )
  
      # test Variables not shared
      assert vars(m.a)["0"] is not vars(m)["b"]
  
      state = m.extract(nnx.Variable)
      assert state["a"]["0"] == m.a[0]
      assert state["a"]["1"] == m.a[1]
      assert state["b"] == m.b
      assert state.variables.b is not state.a.variables[0]
      assert len(state.flat_state()) == 3
  
file_path: flax/flax/experimental/nnx/tests/test_partitioning.py
imports: import jax
  import pytest
  from flax.experimental import nnx
label_desc: Use image node_id 9 for calling the TestPartitioning obj's underlying member method code with example usage: obj.test_get_paritition() without return types