Update NNX pop docs in graph.py

google · Dec 16, 2024 · ae7d010 · ae7d010
1 parent 6bc9858
commit ae7d010
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/flax/nnx/graph.py b/flax/nnx/graph.py
@@ -1526,7 +1526,7 @@ def pop(
 def pop(
   node, *filters: filterlib.Filter
 ) -> tp.Union[GraphState, tuple[GraphState, ...]]:
-  """Pop one or more :class:`Variable` types from the graph node.
+  """Pop one or more :class:`flax.nnx.Variable` types from the graph node.
 
   Example usage::
 
@@ -1555,9 +1555,9 @@ def pop(
 
   Args:
     node: A graph node object.
-    *filters: One or more :class:`Variable` objects to filter by.
+    *filters: One or more :class:`flax.nnx.Variable` objects to filter by.
   Returns:
-    The popped :class:`State` containing the :class:`Variable`
+    The popped :class:`flax.nnx.State` containing the :class:`flax.nnx.Variable`
     objects that were filtered for.
   """
   if len(filters) == 0: