Repository: google/flax

Issue:
Updating subtree with `put_variable` doesn't update sub-scopes' references.
There are rare cases where we want to manually mess with the tree of variables at some point in a model.

If we try to use `get_variable` and `put_variable` to directly modify the variables in a collection, this works _locally_ for variables within a module, however if we try to mess with a sub-module's variables from a parent module, the mutation applied to the outer scope doesn't propagate into the sub-scope's references.

This can be illustrated by the example:

```python
from flax import linen as nn
from jax import random, numpy as jnp

class A(nn.Module):
  def setup(self):
   self.foo = self.param('foo', nn.initializers.zeros, x.shape)
  def dummy(self):
    return None
  def __call__(self, x):
    print(self.foo)  # == [0.] !!
    return x + self.foo

class B(nn.Module):
  @nn.compact
  def __call__(self, x):
    a = A(name="a")
    # trigger setup
    a.dummy()
    # fetch variables under 'a' in params collection
    vs = self.get_variable('params', 'a')
    # update this subtree
    new_vs = jax.tree_map(lambda x: jnp.ones_like(x), vs)
    self.put_variable('params', 'a', new_vs)                
    # now run call and return 
    return a(x)

k = random.PRNGKey(0)
x = jnp.zeros((1,))
y, vs = B().init_with_output(k, x)
y   # DeviceArray([0.], dtype=float32)  # <-- "wrong"
vs  # FrozenDict({'params': {'a': {'foo': DeviceArray([1.], dtype=float32),}}})
```


Discussion:
minimal repro
```python
def test(scope):
  subscope = scope.push(name="a")
  subscope.put_variable('cache', 'x', jnp.array(0.0, jnp.float32))
  # doesn't update subscope._variables but overwrites ref, leaving a "dangling" subscope
  scope.put_variable('cache', 'a', {'x': jnp.array(1.0, jnp.float32)})
  assert scope.variables()['cache']['a']['x'] == subscope.variables()['cache']['x']
k = random.PRNGKey(0)
_, vs = flax.core.init(test)(k)
```

Resolving patch:
diff --git a/flax/core/scope.py b/flax/core/scope.py
--- a/flax/core/scope.py
+++ b/flax/core/scope.py
@@ -693,7 +693,15 @@ def put_variable(self, col: str, name: str, value: Any):
     if not self.is_mutable_collection(col):
       raise errors.ModifyScopeVariableError(col, name, self.path_text)
     variables = self._mutable_collection(col)
-    variables[name] = value
+    # Make sure reference sharing of child variable dictionaries isn't broken
+    def put(target, key, val):
+      if key in target and isinstance(target[key], dict) and isinstance(val, Mapping):
+        for k, v in val.items():
+          put(target[key], k, v)
+      else:
+        target[key] = val
+
+    put(variables, name, value)
 
   def variable(self, col: str, name: str,  # pylint: disable=keyword-arg-before-vararg
                init_fn: Optional[Callable[..., T]] = None,


Test patch:
diff --git a/tests/core/core_lift_test.py b/tests/core/core_lift_test.py
--- a/tests/core/core_lift_test.py
+++ b/tests/core/core_lift_test.py
@@ -190,8 +190,24 @@ def c_fn(scope, x):
     vars = vars.copy(updates)
     self.assertEqual(vars['state'].unfreeze(), {'a_count': 1, 'b_count': 1, 'c_count': 1})
     np.testing.assert_allclose(y1, y3)
-
-
+  
+  def test_subscope_var_aliasing(self):
+    def test(scope, x):
+      subscope = scope.push(name="a")
+      subscope.put_variable('state', 'x', 0.)
+      _ = lift.while_loop(
+        lambda scope, x: False, 
+        lambda scope, x: x, 
+        scope,
+        jnp.array(0, jnp.int32),
+        carry_variables=['state'],
+      )
+      subscope.put_variable('state', 'x', 1.)
+      val0 = scope.variables()['state']['a']['x']
+      val1 = subscope.variables()['state']['x']
+      self.assertEqual(val0, val1)
+      return x
+    init(test)( random.PRNGKey(0), 1.)
 
 
 if __name__ == '__main__':
diff --git a/tests/core/core_scope_test.py b/tests/core/core_scope_test.py
--- a/tests/core/core_scope_test.py
+++ b/tests/core/core_scope_test.py
@@ -209,6 +209,13 @@ def test_variable_no_init(self):
     self.assertEqual(abc.value, 1)
     with self.assertRaises(errors.ScopeVariableNotFoundError):
       root.variable('state', 'test')
+  
+  def test_variable_alias(self):
+    scope = Scope({}, mutable='state')
+    subscope = scope.push(name="a")
+    subscope.put_variable('state', 'x', 0.)
+    scope.put_variable('state', 'a', {'x': jnp.array(1., jnp.float32)})
+    self.assertEqual(scope.variables()['state']['a']['x'], subscope.variables()['state']['x'])
 
 
 if __name__ == '__main__':