Update NNX NodeDef docs in graph.py

flax/nnx/graph.py

@@ -292,9 +292,9 @@ def __treescope_repr__(self, path, subtree_renderer):
 
 @dataclasses.dataclass(frozen=True, repr=False, slots=True)
 class NodeDef(GraphDef[Node], reprlib.Representable):
-  """A dataclass that denotes the tree structure of a
-  :class:`Module`. A ``GraphDef`` can be generated by either
-  calling :func:`split` or :func:`graphdef` on the :class:`Module`."""
+  """A dataclass that denotes the JAX pytree structure of a :class:`flax.nnx.Module`.
+  A :class:`flax.nnx.GraphDef` can be generated by either calling
+  :func:`flax.nnx.split` or :func:`flax.nnx.graphdef` on the ``nnx.Module``."""
 
   type: tp.Type[Node]
   index: int
@@ -1852,4 +1852,4 @@ def _unflatten_pytree(
   type(None),
   flatten=lambda x: ([], None),
   unflatten=lambda _, __: None,  # type: ignore
-)
+)