diff --git a/docs_nnx/api_reference/flax.nnx/graph.rst b/docs_nnx/api_reference/flax.nnx/graph.rst
index d944e3c7bf..763064b736 100644
--- a/docs_nnx/api_reference/flax.nnx/graph.rst
+++ b/docs_nnx/api_reference/flax.nnx/graph.rst
@@ -14,6 +14,7 @@ graph
 .. autofunction:: iter_graph
 .. autofunction:: clone
 .. autofunction:: call
+.. autofunction:: flatten
 
 .. autoclass:: GraphDef
   :members:
@@ -22,4 +23,4 @@ graph
   :members:
 
 .. autofunction:: update_context
-.. autofunction:: current_update_context
\ No newline at end of file
+.. autofunction:: current_update_context