Repository: google/flax

Issue:
Linen modules should be frozen
Currently we don't enforce linen Modules to be frozen after setup. However, this should be the case because Module instances need to be clone-able to work correctly. `__setattr__` should refuse to set attributes after setup is finished.

Update:

Actually there are more sharp edges that can be fixed by freezing correctly.
Currently we accept lists and dicts of sub modules which are registered on assignment. But we can actually freeze them to avoid this common trap:

```
def setup(self):
  self.sub_modules = [Dense()]
  self.sub_modules.append(Dense())
```

We could avoid this by making sub_modules is stored as a tuple and similarly we can avoid the same issue with dicts by transforming them into a FrozenDict



Discussion:
Marking as "pull requests welcome" if anyone wants to take a look at it. This change will help avoid possible footguns for users.

Resolving patch:
diff --git a/flax/linen/linear.py b/flax/linen/linear.py
--- a/flax/linen/linear.py
+++ b/flax/linen/linear.py
@@ -257,8 +257,7 @@ def __call__(self, inputs: Array) -> Array:
       is_single_input = True
       inputs = jnp.expand_dims(inputs, axis=0)
 
-    if self.strides is None:
-      self.strides = (1,) * (inputs.ndim - 2)
+    strides = self.strides or (1,) * (inputs.ndim - 2)
 
     in_features = inputs.shape[-1]
     assert in_features % self.feature_group_count == 0
