Repository: google/flax

Issue:
Improve Error Message: Naming a module in setup
The error `TypeError: Module instance is frozen outside of setup method.` is thrown  at any time you assign module attributes somewhere other than within setup. It is not always clear to users how to resolve this.

-- Example 1

This code
```
def setup(self):
  self.layer = nn.Dense(...)
  self.layer.name = 'dense'
```

Throws the error `TypeError: Module instance is frozen outside of setup method.`, which confuses users.

-- Example 2
See: #936


Discussion:
Similarly, users can be confused about the error message `ValueError: In setup, assign names of Modules via self.<name> and not using keyword argument name="<name>"`. We should provide a more elaborate error message with an example.
A bit more color here:

For submodules defined in-line within a `@nn.compact` method, you either explicitly pass names to submodules via a `name` argument to the submodule constructor, or they are automatically generated if not, e.g.:

```py
# ... inside a module
@nn.compact
def func(self, x):
  dense1 = Dense(features=16)  # submodule name autoassigned to "Dense1"
  dense2 = Dense(features=16, name='final')  # submodule name is "final"
```

For submodules defined inside `setup`, names are always explicit and are derived from the name of the attribute on which they are assigned (via `__setattr__`, following a very similar logic to that of PyTorch):
```py
# ... inside a module
def setup(self):
   self.final = Dense(features=16)  # submodule name is "final"
```

This issue is very closely related to #524.
In short, as summarized by @salayatana66, "really in `setup` the attribute name is the name of the module."

Resolving patch:
diff --git a/flax/core/scope.py b/flax/core/scope.py
--- a/flax/core/scope.py
+++ b/flax/core/scope.py
@@ -377,9 +377,10 @@ def reserve(self, name: str):
       name: the name to reserve.
     """
     if not isinstance(name, str):
-      raise errors.ScopeNameTypeError(name)
+      raise TypeError('The type of scope "{name}" should be string but '
+                     f'it is {type(name)}')
     if name in self.reservations:
-      raise errors.ScopeNameInUseError(name)
+      raise ValueError(f'Duplicate use of scope name: "{name}"')
     self.reservations.add(name)
 
   def default_name(self, prefix: str) -> str:
@@ -502,7 +503,8 @@ def has_rng(self, name: str) -> bool:
 
   def make_rng(self, name: str) -> PRNGKey:
     """Generates A PRNGKey from a PRNGSequence with name `name`."""
-    assert self.has_rng(name), f'Need PRNG for "{name}"'
+    if not self.has_rng(name):
+      raise errors.InvalidRngError(f'{self.name} needs PRNG for "{name}"')
     self._check_valid()
     self._validate_trace_level()
     self.rng_counters[name] += 1
@@ -649,7 +651,8 @@ def bind(variables: VariableDict,
   if not _is_valid_variables(variables):
     raise errors.ApplyScopeInvalidVariablesError()
   if rngs is not None and not _is_valid_rngs(rngs):
-    raise errors.ApplyScopeInvalidRngsError()
+    raise errors.InvalidRngError(
+      'rngs should be a dictionary mapping strings to `jax.PRNGKey`.')
   new_variables = _unfreeze_variables(variables, mutable)
   return Scope(new_variables, rngs=rngs, mutable=mutable)
 
@@ -696,7 +699,9 @@ def init(fn: Callable[..., Any],
   @functools.wraps(fn)
   def wrapper(rngs, *args, **kwargs) -> Tuple[Any, VariableDict]:
     if not _is_valid_rng(rngs) and not _is_valid_rngs(rngs):
-      raise errors.InitScopeInvalidRngsError()
+      raise ValueError('First argument passed to an init function should be a '
+                       '`jax.PRNGKey` or a dictionary mapping strings to '
+                       '`jax.PRNGKey`.')
     if not isinstance(rngs, dict):
       rngs = {'params': rngs}
     return apply(fn, mutable=mutable)({}, *args, rngs=rngs, **kwargs)
diff --git a/flax/errors.py b/flax/errors.py
--- a/flax/errors.py
+++ b/flax/errors.py
@@ -69,77 +69,76 @@ def __init__(self, message):
     error_msg = f'{message} ({error_page}#{module_name}.{class_name})'
     super().__init__(error_msg)
 
+
 #################################################
 # scope.py errors                               #
 #################################################
 
-class InitScopeInvalidRngsError(FlaxError):
-  """
-  When initializing a Module with
-  :meth:`Module.init() <flax.linen.Module.init>`, the first argument can be of
-  two forms:
 
-  1. A single PRNGKey. This is in case only one PRNGKey is needed to initialize
-     the ``params`` collection. Note that this::
+class InvalidRngError(FlaxError):
+  """
+  All rngs used in a Module should be passed to 
+  :meth:`Module.init() <flax.linen.Module.init>` and 
+  :meth:`Module.apply() <flax.linen.Module.apply>` appropriately. We explain
+  both separately using the following example::
 
-       SomeModule(...).init(jax.random.PRNGKey(0), ...)
+    class Bar(nn.Module):
+      @nn.compact
+      def __call__(self, x):
+        some_param = self.param('some_param', nn.initializers.zeros, (1, ))
+        dropout_rng = self.make_rng('dropout')
+        x = nn.Dense(features=4)(x)
+        ...
 
-     Is shorthand for::
+    class Foo(nn.Module):
+      @nn.compact
+      def __call__(self, x):
+        x = Bar()(x)
+        ...
 
-       SomeModule(...).init({'params': jax.random.PRNGKey(0)}, ...)
+  **PRNGs for Module.init()**
+  
+  In this example, two rngs are used:
 
-  2. A directionary mapping collections to the PRNGKey to initialize them with.
-     This is useful if the Module has more rngs than one for ``params``.
-     
-     For instance, suppose an ``EncoderDecoder`` Module that requires an RNG for
-     decoding tokens based on a categorical probability distribution. Then a 
-     typical call looks as follows::
+  * ``params`` is used for initializing the parameters of the model. This rng
+    is used to initialize the ``some_params`` parameter, and for initializing
+    the weights of the ``Dense`` Module used in ``Bar``.
+  
+  * ``dropout`` is used for the dropout rng that is used in ``Bar``.
 
-        EncoderDecoder(...).init({'params': rng1, 'decode': rng2}, ...)
+  So, ``Foo`` is initialized as follows::
+    
+    init_rngs = {'params': random.PRNGKey(0), 'dropout': random.PRNGKey(1)}
+    variables = Foo().init(init_rngs, init_inputs)
 
-     Note that even though they may be used inside submodules, the rngs for the
-     collections should be defined at the top-level. So the ``EncoderDecoder``
-     module above may contain a submodule ``Decoder``, which then uses the 
-     ``decode`` collection. The RNGs will be passed down to submodules
-     automatically.
-  """
-  def __init__(self):
-    super().__init__('First argument passed to an init function should be a '
-                     '`jax.PRNGKey` or a dictionary mapping strings to '
-                     '`jax.PRNGKey`.')
+  If a Module only requires an rng for ``params``, you can use::
 
+       SomeModule().init(rng, ...)  # Shorthand for {'params': rng}
 
-class ApplyScopeInvalidRngsError(FlaxError):
-  """
-  When applying a Module, the `rng` argument should be a dictionary mapping 
-  collections to the PRNGKeys that are used when computing their new values.
 
-  For instance, suppose an ``EncoderDecoder`` Module that requires an RNG for
-  decoding tokens based on a categorical probability distribution. Then a 
-  typical call to :meth:`Module.apply() <flax.linen.Module.apply>` looks as
-  follows::
+  **PRNGs for Module.apply()**
+  
+  When applying ``Foo``, only the rng for ``dropout`` is needed, because 
+  ``params`` is only used for initializing the Module parameters::
 
-     EncoderDecoder(...).apply(params, ... {'decode': rng2}, ...)
+    Foo().apply(variables, inputs, rngs={'dropout': random.PRNGKey(2)})
 
-  Remarks:
+  If a Module only requires an rng for ``params``, you don't have to provide
+  rngs for apply at all::
 
-  * While :meth:`Module.init() <flax.linen.Module.init>` requires a rngs for
-    the collection ``params``, this is not necessary when applying the module,
-    because this collection is only use to initialize the model with.
-  * Even though they may be used inside submodules, the rngs for the collections
-    should be defined at the top-level. So the ``EncoderDecoder`` module above 
-    may contain a submodule ``Decoder``, which then uses the ``decode``
-    collection. The RNGs will be passed down to submodules automatically.
+       SomeModule().apply(variables, inputs)  # rngs=None
   """
-  def __init__(self):
-    super().__init__('rngs should be a dictionary mapping strings to '
-                     '`jax.PRNGKey`.')
-                   
+  def __init__(self, msg):
+    # For this error message we pass the entire message, since there are various
+    # different kinds of RNG errors and we want to be able to be more specific
+    # in the error message, while always linking to the same documentation.
+    super().__init__(msg)
+
 
 class ApplyScopeInvalidVariablesError(FlaxError):
   """
   When calling :meth:`Module.apply() <flax.linen.Module.apply>`, the first
-  argument should be a variable dict. For more explanation on variable direct,
+  argument should be a variable dict. For more explanation on variable dicts,
   please see :mod:`flax.core.variables`.
   """
   def __init__(self):
@@ -166,11 +165,8 @@ def __call__(self, inputs, embed_name='embedding'):
                             (self.num_embeddings, self.features))    
       return embedding[inputs]
 
