From ae7d010cb494cde2e66dce70cf91512d09c8fe82 Mon Sep 17 00:00:00 2001
From: 8bitmp3 <19637339+8bitmp3@users.noreply.github.com>
Date: Mon, 16 Dec 2024 00:01:30 +0000
Subject: [PATCH] Update NNX pop docs in graph.py

---
 flax/nnx/graph.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/flax/nnx/graph.py b/flax/nnx/graph.py
index fec21add2..951aaab21 100644
--- 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: