diff --git a/docs_nnx/api_reference/flax.nnx/graph.rst b/docs_nnx/api_reference/flax.nnx/graph.rst
index d944e3c7bf..9c7acd2f44 100644
--- a/docs_nnx/api_reference/flax.nnx/graph.rst
+++ b/docs_nnx/api_reference/flax.nnx/graph.rst
@@ -10,6 +10,7 @@ graph
 .. autofunction:: update
 .. autofunction:: pop
 .. autofunction:: state
+.. autofunction:: graph
 .. autofunction:: graphdef
 .. autofunction:: iter_graph
 .. autofunction:: clone
@@ -22,4 +23,4 @@ graph
   :members:
 
 .. autofunction:: update_context
-.. autofunction:: current_update_context
\ No newline at end of file
+.. autofunction:: current_update_context