code: def test_tree_map(self):
      class Foo(nnx.Module, experimental_pytree=True):
          def __init__(self):
              self.node = nnx.Param(1)
              self.static = 1
  
      m = Foo()
  
      m = jax.tree_map(lambda x: x + 1, m)
  
      assert m.node == 2
      assert m.static == 1
  
imports: import dataclasses
  from copy import deepcopy
  from typing import Any, TypeVar
  import jax
  import jax.numpy
  import pytest
  from flax.experimental import nnx
label_desc: Use image node_id 1 for calling the TestModulePytree obj's underlying member method code with example usage: obj.test_tree_map() without return types