Avoid passing concrete argument to jax.remat

This argument has had no effect since JAX v0.3.17, aside from raising `NotImplementedError` if it is set to `True`. It will be deprecated in JAX v0.8.2 and eventually removed.

#jax-fixit

PiperOrigin-RevId: 839319133

Author: Jake VanderPlas

flax/core/lift.py

@@ -1458,9 +1458,9 @@ 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,
       static_argnums=static_argnums_,
       prevent_cse=prevent_cse,
       policy=policy,

flax/linen/partitioning.py

@@ -552,9 +552,20 @@ def _repack_remat_args(dyn_args, static_args):
   def inner(scope_fn, repack_fn, variable_groups, rng_groups, *args):
     static_args = tuple(x for i, x in enumerate(args) if i in static_argnums)
     dyn_args = tuple(x for i, x in enumerate(args) if i not in static_argnums)
+    
+    # After JAX v0.3.16, concrete=False is a no-op and concrete=True raises
+    # NotImplementedError. Starting in JAX v0.8.2, the concrete argument is
+    # deprecated and will be removed in the future.
+    if concrete:
+      raise NotImplementedError(
+          "The concrete argument is deprecated. Use static_argnums instead, and"
+          " for more information, see"
+          " https://docs.jax.dev/en/latest/jep/11830-new-remat-checkpoint.html"
+      )
+    del concrete
 
     @functools.partial(
-        jax.remat, concrete=concrete, prevent_cse=prevent_cse, policy=policy
+        jax.remat, prevent_cse=prevent_cse, policy=policy
     )
     @functools.wraps(fn)
     def rematted(variable_groups, rng_groups, *dyn_args):