Repository: google/flax

Issue:
Unexpected ScopeParamNotFoundError when sharing submodules
### System information
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Linux Ubuntu 20.04
- Flax, jax, jaxlib versions (obtain with `pip show flax jax jaxlib`: flax 0.6.5, jax 0.4.3, jaxlib 0.4.3+cuda11.cudnn82
- Python version: 3.10.8
- GPU/TPU model and memory: NVIDIA GeForce GTX 1080ti
- CUDA version (if applicable): 11.7


### Problem you have encountered:
I encounter a `ScopeParamNotFoundError` when a shared submodule is nested under an unshared submodule. Here is a minimal working example:
```python
class Shared(nn.Module):
    @nn.compact
    def __call__(self, x):
        return nn.Dense(1)(x)

class Unshared(nn.Module):
    shared: nn.Module

    def __call__(self, x):
        return self.shared(x)

class Super(nn.Module):
    a: nn.Module
    b: nn.Module

    def run_a(self, x):
        return self.a(x)

    def run_b(self, x):
        return self.b(x)
    
    def __call__(self, x):
        return self.a(x) + self.b(x)


sh = Shared()
a = Unshared(shared=sh)
b = Unshared(shared=sh)
super = Super(a=a, b=b)

rng = jax.random.PRNGKey(0)
params = super.init(rng, jnp.ones(1))["params"]

print(params)
# FrozenDict({
#   a: {
#       shared: {
#            Dense_0: {
#                kernel: Array([[0.13082343]], dtype=float32),
#                bias: Array([0.], dtype=float32),
#            },
#        },
#    },
# })
# 
# as expected, the weights for `shared` only appear under the `a` key, since by my understanding flax caches the
# `shared` object and re-uses the same scope when it's encountered for the second time

print(super.apply({"params": params}, jnp.ones(1)))  # works as expected
print(super.apply({"params": params}, jnp.ones(1), method="run_a"))  # works as expected
print(super.apply({"params": params}, jnp.ones(1), method="run_b"))  # ScopeParamNotFoundError: Could not find parameter named "kernel" in scope "/b/shared/Dense_0"
```

### What you expected to happen:
When I call `apply` with `method="run_b"`, I expect flax to realize that the weights for `/b/shared` are under the path `/a/shared`, just like it does when `__call__` is used, but clearly it is looking in the wrong scope.

Edit: After reading a bit more, I'm thinking this has something to do with lazy binding? When `__call__` is used, first `a.shared` is bound, and then when `b.shared` is bound, flax finds the scope `a/shared` in the cache and so re-uses it for `b.shared`. However, when `run_b` is used, `a.shared` is never bound so the scope `a/shared` is never created.

Is there any way around this, perhaps using some sort of tricks in `setup`? I think there's a pretty strong argument for this use-case: it's nice to have a wrapper around multiple submodules where the submodules all get initialized at once (using `__call__`) but then individual subsets of modules can be used later on, potentially in a more complex way. And I don't see any reason why these submodules shouldn't be able to share parameters.



Discussion:
Hey @kvablack, thanks for reporting this!

I am fixing a related issue in #2748, one of the edge cases I found recently was related to lazy binding, as you suggest this is indeed the problem you are facing here. The solution I am currently exploring is to eagerly bind all sub-modules that are passed as inputs as soon as a scope is available, I suspect this will solve your issue.

I'll be adding the snippet you shared as a test case.

Resolving patch:
diff --git a/flax/linen/module.py b/flax/linen/module.py
--- a/flax/linen/module.py
+++ b/flax/linen/module.py
@@ -490,6 +490,10 @@ def _get_unbound_fn(method_or_fn: Callable[..., Any]) -> Callable[..., Any]:
 
   return method_or_fn
 
+def _map_submodules(fn: Callable[['Module'], Any], tree):
+  """Map a function over all submodules in a tree."""
+  g = lambda _, x: fn(x) if isinstance(x, Module) else x
+  return _freeze_attr(_map_over_modules_in_tree(g, tree))
 
 class SetupState(enum.IntEnum):
   # setup() has not been called.
@@ -1000,6 +1004,12 @@ def __post_init__(self) -> None:
     else:
       raise ValueError('parent must be None, Module or Scope')
 
+    # eagerly bind submodules if scope is available
+    if self.scope is not None:
+        for field in dataclasses.fields(self):
+          if field.name not in ('parent', 'name') and field.init:
+            self._register_submodules(field.name, getattr(self, field.name))
+
     self._state.is_initialized = True
 
   def __repr__(self) -> str:
@@ -1048,8 +1058,8 @@ def _register_submodules(self, name, val):
     _caches[root] = cache
     queue = []
     preserve_adopted_names = config.flax_preserve_adopted_names
-    if hasattr(self, 'preserve_adopted_names'):
-      preserve_adopted_names = self.preserve_adopted_names
+    if hasattr(type(self), 'preserve_adopted_names'):
+      preserve_adopted_names = type(self).preserve_adopted_names
     def adopt_attr_modules(cache, queue, suffix, subvalue):
       if isinstance(subvalue, Module):
         adopted_name = None
@@ -1092,7 +1102,7 @@ def _try_setup(self, shallow: bool = False) -> None:
         # not call the user's setup. This avoids running before a
         # transformation.
         for field in dataclasses.fields(self):
-          if field.name != 'parent' and field.init:
+          if field.name not in ('parent', 'name') and field.init:
             self._register_submodules(field.name, getattr(self, field.name))
         if not shallow:
           self.setup()
@@ -1137,23 +1147,61 @@ def _name_taken(self,
 
   @property
   def _initialization_allowed(self):
-    return self._state.in_setup or self._state.in_compact_method
+    return (not self._state.is_initialized  # allow eager attachment in post-init
+            or self._state.in_setup
+            or self._state.in_compact_method)
 
   def clone(self: M, *,
             parent: Optional[Union[Scope, 'Module']] = None,
+            _deep_clone: Union[bool, weakref.WeakValueDictionary] = False,
             **updates) -> M:
     """Creates a clone of this Module, with optionally updated arguments.
 
     Args:
       parent: The parent of the clone. The clone will have no parent if no
         explicit parent is specified.
+      _deep_clone: A boolean or a weak value dictionary to control deep cloning
+        of submodules. If True, submodules will be cloned recursively. If a
+        weak value dictionary is passed, it will be used to cache cloned
+        submodules. This flag is used by init/apply/bind to avoid scope
+        leakage.
       **updates: Attribute updates.
     Returns:
       A clone of the this Module with the updated attributes and parent.
     """
     attrs = {f.name: getattr(self, f.name) for f in dataclasses.fields(self) if f.init}
+
     attrs.update(parent=parent, **updates)
-    return self.__class__(**attrs)
+
+    # Here we implement deep cloning of submodules, this is necessary to avoid scope leakage
+    # from external submodules into init/apply/bind while preserving sharing-by-reference
+    # relationships between submodules.
+    if _deep_clone != False:
+      # We use a weak value dictionary to cache cloned submodules. When a shared
+      # submodule is cloned, its only cloned once else its fetched from the cache.
+      cache = weakref.WeakValueDictionary() if isinstance(_deep_clone, bool) else _deep_clone
+      def clone_fn(m: Module) -> Module:
+        if hasattr(m, '_id'):
+          key = m._id
+          if key in cache:
+            return cache[key]
+          else:
+            clone = m.clone(_deep_clone=cache)
+            cache[key] = clone
+            return clone
+        else:
+          # If the module doesn't have an _id attribute it could be a mock object
+          # so we return it as is.
+          return m
+
+      # _map_submodules will map over all submodules inside attrs
+      # value here can be any pytree, non-module values are ignored
+      for field_name, value in attrs.items():
+        attrs[field_name] = _map_submodules(clone_fn, value)
+
+    module = self.__class__(**attrs)
+
+    return module
 
   def variable(self, col: str, name: str,
                init_fn: Optional[Callable[..., Any]] = None,
@@ -1371,7 +1419,7 @@ def __call__(self, x):
 
     del args
     scope = core.bind(variables, rngs=rngs, mutable=mutable)
-    return self.clone(parent=scope)
+    return self.clone(parent=scope, _deep_clone=True)
 
   def unbind(self: M) -> Tuple[M, VariableDict]:
     """Returns an unbound copy of a Module and its variables.
@@ -2057,7 +2105,7 @@ def f(foo, x):
   def scope_fn(scope, *args, **kwargs):
     _context.capture_stack.append(capture_intermediates)
     try:
-      return fn(module.clone(parent=scope), *args, **kwargs)
+      return fn(module.clone(parent=scope, _deep_clone=True), *args, **kwargs)
     finally:
       _context.capture_stack.pop()
 
@@ -2118,7 +2166,7 @@ def f(foo, x):
   def scope_fn(scope, *args, **kwargs):
     _context.capture_stack.append(capture_intermediates)
     try:
-      return fn(module.clone(parent=scope), *args, **kwargs)
+      return fn(module.clone(parent=scope, _deep_clone=True), *args, **kwargs)
     finally:
       _context.capture_stack.pop()
 


Test patch:
diff --git a/tests/linen/linen_module_test.py b/tests/linen/linen_module_test.py
--- a/tests/linen/linen_module_test.py
+++ b/tests/linen/linen_module_test.py
@@ -2000,6 +2000,180 @@ def __call__(self):
                                 'Trying to access a property that'):
       foo.apply({})
 
+  def test_nested_external_modules(self):
+    class Baz(nn.Module):
+      a: int
+
+      def setup(self):
+        self.b = self.param('b', lambda k: 2)
+
+      def __call__(self, x):
+        return x + self.a * self.b
+
+    class Bar(nn.Module):
+      baz: Baz
+
+      def __call__(self, x):
+        return self.baz(x)
+
+    class Foo(nn.Module):
+      def setup(self):
+        self.bar = Bar(baz=Baz(a=1))
+
+      def __call__(self, x):
+        return self.bar.baz(x)
+
+    module = Foo()
+    y, variables = module.init_with_output(jax.random.PRNGKey(0), 1)
+    self.assertEqual(y, 3)
+
+  def test_getattribute_triggers_setup(self):
+    class B(nn.Module):
+      def setup(self):
+        self.p1 = self.param('p1', lambda k: jnp.ones((2,)))
+      def fn1(self, x):
+        return self.p1 + x
+    class A(nn.Module):
+      b: nn.Module
+      def __call__(self, x):
+        return self.b.fn1(x)
+    a = A(b=B())
+    k = random.PRNGKey(0)
+    x = jnp.zeros((2,))
+    vs = nn.init(lambda a,x: a(x), a)(k, x)
+    y = nn.apply(lambda a,x: a.b.fn1(x), a)(vs, x)
+    np.testing.assert_array_equal(y, jnp.ones((2,)))
+
+  def test_nested_sequential_in_call(self):
+    class Foo(nn.Module):
+      def setup(self):
+        self.seq = nn.Sequential([nn.Dense(10) for i in range(10)])
+
+      def __call__(self, x):
+        # try calling only the first layer
+        return self.seq.layers[0](x)
+
+    module = Foo()
+    variables = module.init(jax.random.PRNGKey(0), jnp.ones((1, 10)))
+
+  def test_setup_called_bounded_submodules(self):
+    module = nn.Sequential([
+      nn.Sequential([
+        nn.Dense(2),
+        nn.relu,
+        nn.Dense(2),
+      ]),
+      nn.relu,
+      nn.Dense(2),
+    ])
+    x = jnp.ones((1, 3))
+    variables = module.init(jax.random.PRNGKey(0), x)
+    bound_module = module.bind(variables)
+
+    self.assertIsNotNone(bound_module.layers[0].layers[0].scope)
+    self.assertIsNotNone(bound_module.layers[0].layers[2].scope)
+    self.assertIsNotNone(bound_module.layers[2].scope)
+
+  def test_call_bounded_toplevel_mutable(self):
+    class Bar(nn.Module):
+      a: int
+
+      def setup(self):
+        self.b = self.param('b', lambda k: 1)
+
+      def __call__(self, x):
+        return x + self.a * self.b
+
+    class Foo(nn.Module):
+      bars: Sequence[Bar]
+
+      def __call__(self, x):
+        for bar in self.bars:
+          x = bar(x)
+        return x
+
+
+    module = Foo(bars=[])
+    module.bars = [Bar(a=1)]
+
+    variables = module.init(jax.random.PRNGKey(0), jnp.ones(()))
+    bound_module = module.bind(variables)
+
+    bar1 = bound_module.bars[0]
+    self.assertIsNotNone(bar1.scope)
+
+  def test_nested_init(self):
+    class Baz(nn.Module):
+      a: int
+
+      def setup(self):
+        self.b = self.param('b', lambda k: jnp.ones(()))
+
+      def __call__(self, x):
+        return x + self.a * self.b
+
+    class Bar(nn.Module):
+      baz: Baz
+
+      def setup(self):
+        a = 1
+
+      def __call__(self, x):
+        return self.baz(x)
+
+    class Foo(nn.Module):
+
+      def setup(self):
+        self.bar: Bar = Bar(baz=Baz(a=1))
+
+      def __call__(self, x):
+        # y = self.bar(x)
+        y, bar_vars = self.bar.init_with_output(jax.random.PRNGKey(0), x)
+        return y, bar_vars
+
+    # create foo
+    module = Foo()
+
+    # run foo
+    (y, bar_vars), variables = module.init_with_output(
+      jax.random.PRNGKey(0), jnp.ones(()))
+
+    self.assertIn('params', bar_vars)
+
+  def test_nested_shared(self):
+    class Shared(nn.Module):
+      @nn.compact
+      def __call__(self, x):
+        return nn.Dense(1)(x)
+
+    class Unshared(nn.Module):
+      shared: nn.Module
+      def __call__(self, x):
+        return self.shared(x)
+
+    class Super(nn.Module):
+      a: nn.Module
+      b: nn.Module
+      def run_a(self, x):
+        return self.a(x)
+      def run_b(self, x):
+        return self.b(x)
+      def __call__(self, x):
+        return self.a(x) + self.b(x)
+
+
+    sh = Shared()
+    a = Unshared(shared=sh)
+    b = Unshared(shared=sh)
+    module = Super(a=a, b=b)
+
+    rng = jax.random.PRNGKey(0)
+    params = module.init(rng, jnp.ones(1))["params"]
+
+    module.apply({"params": params}, jnp.ones(1))  # works as expected
+    module.apply({"params": params}, jnp.ones(1), method="run_a")  # works as expected
+    module.apply({"params": params}, jnp.ones(1), method="run_b")  # ScopeParamNotFoundError: Could not find parameter named "kernel" in scope "/b/shared/Dense_0"
+
   def test_repr(self):
 
     class Base1(nn.Module):