allow specifying method as string

Author: cgarciae

flax/linen/module.py

@@ -1363,7 +1363,7 @@ def apply(self,
             variables: VariableDict,
             *args,
             rngs: Optional[RNGSequences] = None,
-            method: Optional[Callable[..., Any]] = None,
+            method: Union[Callable[..., Any], str, None] = None,
             mutable: CollectionFilter = False,
             capture_intermediates: Union[bool, Callable[['Module', str], bool]] = False,
             **kwargs) -> Union[Any, Tuple[Any, FrozenVariableDict]]:
@@ -1382,6 +1382,11 @@ def apply(self,
 
       encoded = model.apply({'params': params}, x, method=model.encode)
 
+    You can also pass a string to a callable attribute of the module. For
+    example, the previous can be written as::
+
+      encoded = model.apply({'params': params}, x, method='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::
@@ -1420,7 +1425,14 @@ def other_fn(instance, ...):
     """
     Module._module_checks(self)
 
-    if method is None:
+    if isinstance(method, str):
+      attribute_name = method
+      method = getattr(self, attribute_name)
+      if not callable(method):
+        class_name = type(self).__name__
+        raise TypeError(f'\'{class_name}.{attribute_name}\' must be a callable, got {type(method)}.')
+
+    elif method is None:
       method = self.__call__
     method = _get_unbound_fn(method)
     return apply(

tests/linen/linen_module_test.py

@@ -826,6 +826,7 @@ def __call__(self, x):
   def test_module_apply_method(self):
 
     class Foo(nn.Module):
+      not_callable: int = 1
 
       @nn.compact
       def __call__(self):
@@ -849,9 +850,18 @@ def test(self):
     with self.assertRaisesRegex(errors.ApplyModuleInvalidMethodError, msg):
       Foo().apply({}, method=lambda: True)
 
-    with self.assertRaisesRegex(errors.ApplyModuleInvalidMethodError, msg):
+    # string method names are also allowed.
+    Foo().apply({}, method='test')
+
+    # non-existent attribute names will yield AttributeError.
+    with self.assertRaisesRegex(AttributeError, "allowed_apply_fn"):
       Foo().apply({}, method='allowed_apply_fn')
 
+    # attributes which are not callables yield TypeError.
+    with self.assertRaisesRegex(TypeError, "'Foo.not_callable' must be a callable"):
+      Foo().apply({}, method='not_callable')
+
+
   def test_call_unbound_compact_module_methods(self):
     dense = Dense(3)
     msg = r'Can\'t call compact methods on unbound modules'