code: def test_dataclass_inheritance(self):
      @nnx.dataclass
      class A(nnx.Pytree):
          a: int = nnx.treenode_field(default=1)
          b: int = nnx.field(default=2)
  
      @nnx.dataclass
      class B(A):
          c: int = nnx.treenode_field(default=3)
  
      pytree = B()
      leaves = jax.tree_util.tree_leaves(pytree)
      assert leaves == [1, 3]
  
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 8 for calling the TestPytree obj's underlying member method code with example usage: obj.test_dataclass_inheritance() without return types