diff --git a/docs_nnx/api_reference/flax.nnx/graph.rst b/docs_nnx/api_reference/flax.nnx/graph.rst
index d944e3c7bf..c70ac016f5 100644
--- a/docs_nnx/api_reference/flax.nnx/graph.rst
+++ b/docs_nnx/api_reference/flax.nnx/graph.rst
@@ -18,8 +18,10 @@ graph
 .. autoclass:: GraphDef
   :members:
 
+.. autofunction:: flatten
+
 .. autoclass:: UpdateContext
   :members:
 
 .. autofunction:: update_context
-.. autofunction:: current_update_context
\ No newline at end of file
+.. autofunction:: current_update_context