From e2134afa4e11093c0072861a7b4832b1513074d8 Mon Sep 17 00:00:00 2001
From: Flax Team <no-reply@google.com>
Date: Tue, 7 Jan 2025 00:35:52 -0800
Subject: [PATCH] Makes flax Modules more compatible with IPython auto-reload.

PiperOrigin-RevId: 712802399
---
 flax/linen/module.py | 5 +++++
 1 file changed, 5 insertions(+)

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(