@@ -271,7 +270,7 @@ def __call__(self, inputs: Array) -> Array:
     y = lax.conv_general_dilated(
         inputs,
         kernel,
-        self.strides,
+        strides,
         self.padding,
         lhs_dilation=self.input_dilation,
         rhs_dilation=self.kernel_dilation,
diff --git a/flax/linen/module.py b/flax/linen/module.py
--- a/flax/linen/module.py
+++ b/flax/linen/module.py
@@ -31,7 +31,7 @@
 from flax import serialization
 from flax.core import Scope, apply
 from flax.core.scope import CollectionFilter, Variable, VariableDict
-from flax.core.frozen_dict import freeze
+from flax.core.frozen_dict import FrozenDict, freeze
 
 # from .dotgetter import DotGetter
 
@@ -72,7 +72,7 @@ def _module_repr(module: 'Module', num_spaces: int = 4):
   rep = ''
   attributes = {k: v for k, v in cls.__annotations__.items()
                 if k not in ('parent', 'name')}
-  child_modules = {k: v for k, v in module.children.items()  # pytype: disable=attribute-error
+  child_modules = {k: v for k, v in module._state.children.items()  # pytype: disable=attribute-error
                    if isinstance(v, Module)}
   if attributes:
     rep += '# attributes\n'
@@ -126,7 +126,7 @@ def disable_named_call():
   _use_named_call = False
 
 
-# Utilities for autonaming pytrees of Modules defined inside setup()
+# Utilities for pytrees of Modules defined inside setup()
 # -----------------------------------------------------------------------------
 def _get_suffix_value_pairs(
     tree_or_leaf: Any) -> List[Tuple[str, Type["Module"]]]:
@@ -153,6 +153,15 @@ def _all_names_on_object(obj: Any) -> Set[str]:
   return nameset
 
 
+def _freeze_attr(val: Any) -> Any:
+  if isinstance(val, (dict, FrozenDict)):
+    return FrozenDict({k: _freeze_attr(v) for k, v in val.items()})
+  elif isinstance(val, (list, tuple)):
+    return tuple(_freeze_attr(v) for v in val)
+  else:
+    return val
+
+
 # Method wrapping of "compact methods" and setup()
 # -----------------------------------------------------------------------------
 def compact(fun: Callable) -> Callable:
@@ -268,6 +277,8 @@ class _ModuleInternalState:
   in_setup: bool = False
   last_varname: Optional[str] = None
   autoname_cursor: Optional[dict] = dataclasses.field(default_factory=dict)
+  frozen: bool = False
+  children: Dict[str, Union[str, 'Module']] = dataclasses.field(default_factory=dict)
 
   def reset(self):
     self.in_compact_method = False
@@ -408,6 +419,10 @@ def __setattr__(self, name: str, val: Any):
       name: Attribute to set.
       val: Value of the attribute.
     """
+    if name != '_state' and self._state.frozen:
+      # raises a TypeError just like frozen python dataclasses
+      raise TypeError("Module instance is frozen outside of setup method.")
+
     # We don't mess with the parent module.
     if name == 'parent':
       pass
@@ -416,6 +431,7 @@ def __setattr__(self, name: str, val: Any):
       pass
     # Submodules are being defined and attached in setup()
     else:
+      val = _freeze_attr(val)
       for suffix, subvalue in _get_suffix_value_pairs(val):
         if isinstance(subvalue, Module):
           if not self._state.in_setup:
@@ -454,7 +470,6 @@ def __post_init__(self):
     # this Module at the top-level to variables and rngs.
 
     self._state = _ModuleInternalState()
-    self.children = dict()  # tracks child modules
 
     # Typically we set the parent based on the dynamic module context.
     if self.parent is _unspecified_parent:  # pytype: disable=attribute-error
@@ -488,7 +503,7 @@ def __post_init__(self):
             f"trying to share submodule {self.__class__.__name__} by name "
             f"{self.name}. To share submodules, store module instances as a"
             f" Python object or as an attribute on self and reuse.")
-      self.parent.children[self.name] = self
+      self.parent._state.children[self.name] = self
       self.scope = self.parent.scope.push(self.name)
 
     # Top-level invocation with a functional Scope.
@@ -500,6 +515,7 @@ def __post_init__(self):
 
     # Call the user-defined initialization setup() function.
     self.setup()
+    self._state.frozen = True
 
   def __repr__(self):
     return _module_repr(self)
@@ -590,7 +606,7 @@ def variable(self, col: str, name: str, init_fn, *init_args) -> Variable:
     # ephemeral state for setattr name-equality-check
     self._state.last_varname = name
     v = self.scope.variable(col, name, init_fn, *init_args)
-    self.children[name] = col
+    self._state.children[name] = col
     return v
 
   def param(self, name: str, init_fn: Callable[..., T], *init_args) -> T:
@@ -619,7 +635,7 @@ def param(self, name: str, init_fn: Callable[..., T], *init_args) -> T:
     # ephemeral state for setattr name-equality-check
     self._state.last_varname = name
     v = self.scope.param(name, init_fn, *init_args)
-    self.children[name] = 'params'
+    self._state.children[name] = 'params'
     return v
 
   def has_variable(self, col: str, name: str) -> bool:
diff --git a/flax/linen/transforms.py b/flax/linen/transforms.py
--- a/flax/linen/transforms.py
+++ b/flax/linen/transforms.py
@@ -139,8 +139,6 @@ def core_fn(scopes, *args, **kwargs):
         cloned = set_module_scopes(cloned, scopes)
         cloned._state = copy.deepcopy(self._state)  # pylint: disable=protected-access
         res = fn(cloned, *args, **kwargs)
-        # preserve submodule-tree stripped of scopes/tracers for introspection
-        object.__setattr__(self, 'children', clean_clone(cloned).children)
         self._state = copy.deepcopy(cloned._state)  # pylint: disable=protected-access
         return res
       # here we apply the given lifting transform to the scope-ingesting fn
@@ -172,8 +170,6 @@ def core_fn(scopes, *args, **kwargs):
       cloned = set_module_scopes(self, scopes)
       cloned._state = copy.deepcopy(self._state)  # pylint: disable=protected-access
       res = rewrapped_fn(cloned, *args, **kwargs)
-      # preserve submodule-tree stripped of scopes/tracers for introspection
-      object.__setattr__(self, 'children', clean_clone(cloned).children)
       self._state = copy.deepcopy(cloned._state)  # pylint: disable=protected-access
       return res
     # here we apply the given lifting transform to the scope-ingesting fn
@@ -224,8 +220,6 @@ def core_fn(scopes, *args, **kwargs):
       cloned = set_module_scopes(self, scopes)
       cloned._state = copy.deepcopy(self._state)  # pylint: disable=protected-access
       res = rewrapped_fn(cloned, *args, **kwargs)
-      # preserve submodule-tree stripped of scopes/tracers for introspection
-      object.__setattr__(self, 'children', clean_clone(cloned).children)
       self._state = copy.deepcopy(cloned._state)  # pylint: disable=protected-access
       return res
     # here we apply the given lifting transform to the scope-ingesting fn


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
@@ -721,6 +721,20 @@ 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):
+    class Foo(nn.Module):
+      bar: nn.Dense = dataclasses.field(init=False)
+
+      def setup(self):
+        self.i = 1
+      
+      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))
+
 
 if __name__ == '__main__':
   absltest.main()