code: def test_pytree_nodes(self):
      p = Point(x=1, y=2, meta={"abc": True})
      leaves = jax.tree_util.tree_leaves(p)
      self.assertEqual(leaves, [1, 2])
      new_p = jax.tree_util.tree_map(lambda x: x + x, p)
      self.assertEqual(new_p, Point(x=2, y=4, meta={"abc": True}))
  
file_path: flax/tests/struct_test.py
imports: import dataclasses
  from typing import Any
  import jax
  from absl.testing import absltest
  from jax._src.tree_util import prefix_errors
  from flax import struct
label_desc: Use image node_id 3 for calling the StructTest obj's underlying member method code with example usage: obj.test_pytree_nodes() without return types