code: def test_partition(self):
      m = nnx.Dict(
          a=nnx.Sequence([nnx.Param(1), nnx.BatchStat(2)]),
          b=nnx.Param(2),
          c=100,
      )
  
      params, rest, graphdef = m.split(nnx.Param, ...)
  
      assert len(params) == 2
      assert len(rest) == 1
  
      # check params
      assert params["a"]["0"] == m.a[0]
      assert params["b"] == m.b
  
      # check rest
      assert rest["a"]["1"] == m.a[1]
  
      m2 = graphdef.merge(params, rest)
  
      assert m2.a[0] == m.a[0]
      assert m2.a[1] == m.a[1]
      assert m2.b == m.b
      assert m2.c == 100
  
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 1 for calling the TestPartitioning obj's underlying member method code with example usage: obj.test_partition() without return types