Repository: google/flax

Issue:
Directly calling `module.setup()` should raise an exception
`flax.linen.Module.setup()` should not be called directly because it needs a `flax.linen.Scope` to be set up properly.

Since #653 there is no more exception risen when a user inadvertently calls `flax.linen.Module.setup()` (though there probably will be error messages if the user tries to access the scope...)


Discussion:
An incomplete thought below:

Curiously enough, this recent [discussion](https://github.com/google/flax/discussions/665#discussioncomment-136656) made me think about this and wonder... For most Flax modules (that ultimately define parameters in a compact method), if you define submodules in `setup` but not parameters, I think it may be safe to call `setup` before a module is bound.

In fact, right now we have weird behavior where in Roland's example you can't run `ResNet().backbone`, even though it should probably work just fine (and still be an unbound module).

I'm not yet sure what the best answer is, but perhaps a workaround for Roland for now is actually to explicitly do:

```
resnet = ResNet()
resnet.setup()
resnet.backbone  # does this work?
```

### Update
Currently @avital's example does not work:
```python
resnet = ResNet()
resnet.setup()
resnet.backbone  # does this work?
```
You now get the following error:
```
AssertionError: Trying to register submodules on unbound scope.
```
This behaviour is documented in [The Module lifecycle guide](https://flax.readthedocs.io/en/latest/advanced_topics/module_lifecycle.html#setup). We could go ahead and raise an error if `setup` is called on unbounded Module and add a message with some pointers into how to correctly extract submodules defined in `setup`.

### Proposal
We could recommend the use of `bind` for this use-case, even a bind with empty variables could work:

```python
resnet = ResNet()
backbone = resnet.bind({}).backbone.clone()
```
We could document this pattern in an existing guide or add a short guide about this topic of "Accessing Submodules".

#### Future ideas

If this pattern is sound we could make it a bit more user-friendly in the future via a `.submodules` that would automate the previous:

```python
resnet = ResNet()
backbone = resnet.submodules.backbone
```


Resolving patch:
diff --git a/flax/errors.py b/flax/errors.py
--- a/flax/errors.py
+++ b/flax/errors.py
@@ -538,6 +538,37 @@ class CallCompactUnboundModuleError(FlaxError):
   def __init__(self):
     super().__init__('Can\'t call compact methods on unbound modules')
 
+class CallSetupUnboundModuleError(FlaxError):
+  """
+  This error occurs when you are trying to call `.setup()` directly. For instance, the
+  error will be raised when trying to run this code::
+
+    from flax import linen as nn
+    import jax.numpy as jnp
+
+    class MyModule(nn.Module):
+      def setup(self):
+        self.submodule = MySubModule()
+
+    module = MyModule()
+    module.setup() # <-- ERROR!
+    submodule = module.submodule
+
+  In general you shouldn't call `.setup()` yourself, if you need to get access
+  to a field or submodule defined inside `setup` you can instead create a function
+  to extract it and pass it to `nn.apply`::
+
+    # setup() will be called automatically by `nn.apply`
+    def get_submodule(module):
+      return module.submodule.clone() # avoid leaking the Scope
+
+    empty_variables = {} # you can also use the real variables
+    submodule = nn.apply(get_submodule, module)(empty_variables)
+
+  """
+  def __init__(self):
+    super().__init__('Can\'t call compact methods on unbound modules')
+
 
 class InvalidCheckpointError(FlaxError):
   """
diff --git a/flax/linen/module.py b/flax/linen/module.py
--- a/flax/linen/module.py
+++ b/flax/linen/module.py
@@ -681,6 +681,8 @@ def _call_wrapped_method(self, fun, args, kwargs):
     add_call_info = not is_setup_method and len(_context.call_info_stack) > 0
     # We lazily call setup() only when needed.
     if is_setup_method:
+      if self.scope is None:
+        raise errors.CallSetupUnboundModuleError()
       is_recurrent = self._state.in_setup
       self._state.in_setup = True
     else:


Test patch: