code: def test_pytree(self):
      class Foo(nnx.Pytree, mutable=True):
          def __init__(self, y) -> None:
              self.x = 2
              self.y = nnx.Variable(y)
  
      pytree = Foo(y=3)
  
      leaves = jax.tree_util.tree_leaves(pytree)
      assert leaves == [3]
  
      pytree = jax.tree_map(lambda x: x * 2, pytree)
      assert pytree.x == 2
      assert pytree.y == 6
  
      pytree = pytree.replace(x=3)
      assert pytree.x == 3
      assert pytree.y == 6
  
      # test mutation
      pytree.x = 4
      assert pytree.x == 4
  
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 1 for calling the TestMutablePytree obj's underlying member method code with example usage: obj.test_pytree() without return types