Repository: google/flax

Issue:
bind() - unbind() returns a different module

### System information
- Linux Ubuntu 16.04
- flax 0.7.2
- jax 0.4.14
- jaxlib 0.4.14+cuda11.cudnn86

### Problem you have encountered:

Perform a bind() unbind() round-trip returns a new module unequal the original one.

### What you expected to happen:

This might not be considered a bug, but still an surprise to the user.

Maybe document the reason behind this behavior.

### Steps to reproduce:

```
class MyModule(nn.Module):
    sub_module: nn.Dense
    
    def __call__(self, x):
        return self.sub_module(x)

module = MyModule(nn.Dense(8))
state = module.init(jax.random.PRNGKey(0), np.zeros([8]))
module_new, _ = module.bind(state).unbind()
module == module_new

>>> False

```


Discussion:

Resolving patch:
diff --git a/flax/core/scope.py b/flax/core/scope.py
--- a/flax/core/scope.py
+++ b/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
 
 
diff --git a/flax/linen/module.py b/flax/linen/module.py
--- a/flax/linen/module.py
+++ b/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


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
@@ -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):