Repository: google/flax

Issue:
super().__call__ not generally safe to call in subclass __call__
```
class Foo(nn.Module):
  @nn.compact
  def __call__(self, x):
    return nn.Dense(4)(x)

class Bar(Foo):
  @nn.compact
  def __call__(self, x):
    y = super().__call__(x)
    return nn.Dense(4)(y)

k = random.PRNGKey(0)
x = random.randint(k, (4, 7), 0, 256)

variables = Bar().init(k, x)
y = Bar().apply(variables, x)
```
returns 
```
AssignSubModuleError: Submodule Dense must be defined in `setup()` or in a method wrapped in `@compact` (https://flax.readthedocs.io/en/improve-error/flax.errors.html#flax.errors.AssignSubModuleError)
```

This happens because the parent `super().__call__` is wrapped to mark itself as "compact" upon entry, clearing the "compact" state upon exit, leading to any remaining variable/submodule instantiations in the subclass `__call__` causing an error because the function no longs "looks compact".

I think we could fix this by adding a bit of logic to the `wrap_method_once` function and by passing in the class to `wrap_method_once` during subclass initialization to detect when `self` isn't an instance of the class it's "supposed to be", and having the wrapper acting as a passthrough.


Discussion:
(copying from an offline discussion -- if we "formally" support module inheritence and test it rigorously, i think it may be fine to disallow the base class from having optional attributes. otherwise we'd have to reorder the base class optional attributes and put them at the end -- but that would be very confusing for folks who use positional args as opposed to kwargs)


An alternative fix would be to allow nested compact calls

Resolving patch:
diff --git a/flax/linen/module.py b/flax/linen/module.py
--- a/flax/linen/module.py
+++ b/flax/linen/module.py
@@ -259,12 +259,16 @@ def wrapped_module_method(*args, **kwargs):
     is_compact_method = hasattr(fun, 'compact')
     is_setup_method = fun.__name__ == 'setup'
     # We lazily call setup() only when needed.
-    if not is_setup_method:
+    if is_setup_method:
+      is_recurrent = self._state.in_setup
+      self._state.in_setup = True
+    else:
       self._try_setup()
 
     if is_compact_method:
       if self.scope is None:
         raise errors.CallCompactUnboundModuleError()
+      is_recurrent = self._state.in_compact_method
       self._state.in_compact_method = True
     _context.module_stack.append(self)
     try:
@@ -278,7 +282,10 @@ def wrapped_module_method(*args, **kwargs):
       _context.module_stack.pop()
       if is_compact_method:
         object.__setattr__(self, 'scope', self.scope.rewound())
-      if is_compact_method or is_setup_method:
+      # setup or compact calls can be recurrent for example due to super calls
+      # resetting the state would cause is compact/setup method
+      # to be set to False prematurely.
+      if (is_compact_method or is_setup_method) and not is_recurrent:
         self._state.reset()
   wrapped_module_method.method_handler_wrapped = True
   return wrapped_module_method


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
@@ -1237,6 +1237,51 @@ def __call__(self, x):
     y = Foo().apply(variables, x)
     self.assertEqual(y.shape, (2,))
 
+  def test_super_compact(self):
+    class Foo(nn.Module):
+      @nn.compact
+      def __call__(self, x):
+        return nn.Dense(4)(x)
+
+    class Bar(Foo):
+      @nn.compact
+      def __call__(self, x):
+        y = super().__call__(x)
+        return nn.Dense(3)(y)
+
+    k = random.PRNGKey(0)
+    x = jnp.ones((4, 7))
+
+    variables = Bar().init(k, x)
+    shapes = jax.tree_map(np.shape, variables['params'])
+    self.assertEqual(shapes, {
+      'Dense_0': {'kernel': (7, 4), 'bias': (4,)},
+      'Dense_1': {'kernel': (4, 3), 'bias': (3,)},
+    })
+    y = Bar().apply(variables, x)
+    self.assertEqual(y.shape, (4, 3))
+  
+  def test_super_setup(self):
+    class Foo(nn.Module):
+      def setup(self):
+        self.a = nn.Dense(4)
+
+    class Bar(Foo):
+
+      def setup(self):
+        super().setup()
+        self.b = nn.Dense(3)
+
+      def __call__(self, x):
+        y = self.a(x)
+        return self.b(y)
+
+    k = random.PRNGKey(0)
+    x = jnp.ones((4, 7))
+
+    variables = Bar().init(k, x)
+    y = Bar().apply(variables, x)
+    self.assertEqual(y.shape, (4, 3))
 
 if __name__ == '__main__':
   absltest.main()