allow method argument to accept submodules

Author: cgarciae

flax/linen/module.py

@@ -1911,6 +1911,12 @@ def other_fn(instance, ...):
             f"'{class_name}.{attribute_name}' must be a callable, got"
             f' {type(method)}.'
         )
+      # if the `method` string is a submodule, we create a lambda function
+      # that calls the submodule, forwarding all arguments.
+      if isinstance(method, Module):
+        method = lambda self, *args, **kwargs: getattr(self, attribute_name)(
+            *args, **kwargs
+        )
     elif method is None:
       method = self.__call__
     method = _get_unbound_fn(method)

tests/linen/linen_module_test.py

@@ -887,6 +887,19 @@ def test(self):
       # test same for init.
       Foo().init({}, method='not_callable')
 
+  def test_module_apply_method_submodule(self):
+    class Foo(nn.Module):
+      bar: nn.Module
+
+      @nn.compact
+      def __call__(self, x):
+        return self.bar(x)
+
+    foo = Foo(nn.Dense(3))
+    variables = foo.init(jax.random.PRNGKey(0), jnp.zeros(3))
+
+    foo.apply(variables, jnp.zeros(3), method='bar')
+
   def test_call_unbound_compact_module_methods(self):
     dense = Dense(3)
     msg = r'Can\'t call compact methods on unbound modules'