code: def test_key_paths(self):
      @nnx.dataclass
      class Bar(nnx.Pytree):
          a: int = nnx.treenode_field(default=1)
          b: int = nnx.field(default=2)
  
      @nnx.dataclass
      class Foo(nnx.Pytree):
          x: int = nnx.treenode_field(default=3)
          y: int = nnx.field(default=4)
          z: Bar = nnx.treenode_field(default_factory=Bar)
  
      foo = Foo()
  
      path_values, treedef = jax.tree_util.tree_flatten_with_path(foo)
      path_values = [
          (list(map(str, path)), value) for path, value in path_values
      ]
  
      assert path_values[0] == ([".x", ".value"], 3)
      assert path_values[1] == ([".z", ".value", ".a", ".value"], 1)
  
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 6 for calling the TestPytree obj's underlying member method code with example usage: obj.test_key_paths() without return types