From 013aba987c08442a4c29a8def961472740f617d1 Mon Sep 17 00:00:00 2001
From: Cristian Garcia <cgarcia.e88@gmail.com>
Date: Thu, 23 Jan 2025 13:40:38 -0500
Subject: [PATCH] [nnx] fix fiddle

---
 docs_nnx/conf.py   | 2 ++
 flax/nnx/object.py | 5 ++++-
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/docs_nnx/conf.py b/docs_nnx/conf.py
index 7eee470630..6bfd970472 100644
--- a/docs_nnx/conf.py
+++ b/docs_nnx/conf.py
@@ -35,6 +35,8 @@
 sys.path.insert(0, os.path.abspath('..'))
 # Include local extension.
 sys.path.append(os.path.abspath('./_ext'))
+# Set environment variable to indicate that we are building the docs.
+os.environ['FLAX_DOC_BUILD'] = 'true'
 
 # patch sphinx
 # -- Project information -----------------------------------------------------
diff --git a/flax/nnx/object.py b/flax/nnx/object.py
index b8f35ba752..68b718f7f7 100644
--- a/flax/nnx/object.py
+++ b/flax/nnx/object.py
@@ -16,6 +16,7 @@
 
 import dataclasses
 import inspect
+import os
 import threading
 import typing as tp
 from abc import ABCMeta
@@ -157,7 +158,9 @@ def __init_subclass__(cls) -> None:
       init=cls._graph_node_init,  # type: ignore
     )
 
-    cls.__signature__ = inspect.signature(cls.__init__)
+    if 'FLAX_DOC_BUILD' in os.environ:
+      # set correct signature for sphinx
+      cls.__signature__ = inspect.signature(cls.__init__)
 
   if not tp.TYPE_CHECKING: