Repository: google/flax

Issue:
Improve error message when calling `init` or `apply` on a module class instead of a module object.
A common mistake I see (and do myself frequently) is:
```python
class B(nn.Module):
  @nn.compact
  def __call__(self, x):
    return x

k = random.PRNGKey(0)
x = random.uniform(random.PRNGKey(1), (2,))
B.init(k, x)   # B is module class, not B() a module instance
...
B.apply(vs, x)   # similar issue w. apply called on class instead of instance.
```

We could improve the error message here, instead of allowing the first argument to be treated as `self` accidentally, leading to a pretty mysterious error currently.


Discussion:

Resolving patch:
diff --git a/flax/errors.py b/flax/errors.py
--- a/flax/errors.py
+++ b/flax/errors.py
@@ -569,6 +569,25 @@ def get_submodule(module):
   def __init__(self):
     super().__init__('Can\'t call compact methods on unbound modules')
 
+class InvalidInstanceModuleError(FlaxError):
+  """
+  This error occurs when you are trying to call `.init()`, `.init_with_output()` or `.apply()
+  on the Module class itself, instead of an instance of the Module class.
+  For example, the error will be raised when trying to run this code::
+
+    class B(nn.Module):
+      @nn.compact
+      def __call__(self, x):
+        return x
+
+    k = random.PRNGKey(0)
+    x = random.uniform(random.PRNGKey(1), (2,))
+    B.init(k, x)   # B is module class, not B() a module instance
+    B.apply(vs, x)   # similar issue with apply called on class instead of instance.
+  """
+  def __init__(self):
+    super().__init__('Can only call init, init_with_output or apply methods on an instance of the Module class, not the Module class itself')
+
 
 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
@@ -1255,6 +1255,9 @@ def other_fn(instance, ...):
       mutable, returns ``(output, vars)``, where ``vars`` are is a dict
       of the modified collections.
     """
+    if not isinstance(self, Module):
+      raise errors.InvalidInstanceModuleError()
+
     if method is None:
       method = self.__call__
     method = _get_unbound_fn(method)
@@ -1295,6 +1298,9 @@ def init_with_output(self,
       `(output, vars)``, where ``vars`` are is a dict of the modified
       collections.
     """
+    if not isinstance(self, Module):
+      raise errors.InvalidInstanceModuleError()
+
     if not isinstance(rngs, dict):
       if not core.scope._is_valid_rng(rngs):
         raise errors.InvalidRngError(
@@ -1348,6 +1354,9 @@ def init(self,
     Returns:
       The initialized variable dict.
     """
+    if not isinstance(self, Module):
+      raise errors.InvalidInstanceModuleError()
+
     _, v_out = self.init_with_output(
         rngs,
         *args,


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
@@ -1756,6 +1756,22 @@ def __call__(self):
     self.assertTrue(foo.init_with_output(k)[0])
     self.assertFalse(foo.apply({}))
 
+  def test_throws_invalid_instance_module_error(self):
+
+    class B(nn.Module):
+      @nn.compact
+      def __call__(self, x):
+        return x
+
+    k = random.PRNGKey(0)
+    x = random.uniform(random.PRNGKey(1), (2,))
+
+    with self.assertRaises(errors.InvalidInstanceModuleError):
+      B.init(k, x)   # B is module class, not B() a module instance
+    with self.assertRaises(errors.InvalidInstanceModuleError):
+      B.init_with_output(k, x)
+    with self.assertRaises(errors.InvalidInstanceModuleError):
+      B.apply({}, x)   # similar issue w. apply called on class instead of instance.
 
 class LeakTests(absltest.TestCase):