Repository: google/flax

Issue:
unbind API
As suggested in https://github.com/google/flax/pull/2394#discussion_r978744557, its currently not clear for users that `clone` produces an unbounded copy of the module. To get hold of a submodule and its variables you currently have to use this pattern:

```python
module, variables = load_module()
bounded_module = module.bind(variables)
submodule = bounded_module.some_submodule.clone()
submodule_variables = bounded_module.variables
```
The new proposed `.unbind()` method provide more clarity on how to do this process:

```python
module, variables = load_module()
submodule, submodule_variables = module.bind(variables).some_submodule.unbind()
```
Signature could be:
```python
unbind(self, **clone_kwargs) -> Tuple[Module, FrozenVariableDict]
```

cc @jheek 


Discussion:
How about `unbind()`?
@cgarciae Just so you won't be confused, I renamed the issue "snapshot API" --> "unbind API"
Perfect, edited the original comment with the new name.

Resolving patch:
diff --git a/flax/errors.py b/flax/errors.py
--- a/flax/errors.py
+++ b/flax/errors.py
@@ -606,6 +606,30 @@ def get_submodule(module):
   def __init__(self):
     super().__init__('Can\'t call compact methods on unbound modules')
 
+class CallUnbindOnUnboundModuleError(FlaxError):
+  """This error occurs when you are trying to call ``.unbind()`` on an unbound
+  Module. For instance, when you try running the following example,
+  an error will be raised::
+
+    from flax import linen as nn
+
+    class MyModule(nn.Module):
+      @nn.compact
+      def __call__(self, x):
+        return nn.Dense(features=10)(x)
+
+    module = MyModule()
+    module.unbind() # <-- ERROR!
+
+  Instead, you should ``bind`` the Module to a variable collection before calling
+  ``.unbind()``::
+
+    bound_module = module.bind(variables)
+    ... # do something with bound_module
+    module = bound_module.unbind() # <-- OK!
+  """
+  def __init__(self):
+    super().__init__('Can\'t call `unbind()` on unbound modules')
 
 class InvalidInstanceModuleError(FlaxError):
   """This error occurs when you are trying to call `.init()`, `.init_with_output()`, `.apply() or `.bind()`
diff --git a/flax/linen/module.py b/flax/linen/module.py
--- a/flax/linen/module.py
+++ b/flax/linen/module.py
@@ -53,6 +53,7 @@
 
 T = TypeVar('T')
 K = TypeVar('K')
+M = TypeVar('M', bound='Module')
 _CallableT = TypeVar('_CallableT', bound=Callable)
 
 
@@ -1089,9 +1090,9 @@ def _name_taken(self,
   def _initialization_allowed(self):
     return self._state.in_setup or self._state.in_compact_method
 
-  def clone(self, *,
+  def clone(self: M, *,
             parent: Optional[Union[Scope, 'Module']] = None,
-            **updates) -> 'Module':
+            **updates) -> M:
     """Creates a clone of this Module, with optionally updated arguments.
 
     Args:
@@ -1261,11 +1262,11 @@ def _module_checks(self):
       raise errors.IncorrectPostInitOverrideError()
 
   @traceback_util.api_boundary
-  def bind(self,
+  def bind(self: M,
            variables: VariableDict,
            *args,
            rngs: Optional[RNGSequences] = None,
-           mutable: CollectionFilter = False):
+           mutable: CollectionFilter = False) -> M:
     """Creates an interactive Module instance by binding variables and RNGs.
 
     ``bind`` provides an "interactive" instance of a Module directly without
@@ -1321,6 +1322,42 @@ def __call__(self, x):
     scope = core.bind(variables, rngs=rngs, mutable=mutable)
     return self.clone(parent=scope)
 
+  def unbind(self: M) -> Tuple[M, VariableDict]:
+    """Returns an unbound copy of a Module and its variables.
+
+    ``unbind`` helps create a stateless version of a bound Module.
+
+    An example of a common use case: to extract a sub-Module defined inside
+    ``setup()`` and its corresponding variables: 1) temporarily ``bind`` the parent
+    Module; and then 2) ``unbind`` the desired sub-Module. (Recall that ``setup()``
+    is only called when the Module is bound.)::
+
+      class AutoEncoder(nn.Module):
+        def setup(self):
+          self.encoder = Encoder()
+          self.decoder = Decoder()
+
+        def __call__(self, x):
+          return self.decoder(self.encoder(x))
+
+      module = AutoEncoder()
+      variables = module.init(jax.random.PRNGKey(0), jnp.ones((1, 784)))
+      ...
+      # Extract the Encoder sub-Module and its variables
+      encoder, encoder_vars = module.bind(variables).encoder.unbind()
+
+    Returns:
+      A tuple with an unbound copy of this Module and its variables.
+    """
+    Module._module_checks(self)
+
+    if self.scope is None:
+      raise errors.CallUnbindOnUnboundModuleError()
+
+    variables = self.variables
+    module = self.clone()
+    return module, variables
+
   @traceback_util.api_boundary
   def apply(self,
             variables: VariableDict,


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
@@ -1518,6 +1518,32 @@ def f(foo, x):
         jax.tree_util.tree_leaves(bs_1), jax.tree_util.tree_leaves(bs_2)):
       np.testing.assert_allclose(x, y)
 
+  def test_unbind(self):
+
+    class Foo(nn.Module):
+      def setup(self):
+        self.encoder = nn.Dense(4)
+        self.decoder = nn.Dense(2)
+
+      def __call__(self, x):
+        x = self.encoder(x)
+        return self.decoder(x)
+
+    foo = Foo()
+    x = jnp.ones((2,))
+
+    variables = foo.init(random.PRNGKey(0), x)
+    encoder, encoder_vars = foo.bind(variables).encoder.unbind()
+    decoder, decoder_vars = foo.bind(variables).decoder.unbind()
+
+    self.assertIsInstance(encoder, nn.Dense)
+    self.assertEqual(encoder.features, 4)
+    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'])
+
   def test_passing_mutable_variables(self):
 
     class Foo(nn.Module):