Repository: google/flax

Issue:
remat: wrapper() got an unexpected keyword argument 'use_running_average'
### Problem you have encountered:

The transformed module returned by `remat` does not expect same keyword arguments as original when used as a submodule in `@compact` decorated `__call__` method.

### What you expected to happen:

The transformed module expects the same keyword arguments.

### Logs, error messages, etc:

```
TypeError: wrapper() got an unexpected keyword argument 'use_running_average'
```

### Steps to reproduce:
Whenever possible, please provide a *minimal example*. Please consider submitting it as a Colab link.

```
import jax
import flax.linen as linen 

from jax import numpy as jnp
from typing import Optional

class MyModule(linen.Module):
    
    @linen.compact
    def __call__(self, x, train: Optional[bool]=True):
        return linen.remat(linen.BatchNorm)()(x, use_running_average=not train)

model = MyModule()
key = jax.random.PRNGKey(0)
variables = model.init(key, jnp.ones((10,), jnp.float32))
```
https://colab.research.google.com/drive/1JsmxSn4Msor5D6G5XHpokfHOFlUcsXIJ?usp=sharing


Discussion:
Same. This seems because pack [here](https://github.com/google/flax/blob/095517e679d1687b13e106354e966e418756e535/flax/core/lift.py#L73) returns function(see L83 in the same file below) does not accept keyword arguments. When I manually adds **kwargs to the parameter list of the wrapper, L180,666,671, and 674, the problem is addressed

Resolving patch:
diff --git a/flax/core/lift.py b/flax/core/lift.py
--- a/flax/core/lift.py
+++ b/flax/core/lift.py
@@ -17,6 +17,7 @@
 import collections
 from dataclasses import dataclass
 import functools
+import warnings
 
 
 import jax
@@ -74,13 +75,17 @@ def pack(fn: Callable[..., Any],
          in_variable_filters: Sequence[CollectionFilter],
          out_variable_filters: Sequence[CollectionFilter],
          rng_filters: Sequence[PRNGSequenceFilter],
-         name=None) -> Callable[..., Any]:
+         name=None,
+         enable_kwargs=False) -> Callable[..., Any]:
   """Pack variables and rngs for functional transformations.
 
   The pack function is the building block for all other lifted transformations.
   """
   @functools.wraps(fn)
-  def wrapper(scope_tree: Scope, *args):
+  def wrapper(scope_tree: Scope, *args, **kwargs):
+    if not enable_kwargs and kwargs:
+      msg = 'kwargs are not supported in {}, so \"{}\" is(are) ignored'
+      warnings.warn(msg.format(name, ', '.join(kwargs.keys())), RuntimeWarning)
     # pylint: disable=protected-access
     scopes, treedef = jax.tree_flatten(scope_tree)
     scopes, paths = _dedup_scopes(scopes)
@@ -174,10 +179,16 @@ def repack(inner_scope_tree):
       return _transpose(out_variable_groups_xs)
 
     try:
-      y, out_variable_groups_xs_t = fn(
-          scope_fn, repack,
-          variable_groups_xs_t, rng_groups_xs_t,
-          *args)
+      if enable_kwargs:
+        y, out_variable_groups_xs_t = fn(
+            scope_fn, repack,
+            variable_groups_xs_t, rng_groups_xs_t,
+            *args, **kwargs)
+      else:
+        y, out_variable_groups_xs_t = fn(
+            scope_fn, repack,
+            variable_groups_xs_t, rng_groups_xs_t,
+            *args)
     finally:
       for inner_scope in inner_scopes:
         inner_scope.invalidate()
@@ -672,16 +683,16 @@ def checkpoint(fn: Callable[..., Any],
     A wrapped version of ``fn``. When computing gradients intermediate
     computations will be re-computed when computing gradients.
   """
-  def inner(scope_fn, repack_fn, variable_groups, rng_groups, *args):
+  def inner(scope_fn, repack_fn, variable_groups, rng_groups, *args, **kwargs):
     @functools.partial(jax.remat, concrete=concrete, prevent_cse=prevent_cse)
     @functools.wraps(fn)
-    def rematted(variable_groups, rng_groups, *args):
+    def rematted(variable_groups, rng_groups, *args, **kwargs):
       scope = scope_fn(variable_groups, rng_groups)
-      y = fn(scope, *args)
+      y = fn(scope, *args, **kwargs)
       return y, repack_fn(scope)
 
-    return rematted(variable_groups, rng_groups, *args)
-  return pack(inner, (variables,), (variables,), (rngs,), name='remat')
+    return rematted(variable_groups, rng_groups, *args, **kwargs)
+  return pack(inner, (variables,), (variables,), (rngs,), name='remat', enable_kwargs=True)
 
 
 remat = checkpoint


Test patch:
diff --git a/tests/linen/linen_transforms_test.py b/tests/linen/linen_transforms_test.py
--- a/tests/linen/linen_transforms_test.py
+++ b/tests/linen/linen_transforms_test.py
@@ -121,6 +121,19 @@ def test_remat_decorated(self):
 
     self.assertTrue(np.all(y1 == y2))
 
+  def test_remat_kwargs(self):
+    class ConditionalReLU(nn.Module):
+      @nn.compact
+      def __call__(self, input, apply_relu : bool = False):
+        return nn.relu(input) if apply_relu else input
+    key = random.PRNGKey(0)
+    x = jnp.ones((4, 4)) * -1
+    remat_model = nn.remat(ConditionalReLU)()
+    p = remat_model.init(key, x)
+    y = remat_model.apply(p, x, apply_relu=True)
+
+    self.assertTrue(np.all(y == jnp.zeros_like(x)))
+
   def test_vmap(self):
     key1, key2 = random.split(random.PRNGKey(3), 2)
     x = random.uniform(key1, (4, 4))