Repository: google/flax

Issue:
Bug in error catching
The following code throws an error at the indicated line:

```
class SomeModule(nn.Module):
  @nn.compact
  def __call__(self, x):
    dense = nn.Dense(10)
    dense.features = 20
    dense.new_attr = 20  # <--- ERROR!
    return dense(x)

SomeModule().init(random.PRNGKey(0), jnp.ones((5, 5)))
```

The error is: `Module instance is frozen outside of setup method.` This seems odd: if `dense` were frozen, why are we allowed to modify the existing attribute `features`, but we aren't allowed to add a new one called `new_attr`? It seems we should not be allowed to modify `dense.features` are all and we should already throw an error the line before.


Discussion:

Resolving patch:
diff --git a/flax/linen/module.py b/flax/linen/module.py
--- a/flax/linen/module.py
+++ b/flax/linen/module.py
@@ -329,11 +329,16 @@ class _ModuleInternalState:
   in_compact_method: bool = False
   in_setup: bool = False
   setup_called: bool = False
+  is_initialized: bool = False
   autoname_cursor: Optional[dict] = dataclasses.field(default_factory=dict)
   children: Dict[str, Union[str, 'Module']] = dataclasses.field(default_factory=dict)
 
   def reset(self):
-    """Resets transient state."""
+    """Resets transient state. 
+    
+    This function is called after each module method, so only attributes that
+    are method-dependent are reset.
+    """
     self.in_compact_method = False
     self.in_setup = False
     self.autoname_cursor = dict()
@@ -344,6 +349,7 @@ def export(self):
       in_compact_method=self.in_compact_method,
       in_setup=self.in_setup,
       setup_called=False,  # setup_called is object local, not shared.
+      is_initialized=self.is_initialized,
       autoname_cursor=dict(self.autoname_cursor))
     return cloned
 
@@ -351,6 +357,7 @@ def reimport(self, other):
     """Re-imports transform-preserved state from across transform boundary."""
     self.in_compact_method = other.in_compact_method
     self.in_setup = other.in_setup
+    self.is_initialized = other.is_initialized
     self.autoname_cursor = dict(other.autoname_cursor)
 
 _uninitialized_module_internal_state = _ModuleInternalState()
@@ -504,8 +511,8 @@ def __setattr__(self, name: str, val: Any):
       val: Value of the attribute.
     """
     is_dataclass_attr = name in self.__dataclass_fields__ and self.__dataclass_fields__[name].init  # pytype: disable=attribute-error
-    
-    if not self._state.in_setup and not is_dataclass_attr:
+
+    if not self._state.in_setup and self._state.is_initialized:
       # Raises a TypeError just like frozen python dataclasses.
       raise TypeError("Module instance is frozen outside of setup method.")
     if is_dataclass_attr:
@@ -584,6 +591,8 @@ def __post_init__(self):
     else:
       raise ValueError("parent must be None, Module or Scope")
 
+    self._state.is_initialized = True
+
   def __repr__(self):
     return _module_repr(self)
 


Test patch:
diff --git a/tests/linen/module_test.py b/tests/linen/module_test.py
--- a/tests/linen/module_test.py
+++ b/tests/linen/module_test.py
@@ -691,19 +691,35 @@ def __call__(self, x):
     variables = foo.init(random.PRNGKey(0), x)
     self.assertEqual(variables['params']['bar']['kernel'].shape, (2, 3))
 
-  def test_module_frozen(self):
+  def test_noncompact_module_frozen(self):
     class Foo(nn.Module):
-      bar: nn.Dense = dataclasses.field(init=False)
-
       def setup(self):
-        self.i = 1
+        self.i = 1  # This is allowed (for assigning submodules).
+
+      def __call__(self):
+        self.i = 2  # This is not allowed.
 
+    with self.assertRaisesWithLiteralMatch(TypeError, "Module instance is frozen outside of setup method."):
+      Foo().init(random.PRNGKey(0))
+
+  def test_compact_module_frozen(self):
+    class Foo(nn.Module):
+      @nn.compact
       def __call__(self):
         self.i = 2
 
-    foo = Foo()
     with self.assertRaisesWithLiteralMatch(TypeError, "Module instance is frozen outside of setup method."):
-      foo.init(random.PRNGKey(0))
+      Foo().init(random.PRNGKey(0))
+
+  def test_submodule_frozen(self):
+    class Foo(nn.Module):
+      @nn.compact
+      def __call__(self):
+        dense = nn.Dense(10)
+        dense.features = 20  # <--- This is not allowed
+
+    with self.assertRaisesWithLiteralMatch(TypeError, "Module instance is frozen outside of setup method."):
+      Foo().init(random.PRNGKey(0))
   
   def test_is_mutable_collection(self):
     class EmptyModule(nn.Module):