Makes flax Modules more compatible with IPython auto-reload.

PiperOrigin-RevId: 703062648
google · Jan 6, 2025 · bcab3b9 · bcab3b9
1 parent 53bde74
commit bcab3b9
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/flax/linen/module.py b/flax/linen/module.py
@@ -1274,6 +1274,11 @@ def __setattr__(self, name: str, val: Any):
         object.__setattr__(self, name, val)
         return
       else:
+        # If the attribute is a python special method, we allow setting it (this
+        # is useful e.g. for IPython auto-reload).
+        if name.startswith('__'):
+          object.__setattr__(self, name, val)
+          return
         # We're past all initialization and setup logic:
         # Raises a TypeError just like frozen python dataclasses.
         raise errors.SetAttributeFrozenModuleError(