Removed unused jax_remat_opt_barrier config option

It defaults to True and is not flipped to False by any internal JAX users.

PiperOrigin-RevId: 744754343

Author: superbobry

jax/_src/ad_checkpoint.py

@@ -757,89 +757,34 @@ def _has_effects(effects) -> bool:
   return bool({e for e in effects if not isinstance(e, core.NamedAxisEffect)})
 
 
-def remat_expansion(*args, jaxpr: core.Jaxpr, prevent_cse: bool,
-                    differentiated: bool, is_gpu_platform: bool = False,
-                    **_):
+def remat_expansion(
+    *args, jaxpr: core.Jaxpr, prevent_cse: bool, differentiated: bool, **_
+):
   assert not jaxpr.constvars
 
   if differentiated and prevent_cse:
-    if config.remat_opt_barrier.value:
-      translation_rule = _remat_translation_using_opt_barrier
-    elif is_gpu_platform:
-      translation_rule = _remat_translation_using_while
-    else:
-      translation_rule = _remat_translation_using_cond
+    translation_rule = _remat_translation_using_opt_barrier
   else:
     translation_rule = lambda *args, jaxpr: core.eval_jaxpr(jaxpr, (), *args)
 
   return api.named_call(translation_rule, name="checkpoint")(*args, jaxpr=jaxpr)
 
+
 def _remat_translation_using_opt_barrier(*args, jaxpr: core.Jaxpr):
   args = lax_internal.optimization_barrier(args)
   return core.eval_jaxpr(jaxpr, (), *args)
 
-# TODO(mattjj): add core utility for 'create dummy value for this type'?
-def _dummy_like(aval: core.AbstractValue) -> Any:
-  if aval is core.abstract_token:
-    return lax_internal.create_token()
-  elif isinstance(aval, (core.ShapedArray, core.DShapedArray)):
-    return lax_internal.broadcast(lax_internal.empty(aval.dtype), aval.shape)  # type: ignore
-  else:
-    raise ValueError(aval)
-
-def _remat_translation_using_while(*args, jaxpr: core.Jaxpr):
-  # Implements:
-  #  for(counter=0, result=0; counter < rng(1, 2); counter ++) {
-  #     result = eval_jaxpr(*args)
-  #  }
-  # The loop carry is a tuple: (counter, result, args)
-  from jax._src.lax import control_flow as lax_control_flow
-
-  avals_out = tuple(v.aval for v in jaxpr.outvars)
-  carry_init = (np.int32(0), tuple(map(_dummy_like, avals_out)), args)
-  def cond(carry):
-    counter, _, _ = carry
-    unif = lax_internal.rng_uniform(np.int32(1), np.int32(2), shape=())
-    return counter < unif
-
-  def body(carry):
-    counter, _, args = carry
-    results = core.eval_jaxpr(jaxpr, (), *args)
-    return (counter + 1, tuple(results), args)
-
-  carry_res = lax_control_flow.while_loop(cond, body, carry_init)
-  return carry_res[1]
-
-def _remat_translation_using_cond(*args, jaxpr: core.Jaxpr):
-  # Implements:
-  #  if(rng(0, 1) < 2)
-  #    return eval_jaxpr(*args)
-  #  else:
-  #    return 0
-  from jax._src.lax import control_flow as lax_control_flow
-
-  avals_out = tuple(v.aval for v in jaxpr.outvars)
-
-  def remat_comp(*args):
-    return tuple(core.eval_jaxpr(jaxpr, (), *args))
-  def dummy_comp(*args):
-    return tuple(map(_dummy_like, avals_out))
-
-  unif = lax_internal.rng_uniform(np.float32(0), np.float32(1), shape=())
-  return lax_control_flow.cond(unif < np.float32(2), remat_comp, dummy_comp, *args)
-
-def _remat_lowering(ctx, *args, jaxpr: core.Jaxpr, prevent_cse: bool,
-                   differentiated: bool, policy, is_gpu_platform=False):
+
+def _remat_lowering(
+    ctx,
+    *args,
+    jaxpr: core.Jaxpr,
+    prevent_cse: bool,
+    differentiated: bool,
+    policy,
+):
   jaxpr_args: Sequence[mlir.IrValues]
   if differentiated and prevent_cse:
-    # If we're using the loop or cond lowerings, use the slower lower_fun
-    # based path.
-    if not config.remat_opt_barrier.value:
-      return mlir.lower_fun(remat_expansion, multiple_results=True)(
-          ctx, *args, jaxpr=jaxpr, prevent_cse=prevent_cse,
-          differentiated=differentiated, policy=policy,
-          is_gpu_platform=is_gpu_platform)
-
     arg_types = map(mlir.aval_to_ir_type, ctx.avals_in)
     flat_args = mlir.flatten_ir_values(args)
     barrier_op = hlo.OptimizationBarrierOp(flat_args)
@@ -853,9 +798,8 @@ def _remat_lowering(ctx, *args, jaxpr: core.Jaxpr, prevent_cse: bool,
   ctx.set_tokens_out(tokens_out)
   return outs
 
+
 mlir.register_lowering(remat_p, _remat_lowering)
-mlir.register_lowering(remat_p, partial(_remat_lowering, is_gpu_platform=True),
-                       platform="gpu")
 
 
 def checkpoint_name(x, name):

jax/_src/config.py

@@ -1512,13 +1512,6 @@ def _update_disable_jit_thread_local(val):
     help=('Attempt constant folding during staging.'),
     include_in_jit_key=True)
 
-# This flag is temporary during rollout of the remat barrier.
-# TODO(parkers): Remove if there are no complaints.
-remat_opt_barrier = bool_state(
-    name='jax_remat_opt_barrier',
-    default=True,
-    help=('Enables using optimization-barrier op for lowering remat.'))
-
 enable_remat_opt_pass = bool_state(
     name='jax_compiler_enable_remat_pass',
     default=True,

jax/experimental/jax2tf/jax2tf.py

@@ -3173,12 +3173,11 @@ def select_one_carry(new_c: TfVal, c: TfVal, c_aval: core.ShapedArray) -> TfVal:
     lax_control_flow._scan_impl,
     extra_name_stack="scan")
 
-tf_impl_with_avals[ad_checkpoint.remat_p] = \
-  _convert_jax_impl(partial(ad_checkpoint.remat_expansion,
-                            # TODO: jax2tf cannot discriminate by platform
-                            is_gpu_platform=False),
-                    multiple_results=True,
-                    extra_name_stack="checkpoint")
+tf_impl_with_avals[ad_checkpoint.remat_p] = _convert_jax_impl(
+    ad_checkpoint.remat_expansion,
+    multiple_results=True,
+    extra_name_stack="checkpoint",
+)
 
 tf_impl[ad_checkpoint.name_p] = lambda x, *, name: x
 

jax/experimental/jax2tf/tests/jax2tf_test.py

@@ -832,11 +832,7 @@ def f(x1):
     arg = np.array(3.)
     f_tf = jax2tf.convert(jax.grad(remat_f))
     f_tf_hlo = self.TfToHlo(f_tf, arg)
-    if config.remat_opt_barrier.value:
-      self.assertRegex(f_tf_hlo, r"opt-barrier")
-    else:
-      self.assertRegex(f_tf_hlo,
-                       r'transpose/jax2tf_f_/jvp/checkpoint/cond/branch_1_fun/Sin')
+    self.assertRegex(f_tf_hlo, r"opt-barrier")
 
   def test_remat_free_var(self):
     def f(x):