code: def test_update_from(self):
      m = nnx.Dict(
          a=nnx.Sequence([nnx.Param(1), nnx.BatchStat(3)]),
          b=nnx.Param(2),
          c=100,
      )
  
      state = m.split()[0]
      state = jax.tree_map(lambda x: x * 2, state)
  
      m.update(state)
  
      assert m.a[0] == 2
      assert m.a[1] == 6
      assert m.b == 4
      assert m.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 6 for calling the TestPartitioning obj's underlying member method code with example usage: obj.test_update_from() without return types