-  vars = Embed(4, 8).init(random.PRNGKey(0), jnp.ones((5, 5, 1)))
-  print(jax.tree_map(lambda x : x.shape, vars))
-  _ = NoBiasDense().apply(vars, jnp.ones((5, 5, 1)), 'embed')
-  
-
+    variables = Embed(4, 8).init(random.PRNGKey(0), jnp.ones((5, 5, 1)))
+    _ = NoBiasDense().apply(variables, jnp.ones((5, 5, 1)), 'embed')
   """
   def __init__(self, param_name, scope_path):
     super().__init__(f'No parameter named "{param_name}" exists in '
@@ -201,8 +197,8 @@ def __call__(self, x):
                             (((x.ndim - 1,), (0,)), ((), ())))
         return y
 
-    vars = NoBiasDense().init(random.PRNGKey(0), jnp.ones((5, 5, 1)))
-    _ = NoBiasDense().apply(vars, jnp.ones((5, 5)))
+    variables = NoBiasDense().init(random.PRNGKey(0), jnp.ones((5, 5, 1)))
+    _ = NoBiasDense().apply(variables, jnp.ones((5, 5)))
   """
   def __init__(self, param_name, scope_path, value_shape, init_shape):
     super().__init__('Inconsistent shapes between value and initializer '
@@ -214,7 +210,7 @@ class ScopeVariableNotFoundError(FlaxError):
   """
   This error is thrown when trying to use a variable in a Scope in a collection
   that is immutable. In order to create this variable, mark the collection as
-  mutable explicitly using the `mutable` keyword in
+  mutable explicitly using the ``mutable`` keyword in
   :meth:`Module.apply() <flax.linen.Module.apply>`.
   """
   def __init__(self, name, col, scope_path):
@@ -257,35 +253,38 @@ def __call__(self, x):
         var.value = ...
         ...
     
-    vars = MyModule.init(...)
+    v = MyModule.init(...)
     ...
-    logits = MyModule.apply(vars, batch)  # This throws an error.
-    logits = MyModule.apply(vars, batch, mutable=['batch_stats'])  # This works.
+    logits = MyModule.apply(v, batch)  # This throws an error.
+    logits = MyModule.apply(v, batch, mutable=['batch_stats'])  # This works.
   """
   def __init__(self, col, variable_name, scope_path):
     super().__init__(f'Cannot update variable "{variable_name}" in '
                      f'"{scope_path}" because collection "{col}" is immutable.')
 
 
-class ScopeNameTypeError(FlaxError):
-  """
-  Scope names should be strings.
-  """
-  def __init__(self, scope_name):
-    super().__init__(f'The type of scope "{scope_name}" should be string but '
-                     f'it is {type(scope_name)}')
+#################################################
+# module.py errors                              #
+#################################################
 
 
-class ScopeNameInUseError(FlaxError):
+class NameInUseError(FlaxError):
   """
-  Module names are unique within a subscope::
+  This error is raised when trying to create a submodule, param, or variable
+  with an existing name. They are all considered to be in the same namespace.
 
-    class MyModule(nn.Module):
-      @nn.compact
-      def __call__(self, x):
-        x = MySubModule(name='m1')(x)
-        x = MySubModule(name='m1')(x)  # This is not allowed.
-        return x
+  **Sharing Submodules**
+
+  This is the wrong pattern for sharing submodules::
+
+    y = nn.Dense(feature=3, name='bar')(x)
+    z = nn.Dense(feature=3, name='bar')(x+epsilon)
+
+  Instead, modules should be shared by instance::
+
+    dense = nn.Dense(feature=3, name='bar')
+    y = dense(x)
+    z = dense(x+epsilon)
 
   If submodules are not provided with a name, a unique name will be given to
   them automatically::
@@ -296,9 +295,226 @@ def __call__(self, x):
         x = MySubModule()(x)
         x = MySubModule()(x)  # This is fine.
         return x
+
+  **Parameters and Variables**
+
+  A parameter name can collide with a submodule or variable, since they are all
+  stored in the same variable dict::
+
+    class Foo(nn.Module):
+      @nn.compact
+      def __call__(self, x):
+        bar = self.param('bar', nn.initializers.zeros, (1, ))
+        embed = nn.Embed(num_embeddings=2, features=5, name='bar')  # <-- ERROR!
+  
+  Variables should also have unique names, even if they have their own
+  collection::
+
+    class Foo(nn.Module):
+      @nn.compact
+      def __call__(self, inputs):
+        _ = self.param('mean', initializers.lecun_normal(), (2, 2))
+        _ = self.variable('stats', 'mean', initializers.zeros, (2, 2))
   """
-  def __init__(self, scope_name):
-    super().__init__(f'Duplicate use of scope name: "{scope_name}"')
+  def __init__(self, key_type, value, module_name):
+    # key_type is in {param, variable, submodule}.
+    super().__init__(f'Could not create {key_type} "{value}" in Module '
+                     f'{module_name}: Name in use.')
+
+
+class AssignSubModuleError(FlaxError):
+  """
+  You are only allowed to create submodules in two places:
+
+  1.  If your Module is noncompact: inside 
+      :meth:`Module.setup() <flax.linen.Module.setup>`.
+  2.  If your Module is compact: inside the method wrapped in 
+      :meth:`nn.compact() <flax.linen.compact>`.
+
+  For instance, the following code throws this error, because ``nn.Conv`` is
+  created in ``__call__``, which is not marked as compact::
+
+    class Foo(nn.Module):
+      def setup(self):
+        pass
+
+      def __call__(self, x):
+        conv = nn.Conv(features=3, kernel_size=3)
+
+    Foo().init(random.PRNGKey(0), jnp.zeros((1,)))
+
+  Note that this error is also thrown if you partially defined a Module inside
+  setup::
+
+    class Foo(nn.Module):
+      def setup(self):
+        self.conv = functools.partial(nn.Conv, features=3)
+
+      def __call__(self, x):
+        x = self.conv(kernel_size=4)(x)
+        return x
+
+    Foo().init(random.PRNGKey(0), jnp.zeros((1,)))
+
+  In this case, ``self.conv(kernel_size=4)`` is called from ``__call__``, which
+  is disallowed beause it's neither within ``setup`` nor a method wrapped in
+  x``nn.compact``.
+  """
+  def __init__(self, cls):
+    super().__init__(f'Submodule {cls} must be defined in `setup()` or in a '
+                     'method wrapped in `@compact`')
+
+
+class SetAttributeInModuleSetupError(FlaxError):
+  """
+  You are not allowed to modify Module class attributes in
+  :meth:`Module.setup() <flax.linen.Module.setup>`::
+
+    class Foo(nn.Module):
+      features: int = 6
+
+      def setup(self):
+        self.features = 3  # <-- ERROR
+
+      def __call__(self, x):
+        return nn.Dense(self.features)(x)
+
+    variables = SomeModule().init(random.PRNGKey(0), jnp.ones((1, )))
+
+  Instead, these attributes should be set when initializing the Module::
+
+    class Foo(nn.Module):
+      features: int = 6
+
+      @nn.compact
+      def __call__(self, x):
+        return nn.Dense(self.features)(x)
+
+    variables = SomeModule(features=3).init(random.PRNGKey(0), jnp.ones((1, )))
+  
+  TODO(marcvanzee): Link to a design note explaining why it's necessary for
+  modules to stay frozen (otherwise we can't safely clone them, which we use for
+  lifted transformations).
+  """
+  def __init__(self):
+    super().__init__(f'Module construction attributes are frozen.')
+
+
+class SetAttributeFrozenModuleError(FlaxError):
+  """
+  You can only assign Module attributes to ``self`` inside
+  :meth:`Module.setup() <flax.linen.Module.setup>`. Outside of that method, the
+  Module instance is frozen (i.e., immutable). This behavior is similar to
+  frozen Python dataclasses.
+  
+  For instance, this error is raised in the following case::
+
+    class SomeModule(nn.Module):
+      @nn.compact
+      def __call__(self, x, num_features=10):
+        self.num_features = num_features  # <-- ERROR!
+        x = nn.Dense(self.num_features)(x)
+        return x
+
+    s = SomeModule().init(random.PRNGKey(0), jnp.ones((5, 5)))
+
+  Similarly, the error is raised when trying to modify a submodule's attributes
+  after constructing it, even if this is done in the ``setup()`` method of the
+  parent module::
+
+    class Foo(nn.Module):
+        def setup(self):
+          self.dense = nn.Dense(features=10)
+          self.dense.features = 20  # <--- This is not allowed
+        
+        def __call__(self, x):
+          return self.dense(x)
+  """
+  def __init__(self, module_cls, attr_name, attr_val):
+    super().__init__(f'Can\'t set {attr_name}={attr_val} for Module of type '
+                    f'{module_cls}: Module instance is frozen outside of '
+                     'setup method.')
+
+
+class MultipleMethodsCompactError(FlaxError):
+  """
+  The ``@compact`` decorator may only be added to at most one method in a Flax
+  module. In order to resolve this, you can:
+  
+  * remove ``@compact`` and define submodules and variables using 
+    :meth:`Module.setup() <flax.linen.Module.setup>`.
+  * Use two separate modules that both have a unique ``@compact`` method.
+
+  TODO(marcvanzee): Link to a design note explaining the motivation behind this.
+  There is no need for an equivalent to `hk.transparent` and it makes submodules
+  much more sane because there is no need to prefix the method names.
+  """
+  def __init__(self):
+    super().__init__(f'Only one method per class can be @compact')
+
+class ReservedModuleAttributeError(FlaxError):
+  """
+  This error is thrown when creating a Module that is using reserved attributes.
+  The following attributes are reserved:
+  
+  * ``parent``: The parent Module of this Module.
+  * ``name``: The name of this Module.
+  """
+  def __init__(self, annotations):
+    super().__init__(f'properties `parent` and `name` are reserved: '
+                     f'{annotations}')
+
+
+class ApplyModuleInvalidMethodError(FlaxError):
+  """
+  When calling :meth:`Module.apply() <flax.linen.Module.apply>`, you can specify
+  the method to apply using parameter ``method``. This error is thrown if the
+  provided parameter is not a method in the Module and not a function with at
+  least one argument.
+
+  Learn more on the reference docs for
+  :meth:`Module.apply() <flax.linen.Module.apply>`.
+  """
+  def __init__(self, method):
+    super().__init__(f'Cannot call apply(): {method} is not a valid function '
+                     'for apply().')
+
+
+class CallCompactUnboundModuleError(FlaxError):
+  """
+  This error occurs when you are trying to call a Module directly, rather than
+  through :meth:`Module.apply() <flax.linen.Module.apply>`. For instance, the 
+  error will be raised when trying to run this code::
+
+    from flax import linen as nn
+    import jax.numpy as jnp
+
+    test_dense = nn.Dense(10)
+    test_dense(jnp.ones((5,5)))
+
+  Instead, you should pass the variables (parameters and other state) via 
+  :meth:`Module.apply() <flax.linen.Module.apply>` (or use 
+  :meth:`Module.init() <flax.linen.Module.init>` to get initial variables)::
+
+    from jax import random
+    variables = test_dense.init(random.PRNGKey(0), jnp.ones((5,5)))
+
+    y = test_dense.apply(variables, jnp.ones((5,5)))
+  """
+  def __init__(self):
+    super().__init__('Can\'t call compact methods on unbound modules')
+
+
+class JaxOmnistagingError(FlaxError):
+  """
+  The Flax linen API requires JAX omnistaging to be enabled. In order to enable
+  this, add this to your imports::
+    
+    from jax.config import config
+    config.enable_omnistaging()
+  """
+  def __init__(self):
+    super().__init__(f'Flax Linen requires Omnistaging to be enabled')
 
 
 class InvalidCheckpointError(FlaxError):
@@ -310,4 +526,4 @@ class InvalidCheckpointError(FlaxError):
   overwrite existing checkpoints in the target directory.
   """
   def __init__(self, path, step):
-    super().__init__(f'Trying to save an outdated checkpoint at step: "{step}" and path: "{path}".')
\ No newline at end of file
+    super().__init__(f'Trying to save an outdated checkpoint at step: "{step}" and path: "{path}".')
diff --git a/flax/linen/module.py b/flax/linen/module.py
--- a/flax/linen/module.py
+++ b/flax/linen/module.py
@@ -30,6 +30,7 @@
 import numpy as np
 
 import flax
+from flax import errors
 from flax import traverse_util
 from flax import serialization
 from flax import core
@@ -53,10 +54,7 @@
 
 def _check_omnistaging():
   if not jax.config.omnistaging_enabled:
-    raise RuntimeError(
-        "Flax linen API requires JAX omnistaging to be enabled:\n"
-        "  from jax.config import config\n"
-        "  config.enable_omnistaging()")
+    raise errors.JaxOmnistagingError()
 
 
 def _indent(x: str, num_spaces: int):
@@ -271,7 +269,7 @@ def wrapped_module_method(*args, **kwargs):
 
     if is_compact_method:
       if self.scope is None:
-        raise ValueError("Can't call compact methods on unbound modules")
+        raise errors.CallCompactUnboundModuleError()
       self._state.in_compact_method = True
     _context.module_stack.append(self)
     try:
@@ -303,8 +301,9 @@ def wrapped(self):
 def _get_unbound_fn(method_or_fn: Callable[..., Any]) -> Callable[..., Any]:
   """Returns an unbound function from a method that is possibly bound.
   
-  This means that the returned function does no longer depend on the instance
-  of the class, which is passed as it first argument. 
+  This means that if the passed function belongs of an instance of a class, then
+  the returned function does no longer depend on the instance, which is passed
+  as the first argument to the function.
 
   Args:
     method_or_fn: A class method or function.
@@ -312,12 +311,16 @@ def _get_unbound_fn(method_or_fn: Callable[..., Any]) -> Callable[..., Any]:
     An unbound version of input function.
   """
   if inspect.ismethod(method_or_fn):
-    return method_or_fn.__func__  # pytype: disable=attribute-error
-  elif callable(method_or_fn):
-    return method_or_fn
-  else:
-    raise ValueError('Expect a function or method.')
+    method_or_fn = method_or_fn.__func__  # pytype: disable=attribute-error
 
+  # The method should be callable, and it should have at least one argument
+  # representing the class that is passed in.
+  if (not callable(method_or_fn) or 
+      len(inspect.signature(method_or_fn).parameters) < 1):
+    raise errors.ApplyModuleInvalidMethodError(method_or_fn)
+  
+  return method_or_fn
+  
 
 @dataclasses.dataclass
 class _ModuleInternalState:
@@ -443,8 +446,7 @@ def _customized_dataclass_transform(cls):
     # Use cls.__dict__ to get annotations of cls itself (no parent class).
     annotations = dict(cls.__dict__.get('__annotations__', {}))
     if 'parent' in annotations or 'name' in annotations:
-      raise ValueError(
-          f'properties `parent` and `name` are reserved: {annotations}')
+      raise errors.ReservedModuleAttributeError(annotations)
     # Add `parent` and `name` default fields at end.
     # We temporarily modify base class __dataclass_fields__ to force desired
     # argument behavior and ordering from dataclass class-transform.
@@ -475,10 +477,7 @@ def _verify_single_or_no_compact(cls):
     n_compact_fns = len([method_name for method_name in methods
                          if hasattr(getattr(cls, method_name), 'compact')])
     if n_compact_fns > 1:
-      raise RuntimeError(
-          'Only one method per class can be @compact. You can remove @compact '
-          'and define submodules and variables in setup(), or use two '
-          'separate modules.')
+      raise errors.MultipleMethodsCompactError()
 
   @classmethod
   def _wrap_module_methods(cls):
@@ -499,7 +498,7 @@ def __setattr__(self, name: str, val: Any):
     """Sets an attribute on this Module.
     
     We overload setattr solely to support pythonic naming via assignment of 
-    submodules in the special setup() function::
+    submodules in the special :meth:`setup` function::
 
       self.submodule_name = MyModule(...)
 
@@ -515,10 +514,11 @@ def __setattr__(self, name: str, val: Any):
     
     if not self._state.in_setup and self._state.is_initialized:
       # Raises a TypeError just like frozen python dataclasses.
-      raise TypeError("Module instance is frozen outside of setup method.")
+      raise errors.SetAttributeFrozenModuleError(self.__class__.__name__, name, 
+                                                 val)
     if is_dataclass_attr:
       if self._state.in_setup:
-        raise TypeError("Module construction attributes are frozen.")
+        raise errors.SetAttributeInModuleSetupError()
       object.__setattr__(self, name, val)
     # Submodules are being defined and attached in setup()
     else:
@@ -534,7 +534,7 @@ def __getattr__(self, name: str) -> Any:
       return self.__dict__[name]
     else:
       raise AttributeError(
-          f"'{self.__class__.__name__}' object has no attribute '{name}'")
+          f'"{self.__class__.__name__}" object has no attribute "{name}"')
 
   def __dir__(self) -> List[str]:
     """Call setup() before listing attributes."""
@@ -568,9 +568,7 @@ def __post_init__(self):
       if self.parent._state.in_setup and self.name is None:  # pytype: disable=attribute-error
         return
       if not self.parent._initialization_allowed:
-        raise ValueError(
-            'Submodules must be defined in `setup()` or in a method wrapped '
-            'in `@compact`')
+        raise errors.AssignSubModuleError(self.__class__.__name__)
       # Autonaming of submodules.
       if self.name is None:  # pytype: disable=attribute-error
         prefix = f"{self.__class__.__name__}"
@@ -578,11 +576,8 @@ def __post_init__(self):
         self.name = f"{prefix}_{cursor}"
         self.parent._state.autoname_cursor[prefix] = cursor + 1
       if self.parent._name_taken(self.name, self):
-        raise ValueError(
-            f"A variable of name {self.name} exists already, or "
-            f"trying to share submodule {self.__class__.__name__} by name "
-            f"{self.name}. To share submodules, store module instances as a"
-            f" Python object or as an attribute on self and reuse.")
+        parent_class = self.parent.__class__.__name__
+        raise errors.NameInUseError('submodule', self.name, parent_class)
       self.parent._state.children[self.name] = self
       object.__setattr__(self, 'scope', self.parent.scope.push(self.name))
 
@@ -737,8 +732,7 @@ def variable(self, col: str, name: str, init_fn, *init_args) -> Variable:
           'Variables must be initialized in `setup()` or in a method '
           'wrapped in `@compact`')
     if self._name_taken(name):
-      raise ValueError(
-          f'Name {name} already in use in {self.__class__.__name__}.')
+      raise errors.NameInUseError('variable', name, self.__class__.__name__)
     v = self.scope.variable(col, name, init_fn, *init_args)
     self._state.children[name] = col
     return v
@@ -774,8 +768,7 @@ def param(self, name: str, init_fn: Callable[..., T], *init_args) -> T:
           'Parameters must be initialized in `setup()` or in a method '
           'wrapped in `@compact`')
     if self._name_taken(name):
-      raise ValueError(
-          f'Name {name} already in use in {self.__class__.__name__}.')
+      raise errors.NameInUseError('param', name, self.__class__.__name__)
     v = self.scope.param(name, init_fn, *init_args)
     self._state.children[name] = 'params'
     return v
@@ -871,12 +864,27 @@ def apply(self, variables: VariableDict, *args, rngs: RNGSequences = None,
     """Applies a module method to variables and returns output and modified variables.
 
     Note that `method` should be set if one would like to call `apply` on a
-    different class method than ``__call__``. For instance, suppose a Transformer
-    modules has a method called `encode`, then the following calls `apply` on
-    that method::
+    different class method than ``__call__``. For instance, suppose a
+    Transformer modules has a method called `encode`, then the following calls
+    `apply` on that method::
+
+      model = Transformer()
+      encoded = model.apply({'params': params}, x, method=Transformer.encode)
+  
+    If a function instance is provided, the unbound function is used. For 
+    instance, the example below is equivalent to the one above::
+
+      encoded = model.apply({'params': params}, x, method=model.encode)
+
+    Note ``method`` can also be a function that is not defined in
+    ``Transformer``. In that case, the function should have at least one
+    argument representing an instance of the Module class::
+
+      def other_fn(instance, ...):
+        instance.some_module_attr(...)
+        ...
 
-      model = models.Transformer(config)
-      encoded = model.apply({'params': params}, inputs, method=model.encode)
+      model.apply({'params': params}, x, method=other_fn)
 
     Args:
       variables: A dictionary containing variables keyed by variable
@@ -884,8 +892,9 @@ def apply(self, variables: VariableDict, *args, rngs: RNGSequences = None,
         about variables.
       rngs: a dict of PRNGKeys to initialize the PRNG sequences.
         The "params" PRNG sequence is used to initialize parameters.
-      method: The literal name of a method in this class. If provided, applies
-        this method. If not provided, applies the ``__call__`` method.
+      method: A function to call apply on. This is generally a function in the
+        module. If provided, applies this method. If not provided, applies the
+        ``__call__`` method of the module.
       mutable: Can be bool, str, or list. Specifies which collections should be
                treated as mutable: ``bool``: all/no collections are mutable.
                ``str``: The name of a single mutable collection. ``list``: A
@@ -924,7 +933,10 @@ def init_with_output(self, rngs: Union[PRNGKey, RNGSequences], *args,
       collections.
     """
     if not isinstance(rngs, dict):
-      assert rngs.shape == (2,)
+      if rngs.shape != (2,):
+        raise errors.InvalidRngError(
+            'RNGs should be of shape (2,) in Module '
+            f'{self.__class__.__name__}, but rngs are: {rngs}')
       rngs = {'params': rngs}
     return self.apply(
         {}, *args, rngs=rngs, method=method, mutable=True, **kwargs)


Test patch:
diff --git a/tests/linen/module_test.py b/tests/linen/module_test.py
--- a/tests/linen/module_test.py
+++ b/tests/linen/module_test.py
@@ -253,35 +253,38 @@ def __call__(self, x):
         return x + self.bias
     x = jnp.array([1.])
     scope = Scope({}, {'params': rngkey}, mutable=['params'])
-    with self.assertRaisesRegex(ValueError, 'bias already in use'):
+    msg = 'Could not create param "bias" in Module Dummy: Name in use'
+    with self.assertRaisesRegex(errors.NameInUseError, msg):
       y = Dummy(x.shape, parent=scope)(x)
 
-  def test_setup_var_collision(self):
+  def test_call_var_collision(self):
     rngkey = jax.random.PRNGKey(0)
     class Dummy(nn.Module):
       xshape: Tuple[int]
-      def setup(self):
-        self.bias = self.param('bias', initializers.ones, self.xshape)
-        self.bias = self.param('bias', initializers.ones, self.xshape)
+      @compact
       def __call__(self, x):
-        return x + self.bias
+        bias = self.param('bias', initializers.ones, self.xshape)
+        bias = self.param('bias', initializers.ones, self.xshape)
+        return x + bias
     x = jnp.array([1.])
     scope = Scope({}, {'params': rngkey}, mutable=['params'])
-    with self.assertRaisesRegex(ValueError, 'bias already in use'):
+    msg = 'Could not create param "bias" in Module Dummy: Name in use'
+    with self.assertRaisesRegex(errors.NameInUseError, msg):
       y = Dummy(x.shape, parent=scope)(x)
 
-  def test_call_var_collision(self):
+  def test_setup_var_collision(self):
     rngkey = jax.random.PRNGKey(0)
     class Dummy(nn.Module):
       xshape: Tuple[int]
-      @compact
+      def setup(self):
+        self.bias = self.param('bias', initializers.ones, self.xshape)
+        self.bias = self.param('bias', initializers.ones, self.xshape)
       def __call__(self, x):
-        bias = self.param('bias', initializers.ones, self.xshape)
-        bias = self.param('bias', initializers.ones, self.xshape)
-        return x + bias
+        return x + self.bias
     x = jnp.array([1.])
     scope = Scope({}, {'params': rngkey}, mutable=['params'])
-    with self.assertRaisesRegex(ValueError, 'bias already in use'):
+    msg = 'Could not create param "bias" in Module Dummy: Name in use'
+    with self.assertRaisesRegex(errors.NameInUseError, msg):
       y = Dummy(x.shape, parent=scope)(x)
 
   def test_setattr_name_var_disagreement_allowed_in_lists(self):
@@ -320,43 +323,66 @@ def __call__(self, x):
     y = Dummy(x.shape, parent=scope)(x)
     self.assertEqual(y, jnp.array([2.]))
 
-  def test_submodule_var_collision(self):
+  def test_submodule_var_collision_with_scope(self):
     rngkey = jax.random.PRNGKey(0)
+
     class Dummy(nn.Module):
       xshape: Tuple[int]
+
       def setup(self):
         self.bias = self.param('bias', initializers.ones, self.xshape)
         self.bias = DummyModule()
+
       def __call__(self, x):
         return x + self.bias
+
     x = jnp.array([1.])
     scope = Scope({}, {'params': rngkey}, mutable=['params'])
-    msg = r'Duplicate use of scope name: "bias"'
-    with self.assertRaisesRegex(errors.ScopeNameInUseError, msg):
+
+    msg = 'Duplicate use of scope name: "bias"'
+    with self.assertRaisesWithLiteralMatch(ValueError, msg):
       y = Dummy(x.shape, parent=scope)(x)
+
+  def test_submodule_var_collision_with_submodule(self):
+    rngkey = jax.random.PRNGKey(0)
+
     class Dummy(nn.Module):
       xshape: Tuple[int]
+
       def setup(self):
         self.bias = self.param('bias', initializers.ones, self.xshape)
+
       @compact
       def __call__(self, x):
         bias = DummyModule(name='bias')
         return x + self.bias
+
     x = jnp.array([1.])
     scope = Scope({}, {'params': rngkey}, mutable=['params'])
-    with self.assertRaisesRegex(ValueError, 'name bias exists already'):
+
+    msg = 'Could not create submodule "bias" in Module Dummy: Name in use'
+    with self.assertRaisesRegex(errors.NameInUseError, msg):
       y = Dummy(x.shape, parent=scope)(x)
+
+  def test_submodule_var_collision_with_params(self):
+    rngkey = jax.random.PRNGKey(0)
+
     class Dummy(nn.Module):
       xshape: Tuple[int]
+
       def setup(self):
         self.bias = DummyModule()
+
       @compact
       def __call__(self, x):
         bias = self.param('bias', initializers.ones, self.xshape)
         return x + self.bias
+
     x = jnp.array([1.])
     scope = Scope({}, {'params': rngkey}, mutable=['params'])
-    with self.assertRaisesRegex(ValueError, 'bias already'):
+
+    msg = 'Could not create param "bias" in Module Dummy: Name in use'
+    with self.assertRaisesRegex(errors.NameInUseError, msg):
       y = Dummy(x.shape, parent=scope)(x)
 
   def test_attr_param_name_collision(self):
@@ -369,7 +395,8 @@ def __call__(self, x):
         return x + self.bias
     x = jnp.array([1.])
     scope = Scope({}, {'params': rngkey}, mutable=['params'])
-    with self.assertRaisesRegex(ValueError, 'Name bias already in use'):
+    msg = 'Could not create param "bias" in Module Dummy: Name in use'
+    with self.assertRaisesRegex(errors.NameInUseError, msg):
       y = Dummy(x.shape, parent=scope)(x)
 
   def test_attr_submodule_name_collision(self):
@@ -382,11 +409,13 @@ def __call__(self, x):
         return self.bias(x)
     x = jnp.array([1.])
     scope = Scope({}, {'params': rngkey}, mutable=['params'])
-    with self.assertRaisesRegex(ValueError, 'bias exists already'):
+    msg = 'Could not create submodule "bias" in Module Dummy: Name in use'
+    with self.assertRaisesRegex(errors.NameInUseError, msg):
       y = Dummy(x.shape, parent=scope)(x)
 
   def test_only_one_compact_method(self):
-    with self.assertRaisesRegex(RuntimeError, '@compact'):
+    msg = 'Only one method per class can be @compact'
+    with self.assertRaisesRegex(errors.MultipleMethodsCompactError, msg):
       class Dummy(nn.Module):
         @compact
         def call1(self):
@@ -424,7 +453,9 @@ def __call__(self, x):
         x = bar(x)
         x = bar(x)
         return x
-    with self.assertRaisesRegex(ValueError, '@compact'):
+    msg = (r'Submodule Dense must be defined in `setup\(\)` or in a method '
+            'wrapped in `@compact`')
+    with self.assertRaisesRegex(errors.AssignSubModuleError, msg):
       Foo().init(random.PRNGKey(0), jnp.ones((1, 3)))
 
   def test_forgotten_compact_annotation_with_explicit_parent(self):
@@ -440,7 +471,9 @@ def __call__(self, x):
         x = bar(x)
         return x
 
-    with self.assertRaisesRegex(ValueError, '@compact'):
+    msg = (r'Submodule Dense must be defined in `setup\(\)` or in a method '
+            'wrapped in `@compact`')
+    with self.assertRaisesRegex(errors.AssignSubModuleError, msg):
       Foo().init(random.PRNGKey(0), jnp.ones((1, 3)))
 
   def test_numpy_array_shape_class_args(self):
@@ -568,7 +601,8 @@ def test_module_is_hashable(self):
 
   def test_module_with_scope_is_not_hashable(self):
     module_a = nn.Dense(10, parent=Scope({}))
-    with self.assertRaisesWithLiteralMatch(ValueError, 'Can\'t call __hash__ on modules that hold variables.'):
+    msg = 'Can\'t call __hash__ on modules that hold variables.'
+    with self.assertRaisesWithLiteralMatch(ValueError, msg):
       hash(module_a)
 
   def test_module_trace(self):
@@ -615,9 +649,38 @@ def __call__(self, x):
     self.assertEqual(trace, expected_trace)
 
 
+  def test_module_apply_method(self):
+    class Foo(nn.Module):
+      @nn.compact
+      def __call__(self):
+        pass
+      
+      def test(self):
+        pass
+
+    # We can use both instance and class methods in apply.
+    Foo().apply({}, method=Foo.test)
+    Foo().apply({}, method=Foo().test)
+
+    # We also use a function that is not in the provided Module, although it
+    # should have a first argument representing an instance of the Module (Foo
+    # in this case).
+    x = Foo().apply({}, method=lambda foo_instance: foo_instance)
+    self.assertEqual(type(x), type(Foo()))
+
+    # This is not allowed.
+    msg = 'Cannot call apply()'
+    with self.assertRaisesRegex(errors.ApplyModuleInvalidMethodError, msg):
+      Foo().apply({}, method=lambda: True)
+
+    with self.assertRaisesRegex(errors.ApplyModuleInvalidMethodError, msg):
+      Foo().apply({}, method='allowed_apply_fn')
+
+
   def test_call_unbound_compact_module_methods(self):
     dense = Dense(3)
-    with self.assertRaisesRegex(ValueError, "compact.*unbound module"):
+    msg = r'Can\'t call compact methods on unbound modules'
+    with self.assertRaisesRegex(errors.CallCompactUnboundModuleError, msg):
       dense(jnp.ones((1, )))
 
 
@@ -660,22 +723,23 @@ def bar(self):
 
     empty = EmptyModule()
     # It's fine to call methods of unbound methods that don't depend on
-    # attributes defined during `setup`
+    # attributes defined during `setup`.
     self.assertEqual(empty.bar(), 3)
 
 
-  def test_call_unbound_noncompact_module_methods(self):
+  def test_call_unbound_noncompact_module_methods_depending_on_setup(self):
     class EmptyModule(nn.Module):
-      foo: int = 3
+      def setup(self):
+        self.foo = 2
 
       def bar(self):
         return self.foo
 
     empty = EmptyModule()
-    # It's fine to call methods of unbound methods that don't depend on
-    # attributes defined during `setup`
-    self.assertEqual(empty.bar(), 3)
-
+    msg = r'"EmptyModule" object has no attribute "foo"'
+    with self.assertRaisesRegex(AttributeError, msg):
+      empty.bar()
+     
 
   def test_module_with_attrs(self):
     class Foo(nn.Module):
@@ -700,18 +764,24 @@ def setup(self):
       def __call__(self):
         self.i = 2  # This is not allowed.
 
-    with self.assertRaisesWithLiteralMatch(TypeError, "Module instance is frozen outside of setup method."):
+    msg = ('Can\'t set i=2 for Module of type Foo: Module instance is frozen '
+           'outside of setup method.')
+    with self.assertRaisesRegex(errors.SetAttributeFrozenModuleError, msg):
       Foo().init(random.PRNGKey(0))
 
+
   def test_compact_module_frozen(self):
     class Foo(nn.Module):
       @nn.compact
       def __call__(self):
         self.i = 2
 
-    with self.assertRaisesWithLiteralMatch(TypeError, "Module instance is frozen outside of setup method."):
+    msg = ('Can\'t set i=2 for Module of type Foo: Module instance is frozen '
+           'outside of setup method.')
+    with self.assertRaisesRegex(errors.SetAttributeFrozenModuleError, msg):
       Foo().init(random.PRNGKey(0))
 
+
   def test_submodule_frozen(self):
     class Foo(nn.Module):
       @nn.compact
@@ -719,7 +789,9 @@ def __call__(self):
         dense = nn.Dense(10)
         dense.features = 20  # <--- This is not allowed
 
-    with self.assertRaisesWithLiteralMatch(TypeError, "Module instance is frozen outside of setup method."):
+    msg = ('Can\'t set features=20 for Module of type Dense: Module instance '
+           'is frozen outside of setup method.')
+    with self.assertRaisesRegex(errors.SetAttributeFrozenModuleError, msg):
       Foo().init(random.PRNGKey(0))
 
 
@@ -727,10 +799,11 @@ def test_module_call_not_implemented(self):
     class Foo(nn.Module):
       pass
 
-    foo = Foo()
-    with self.assertRaisesWithLiteralMatch(AttributeError, "'Foo' object has no attribute '__call__'"):
-      foo.init(random.PRNGKey(0))
-  
+    msg = '"Foo" object has no attribute "__call__"'
+    with self.assertRaisesRegex(AttributeError, msg):
+      Foo().init(random.PRNGKey(0))
+
+
   def test_is_mutable_collection(self):
     class EmptyModule(nn.Module):
       def __call__(self):
@@ -795,7 +868,8 @@ class B(nn.Module):
       def setup(self):
         self.c = nn.Dense(2)
 
-    with self.assertRaisesWithLiteralMatch(AttributeError, "'B' object has no attribute 'c'"):
+    msg = '"B" object has no attribute "c"'
+    with self.assertRaisesRegex(AttributeError, msg):
       A().init(random.PRNGKey(0))
 
   def test_unbound_setup_call(self):