Makes flax Modules more compatible with IPython auto-reload.

PiperOrigin-RevId: 712802399
google · Jan 7, 2025 · e2134af · e2134af
1 parent dd70960
commit e2134af
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(