fixed bind-unbind bug

flax/core/scope.py

@@ -1006,7 +1006,7 @@ def _unfreeze_variables(variables, mutable):
     if in_filter(mutable, key):
       new_variables[key] = unfreeze(value)
     else:
-      new_variables[key] = freeze(value)
+      new_variables[key] = value
   return new_variables
 
 

flax/linen/module.py

@@ -1432,6 +1432,7 @@ def clone(
       *,
       parent: Optional[Union[Scope, 'Module']] = None,
       _deep_clone: Union[bool, weakref.WeakValueDictionary] = False,
+      _reset_names: bool = False,
       **updates,
   ) -> M:
     """Creates a clone of this Module, with optionally updated arguments.
@@ -1443,6 +1444,8 @@ def clone(
         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.
+      _reset_names: If True, `name=None` is also passed to submodules when
+        cloning. Resetting names in submodules is necessary when calling `.unbind`.
       **updates: Attribute updates.
 
     Returns:
@@ -1474,7 +1477,10 @@ def clone_fn(m: Module) -> Module:
           if key in cache:
             return cache[key]
           else:
-            clone = m.clone(_deep_clone=cache)
+            if _reset_names:
+              clone = m.clone(_deep_clone=cache, _reset_names=_reset_names, name=None)
+            else:
+              clone = m.clone(_deep_clone=cache)
             cache[key] = clone
             return clone
         else:
@@ -1825,7 +1831,7 @@ def __call__(self, x):
       raise errors.CallUnbindOnUnboundModuleError()
 
     variables = self.variables
-    module = self.clone()
+    module = self.clone(_deep_clone=True, _reset_names=True, name=None)
     return module, variables
 
   @traceback_util.api_boundary

tests/linen/linen_module_test.py

@@ -1622,12 +1622,27 @@ def __call__(self, x):
     self.assertIsInstance(decoder, nn.Dense)
     self.assertEqual(decoder.features, 2)
 
-    np.testing.assert_equal(
-        variables['params']['encoder'], encoder_vars['params']
-    )
-    np.testing.assert_equal(
-        variables['params']['decoder'], decoder_vars['params']
-    )
+    self.assertTrue(jax.tree_util.tree_all(jax.tree_map(lambda v1, v2: (v1 == v2).all(), variables['params']['encoder'], encoder_vars['params'])))
+    self.assertTrue(jax.tree_util.tree_all(jax.tree_map(lambda v1, v2: (v1 == v2).all(), variables['params']['decoder'], decoder_vars['params'])))
+
+  def test_bind_unbind_equality(self):
+    class Foo(nn.Module):
+        sub_module: Any
+        @nn.compact
+        def __call__(self, x):
+            x = nn.Dense(2)(x)
+            return self.sub_module(x)
+
+    sub_module = Foo(nn.Dense(3))
+    module = Foo(sub_module)
+    x = jnp.ones((1, 2))
+    variables = module.init(jax.random.PRNGKey(0), x)
+
+    bound_module = module.bind(variables)
+    self.assertTrue((module.apply(variables, x) == bound_module(x)).all())
+    new_module, new_variables = bound_module.unbind()
+    self.assertTrue(jax.tree_util.tree_all(jax.tree_map(lambda v1, v2: (v1 == v2).all(), variables, new_variables)))
+    self.assertEqual(module, new_module)
 
   def test_passing_mutable_variables(self):
     class Foo(nn.Module):