Repository: google/flax

Issue:
flax.linen.remat with concrete=True doesn't work with jax 0.3.17
### Problem you have encountered:

This may already be on the Flax team's radar, but I noticed that when using flax.linen.remat, setting concrete=True doesn't work with Jax 0.3.17, for the reasons discussed [here](https://jax.readthedocs.io/en/latest/jep/11830-new-remat-checkpoint.html). 

As of version 0.6.0: flax.linen.remat
(1) passes the argument ```concrete=True``` to ```jax.remat```, which leads to an error message. 
(2) does not accept an argument ```static_argnums```, as used in the latest ```jax.remat```.

Interestingly, pip's constraint solver did not seem to be aware of this incompatibility; running ```pip install jax, flax``` allowed me to install flax==0.6.0 with jax==0.3.17, leading to the observed problem. 

As a workaround, I've downgraded to jax==0.3.16, and am running ```jax.config.update("jax_new_checkpoint", False)``` at the top of my scripts, as suggested by the link above.

### What you expected to happen:

To ensure compatibility with Jax's remat functionality, future versions of flax.linen.remat would ideally accept an argument ```static_argnums```, which can be passed to the jax.remat implementation. 

In the traceback triggered by Flax passing ```concrete=True```, the Jax developers also remark that 

> If jax.numpy operations need to be performed on static arguments, we can use the `jax.ensure_compile_time_eval()` context manager. 

which may also be relevant to the future design of flax.linen.remat. 

### Steps to reproduce:
The problem can be reproduced by running the script

```
import flax.linen as nn
import jax

class Foo(nn.Module):
    def setup(self):
        self.linear = nn.remat(nn.Dense, concrete=True)(100, use_bias=False)

    def __call__(self, inputs):
        return self.linear(inputs)

if __name__ == '__main__':
   rng = jax.random.PRNGKey(0)
   rng, sk1, sk2 = jax.random.split(rng, 3)
   foo = Foo()
   input = jax.random.normal(sk1, [1, 10])
   params = foo.init({"params": sk2}, input)["params"]
   out = foo.apply({"params": params}, input)
```

### Logs, error messages, etc:

When I run the above script, I obtain the following traceback: 
<details>
<summary>toggle to show</summary>

```
Traceback (most recent call last):
  File "/Users/lucaslingle/PycharmProjects/project123/src/project123/nn/generic_module.py", line 17, in <module>
    params = foo.init({"params": sk2}, input)["params"]
  File "/Users/lucaslingle/opt/miniconda3/envs/project123/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 162, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/Users/lucaslingle/opt/miniconda3/envs/project123/lib/python3.10/site-packages/flax/linen/module.py", line 1273, in init
    _, v_out = self.init_with_output(
  File "/Users/lucaslingle/opt/miniconda3/envs/project123/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 162, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/Users/lucaslingle/opt/miniconda3/envs/project123/lib/python3.10/site-packages/flax/linen/module.py", line 1229, in init_with_output
    return init_with_output(
  File "/Users/lucaslingle/opt/miniconda3/envs/project123/lib/python3.10/site-packages/flax/core/scope.py", line 897, in wrapper
    return apply(fn, mutable=mutable, flags=init_flags)({}, *args, rngs=rngs,
  File "/Users/lucaslingle/opt/miniconda3/envs/project123/lib/python3.10/site-packages/flax/core/scope.py", line 865, in wrapper
    y = fn(root, *args, **kwargs)
  File "/Users/lucaslingle/opt/miniconda3/envs/project123/lib/python3.10/site-packages/flax/linen/module.py", line 1647, in scope_fn
    return fn(module.clone(parent=scope), *args, **kwargs)
  File "/Users/lucaslingle/opt/miniconda3/envs/project123/lib/python3.10/site-packages/flax/linen/module.py", line 361, in wrapped_module_method
    return self._call_wrapped_method(fun, args, kwargs)
  File "/Users/lucaslingle/opt/miniconda3/envs/project123/lib/python3.10/site-packages/flax/linen/module.py", line 657, in _call_wrapped_method
    y = fun(self, *args, **kwargs)
  File "/Users/lucaslingle/PycharmProjects/project123/src/project123/nn/generic_module.py", line 9, in __call__
    return self.linear(inputs)
  File "/Users/lucaslingle/opt/miniconda3/envs/project123/lib/python3.10/site-packages/flax/linen/transforms.py", line 316, in wrapped_fn
    ret = trafo_fn(module_scopes, *args, **kwargs)
  File "/Users/lucaslingle/opt/miniconda3/envs/project123/lib/python3.10/site-packages/flax/core/lift.py", line 213, in wrapper
    y, out_variable_groups_xs_t = fn(
  File "/Users/lucaslingle/opt/miniconda3/envs/project123/lib/python3.10/site-packages/flax/core/lift.py", line 1177, in inner
    def rematted(variable_groups, rng_groups, *args, **kwargs):
  File "/Users/lucaslingle/opt/miniconda3/envs/project123/lib/python3.10/site-packages/jax/_src/api.py", line 3084, in checkpoint
    raise NotImplementedError(msg)
jax._src.traceback_util.UnfilteredStackTrace: NotImplementedError: The 'concrete' option to jax.checkpoint / jax.remat is deprecated; in its place, you can use its `static_argnums` option, and if necessary the `jax.ensure_compile_time_eval()` context manager.

For example, if using `concrete=True` for an `is_training` flag:

  from functools import partial

  @partial(jax.checkpoint, concrete=True)
  def foo(x, is_training):
    if is_training:
      return f(x)
    else:
      return g(x)

replace it with a use of `static_argnums`:

  @partial(jax.checkpoint, static_argnums=(1,))
  def foo(x, is_training):
    ...

If jax.numpy operations need to be performed on static arguments, we can use the `jax.ensure_compile_time_eval()` context manager. For example, we can replace this use of `concrete=True`
:
  @partial(jax.checkpoint, concrete=True)
  def foo(x, y):
    if y > 0:
      return f(x)
    else:
      return g(x)

with this combination of `static_argnums` and `jax.ensure_compile_time_eval()`:

  @partial(jax.checkpoint, static_argnums=(1,))
  def foo(x, y):
    with jax.ensure_compile_time_eval():
      y_pos = y > 0
    if y_pos:
      return f(x)
    else:
      return g(x)

See https://jax.readthedocs.io/en/latest/jep/11830-new-remat-checkpoint.html

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/Users/lucaslingle/PycharmProjects/project123/src/project123/nn/generic_module.py", line 17, in <module>
    params = foo.init({"params": sk2}, input)["params"]
  File "/Users/lucaslingle/PycharmProjects/project123/src/project123/nn/generic_module.py", line 9, in __call__
    return self.linear(inputs)
NotImplementedError: The 'concrete' option to jax.checkpoint / jax.remat is deprecated; in its place, you can use its `static_argnums` option, and if necessary the `jax.ensure_compile_time_eval()` context manager.

For example, if using `concrete=True` for an `is_training` flag:

  from functools import partial

  @partial(jax.checkpoint, concrete=True)
  def foo(x, is_training):
    if is_training:
      return f(x)
    else:
      return g(x)

replace it with a use of `static_argnums`:

  @partial(jax.checkpoint, static_argnums=(1,))
  def foo(x, is_training):
    ...

If jax.numpy operations need to be performed on static arguments, we can use the `jax.ensure_compile_time_eval()` context manager. For example, we can replace this use of `concrete=True`
:
  @partial(jax.checkpoint, concrete=True)
  def foo(x, y):
    if y > 0:
      return f(x)
    else:
      return g(x)

with this combination of `static_argnums` and `jax.ensure_compile_time_eval()`:

  @partial(jax.checkpoint, static_argnums=(1,))
  def foo(x, y):
    with jax.ensure_compile_time_eval():
      y_pos = y > 0
    if y_pos:
      return f(x)
    else:
      return g(x)

See https://jax.readthedocs.io/en/latest/jep/11830-new-remat-checkpoint.html
```
</details>

### System information
- OS Platform and Distribution: ```MacOS Catalina 10.15.7```
- Flax, jax, jaxlib versions: ```flax==0.6.0, jax==0.3.17, jaxlib==0.3.15```
- Python version: ```3.10```
- GPU/TPU model and memory: ```N/A```
- CUDA version (if applicable): ```N/A```





Discussion:
Hey @lucaslingle, thanks for bringing this up! I've opened #2457 with a fix for this.

Resolving patch:
diff --git a/flax/core/lift.py b/flax/core/lift.py
--- a/flax/core/lift.py
+++ b/flax/core/lift.py
@@ -1138,6 +1138,7 @@ def checkpoint(fn: Callable[..., Any],
                rngs: PRNGSequenceFilter = True,
                concrete: bool = False,
                prevent_cse: bool = True,
+               static_argnums: Union[int, Tuple[int, ...]] = (),
                policy: Optional[Callable[..., bool]] = None,
                ) -> Callable[..., Any]:
   """Lifted version of ``jax.checkpoint``.
@@ -1164,15 +1165,21 @@ def checkpoint(fn: Callable[..., Any],
       ``pmap``, CSE can defeat the purpose of this decorator. But in some
       settings, like when used inside a ``scan``, this CSE prevention mechanism
       is unnecessary, in which case ``prevent_cse`` can be set to False.
+    static_argnums: Optional, int or sequence of ints, indicates which argument 
+      values on which to specialize for tracing and caching purposes. Specifying 
+      arguments as static can avoid ConcretizationTypeErrors when tracing, but 
+      at the cost of more retracing overheads.
     policy: Experimental checkpoint policy, see ``jax.checkpoint``.
   Returns:
     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, **kwargs):
+    # add 2 to each static_argnums because we add two initial arguments to rematted
+    static_argnums_ = jax.tree_util.tree_map(lambda x: x + 2, static_argnums)
     @functools.partial(jax.remat,
-                       concrete=concrete, prevent_cse=prevent_cse,
-                       policy=policy)
+                       concrete=concrete, static_argnums=static_argnums_, 
+                       prevent_cse=prevent_cse, policy=policy)
     @functools.wraps(fn)
     def rematted(variable_groups, rng_groups, *args, **kwargs):
       scope = scope_fn(variable_groups, rng_groups)
diff --git a/flax/linen/transforms.py b/flax/linen/transforms.py
--- a/flax/linen/transforms.py
+++ b/flax/linen/transforms.py
@@ -572,6 +572,7 @@ def checkpoint(target: Target,
                rngs: lift.PRNGSequenceFilter = True,
                concrete: bool = False,
                prevent_cse: bool = True,
+               static_argnums: Union[int, Tuple[int, ...]] = (),
                policy: Optional[Callable[..., bool]] = None,
                methods=None) -> Target:
   """Lifted version of ``jax.checkpoint``.
@@ -599,6 +600,10 @@ def checkpoint(target: Target,
       ``pmap``, CSE can defeat the purpose of this decorator. But in some
       settings, like when used inside a ``scan``, this CSE prevention mechanism
       is unnecessary, in which case ``prevent_cse`` should be set to False.
+    static_argnums: Optional, int or sequence of ints, indicates which argument 
+      values on which to specialize for tracing and caching purposes. Specifying 
+      arguments as static can avoid ConcretizationTypeErrors when tracing, but 
+      at the cost of more retracing overheads.
     policy: Experimental checkpoint policy, see ``jax.checkpoint``.
     methods: If `target` is a `Module`, the methods of `Module` to checkpoint.
 
@@ -606,9 +611,13 @@ def checkpoint(target: Target,
     A wrapped version of ``target``. When computing gradients intermediate
     computations will be re-computed on the backward pass.
   """
+  # subtract 1 from each static_argnums because 'self' is not passed to the
+  # lifted function
+  static_argnums = jax.tree_util.tree_map(lambda x: x - 1, static_argnums)
   return lift_transform(
       lift.checkpoint, target,
       variables=variables, rngs=rngs, concrete=concrete,
+      static_argnums=static_argnums,
       prevent_cse=prevent_cse, policy=policy,
       methods=methods)
 


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
@@ -145,6 +145,73 @@ def __call__(self, input, apply_relu : bool = False):
     # This next line crashes with a concretization error
     _ = jax.grad(lambda x: remat_model.apply(p, x, apply_relu=True))(x)
 
+  def test_remat_static_argnums(self):
+    test = self
+    
+    class Foo(nn.Module):
+      train_is_static: bool
+
+      @nn.compact
+      def __call__(self, inputs, train: bool):
+        if self.train_is_static:
+          test.assertTrue(isinstance(train, bool))
+        else:
+          test.assertTrue(isinstance(train, jnp.ndarray))
+        
+        return nn.Dense(3, use_bias=False)(inputs)
+
+    # set train as a static argument
+    FooRemat = nn.remat(Foo, static_argnums=(2,))
+    foo = FooRemat(train_is_static=True)
+
+    x = jnp.empty((1, 2))
+    variables = foo.init(random.PRNGKey(0), x, True)
+    y = foo.apply(variables, x, False)
+    self.assertEqual(y.shape, (1, 3))
+
+    # set train as a non-static arguments
+    FooRemat = nn.remat(Foo, static_argnums=())
+    foo = FooRemat(train_is_static=False)
+
+    variables = foo.init(random.PRNGKey(0), x, True)
+    y = foo.apply(variables, x, False)
+    self.assertEqual(y.shape, (1, 3))
+  
+  def test_remat_decorator_static_argnums(self):
+    test = self
+
+    class FooTrainStatic(nn.Module):
+      @partial(nn.remat, static_argnums=(2,))
+      @nn.compact
+      def __call__(self, inputs, train: bool):
+        test.assertTrue(isinstance(train, bool))
+        
+        return nn.Dense(3, use_bias=False)(inputs)
+
+    # set train as a static argument
+    foo = FooTrainStatic()
+
+    x = jnp.empty((1, 2))
+    variables = foo.init(random.PRNGKey(0), x, True)
+    y = foo.apply(variables, x, False)
+    self.assertEqual(y.shape, (1, 3))
+
+    class FooTrainDynamic(nn.Module):
+      @partial(nn.remat, static_argnums=())
+      @nn.compact
+      def __call__(self, inputs, train: bool):
+        test.assertTrue(isinstance(train, jnp.ndarray))
+        
+        return nn.Dense(3, use_bias=False)(inputs)
+
+    # set train as a non-static arguments
+    foo = FooTrainDynamic()
+
+    variables = foo.init(random.PRNGKey(0), x, True)
+    y = foo.apply(variables, x, False)
+    self.assertEqual(y.shape, (1, 3))
+    
+
   def test_vmap(self):
     key1, key2 = random.split(random.PRNGKey(3), 2)
     x = random.uniform(key1, (4, 4))