diff --git a/flax/linen/module.py b/flax/linen/module.py index f8a57b9546..52e5a0594b 100644 --- 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(