From 34b8c6b01dc9c4879dfe448f7009c93ce1e787c9 Mon Sep 17 00:00:00 2001
From: 8bitmp3 <19637339+8bitmp3@users.noreply.github.com>
Date: Fri, 4 Oct 2024 23:07:57 +0000
Subject: [PATCH 1/2] Add flatten to nnx.graph autosummary in graph.rst

---
 docs_nnx/api_reference/flax.nnx/graph.rst | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

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

From 149144e047eb8ab0b68db7b4af6412f1bbfc1db5 Mon Sep 17 00:00:00 2001
From: 8bitmp3 <19637339+8bitmp3@users.noreply.github.com>
Date: Fri, 4 Oct 2024 23:33:18 +0000
Subject: [PATCH 2/2] Update graph.rst

---
 docs_nnx/api_reference/flax.nnx/graph.rst | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/docs_nnx/api_reference/flax.nnx/graph.rst b/docs_nnx/api_reference/flax.nnx/graph.rst
index c70ac016f5..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
@@ -18,8 +19,6 @@ graph
 .. autoclass:: GraphDef
   :members:
 
-.. autofunction:: flatten
-
 .. autoclass:: UpdateContext
   :members: