Add is_initializing API

PiperOrigin-RevId: 728413664
google · Feb 19, 2025 · 8d3c22b · 8d3c22b
1 parent c01ff11
commit 8d3c22b
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/flax/nnx/bridge/module.py b/flax/nnx/bridge/module.py
@@ -460,13 +460,16 @@ def init_with_output(
       **kwargs,
     )
 
+  def is_initializing(self) -> bool:
+    return self._object__state._initializing
+
 
 def compact(f: F) -> F:
   @functools.wraps(f)
   def compact_wrapper(self, *args, **kwargs):
     if not isinstance(self, Module):
       raise ValueError(
-        f"Expected 'self' to be a nnx.compat.Module, got {type(self).__name__}"
+        f"Expected 'self' to be a nnx.bridge.Module, got {type(self).__name__}"
       )
 
     MODULE_CONTEXT.parent_stack.append(CompactContext(self))