code: def test_flax_serialization(self):
      class Bar(nnx.Pytree):
          def __init__(self, a, b):
              self.a = a
              self.b = nnx.Variable(b)
  
      @nnx.dataclass
      class Foo(nnx.Pytree):
          bar: Bar
          c: int = nnx.treenode_field()
          d: int = nnx.field()
  
      foo: Foo = Foo(bar=Bar(a=1, b=2), c=3, d=4)
  
      state_dict = serialization.to_state_dict(foo)
  
      assert state_dict == {
          "bar": {
              "b": 2,
          },
          "c": 3,
      }
  
      state_dict["bar"]["b"] = 5
  
      foo = serialization.from_state_dict(foo, state_dict)
  
      assert foo.bar.b == 5
  
      del state_dict["bar"]["b"]
  
      with pytest.raises(ValueError, match="Missing field"):
          serialization.from_state_dict(foo, state_dict)
  
      state_dict["bar"]["b"] = 5
  
      # add unknown field
      state_dict["x"] = 6
  
      with pytest.raises(ValueError, match="Unknown field"):
          serialization.from_state_dict(foo, state_dict)
  
file_path: flax/flax/experimental/nnx/tests/test_pytree.py
imports: from typing import Generic, TypeVar
  import jax
  import pytest
  from flax import serialization
  from flax.experimental import nnx
label_desc: Use image node_id 4 for calling the TestPytree obj's underlying member method code with example usage: obj.test_flax_serialization() without return types