Merge pull request #33660 from jakevdp:ad_checkpoint

PiperOrigin-RevId: 839358516

Author: Google-ML-Automation

docs/gradient-checkpointing.md

@@ -359,7 +359,7 @@ Another policy which refers to names is `jax.checkpoint_policies.save_only_these
 You may consider offloading to CPU memory instead of recomputing when checkpointing to save accelerator memory. `jax.checkpoint_policies.offload_dot_with_no_batch_dims` can offload the results of matrix multiplications with no batch dimensions to the CPU.
 
 ```{code-cell}
-from jax.ad_checkpoint import checkpoint
+from jax import checkpoint
 
 def checkpoint_offload_dot_with_no_batch_dims(self):
   policy = jax.checkpoint_policies.offload_dot_with_no_batch_dims(
@@ -380,7 +380,8 @@ def checkpoint_offload_dot_with_no_batch_dims(self):
 One of JAX's checkpoint policies allows specified checkpoint names to be offloaded to CPUs. This policy is implemented through `jax.checkpoint_policies.save_and_offload_only_these_names`, which has four arguments: `names_which_can_be_saved`, `names_which_can_be_offloaded`, the offloading source, and destination. Names listed in `names_which_can_be_saved` are kept on the device, names listed in `names_which_can_be_offloaded` are moved to CPU memory, and other names or operations without names are recomputed. For example, if we have checkpoint names `y`, `z`, and `w`, `y` can be saved on the device, `z` can be offloaded to CPU memory, and `w` can be recomputed.
 
 ```{code-cell}
-from jax.ad_checkpoint import checkpoint, checkpoint_name
+from jax import checkpoint
+from jax.ad_checkpoint import checkpoint_name
 from jax._src import test_util as jtu
 
 def checkpoint_names_saved_offloaded_recomputed(self):

jax/__init__.py

@@ -85,8 +85,9 @@
 from jax._src.core import typeof as typeof
 from jax._src.api import effects_barrier as effects_barrier
 from jax._src.api import block_until_ready as block_until_ready
-from jax._src.ad_checkpoint import checkpoint_wrapper as checkpoint  # noqa: F401
+from jax._src.ad_checkpoint import checkpoint as checkpoint
 from jax._src.ad_checkpoint import checkpoint_policies as checkpoint_policies
+from jax._src.ad_checkpoint import remat as remat
 from jax._src.api import clear_caches as clear_caches
 from jax._src.api import copy_to_host_async as copy_to_host_async
 from jax._src.custom_derivatives import closure_convert as closure_convert
@@ -127,7 +128,6 @@
 from jax._src.xla_bridge import process_index as process_index
 from jax._src.xla_bridge import process_indices as process_indices
 from jax._src.callback import pure_callback as pure_callback
-from jax._src.ad_checkpoint import checkpoint_wrapper as remat  # noqa: F401
 from jax._src.core import ShapeDtypeStruct as ShapeDtypeStruct
 from jax._src.api import value_and_grad as value_and_grad
 from jax._src.api import vjp as vjp

jax/_src/ad_checkpoint.py

@@ -15,7 +15,6 @@
 from __future__ import annotations
 
 from collections.abc import Callable, Sequence
-import functools
 from functools import partial
 import logging
 from typing import Any
@@ -203,7 +202,7 @@ def policy(prim, *args, **params):
 def checkpoint(fun: Callable, *, prevent_cse: bool = True,
                policy: Callable[..., bool] | None = None,
                static_argnums: int | tuple[int, ...] = (),
-               ) -> Callable:
+               concrete: bool = False) -> Callable:
   """Make ``fun`` recompute internal linearization points when differentiated.
 
   The :func:`jax.checkpoint` decorator, aliased to :func:`jax.remat`, provides a
@@ -257,6 +256,8 @@ def checkpoint(fun: Callable, *, prevent_cse: bool = True,
       returns a boolean indicating whether the corresponding output value(s) can
       be saved as residuals (or instead must be recomputed in the (co)tangent
       computation if needed).
+    concrete: Optional boolean; deprecated. Passing a non-False value will
+      result in NotImplementedError.
 
   Returns:
     A function (callable) with the same input/output behavior as ``fun`` but
@@ -344,6 +345,11 @@ def foo(x, y):
   ``jax.ensure_compile_time_eval``), it may be easier to compute some values
   outside the :func:`jax.checkpoint`-decorated function and then close over them.
   """
+  if concrete:
+    raise NotImplementedError(
+      "The concrete option to jax.checkpoint has been deprecated. In its"
+      " place please use ``static_argnums``; for details refer to"
+      " https://docs.jax.dev/en/latest/jep/11830-new-remat-checkpoint.html.")
   if isinstance(static_argnums, int):
     static_argnums = static_argnums,
   if isinstance(prevent_cse, list):
@@ -373,7 +379,14 @@ def fun_remat(*args, **kwargs):
     return tree_unflatten(out_tree, out_flat)
   return fun_remat
 
-remat = checkpoint  # alias
+
+def remat(fun: Callable, *, prevent_cse: bool = True,
+          policy: Callable[..., bool] | None = None,
+          static_argnums: int | tuple[int, ...] = (),
+          concrete: bool = False) -> Callable:
+  """Alias of :func:`jax.checkpoint`."""
+  return checkpoint(fun, prevent_cse=prevent_cse, policy=policy,
+                    static_argnums=static_argnums, concrete=concrete)
 
 # This function is similar to api_util.argnums_partial, except the error
 # messages are specific to jax.remat (and thus more actionable), the
@@ -894,66 +907,6 @@ def name_batcher(args, dims, *, name):
 batching.primitive_batchers[name_p] = name_batcher
 
 
-@functools.wraps(checkpoint)
-def checkpoint_wrapper(
-    fun: Callable,
-    *,
-    concrete: bool = False,
-    prevent_cse: bool = True,
-    static_argnums: int | tuple[int, ...] = (),
-    policy: Callable[..., bool] | None = None,
-) -> Callable:
-  if concrete:
-    msg = ("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.\n"
-           "\n"
-           "For example, if using `concrete=True` for an `is_training` flag:\n"
-           "\n"
-           "  from functools import partial\n"
-           "\n"
-           "  @partial(jax.checkpoint, concrete=True)\n"
-           "  def foo(x, is_training):\n"
-           "    if is_training:\n"
-           "      return f(x)\n"
-           "    else:\n"
-           "      return g(x)\n"
-           "\n"
-           "replace it with a use of `static_argnums`:\n"
-           "\n"
-           "  @partial(jax.checkpoint, static_argnums=(1,))\n"
-           "  def foo(x, is_training):\n"
-           "    ...\n"
-           "\n"
-           "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`\n:"
-           "\n"
-           "  @partial(jax.checkpoint, concrete=True)\n"
-           "  def foo(x, y):\n"
-           "    if y > 0:\n"
-           "      return f(x)\n"
-           "    else:\n"
-           "      return g(x)\n"
-           "\n"
-           "with this combination of `static_argnums` and "
-           "`jax.ensure_compile_time_eval()`:\n"
-           "\n"
-           "  @partial(jax.checkpoint, static_argnums=(1,))\n"
-           "  def foo(x, y):\n"
-           "    with jax.ensure_compile_time_eval():\n"
-           "      y_pos = y > 0\n"
-           "    if y_pos:\n"
-           "      return f(x)\n"
-           "    else:\n"
-           "      return g(x)\n"
-           "\n"
-           "See https://docs.jax.dev/en/latest/jep/11830-new-remat-checkpoint.html\n")
-    raise NotImplementedError(msg)
-  return checkpoint(fun, prevent_cse=prevent_cse, policy=policy,
-                    static_argnums=static_argnums)
-
-
 @discharge.register_discharge_rule(remat_p)
 def _remat_state_discharge_rule(
     in_avals, out_avals, *args, jaxpr, **params):

jax/ad_checkpoint.py

@@ -13,14 +13,36 @@
 # limitations under the License.
 
 from jax._src.ad_checkpoint import (
-  checkpoint as checkpoint,
+  checkpoint as _deprecated_checkpoint,
   checkpoint_policies as checkpoint_policies,
   checkpoint_name as checkpoint_name,
   print_saved_residuals as print_saved_residuals,
-  remat as remat,
 )
 from jax._src.interpreters.partial_eval import (
   Recompute as Recompute,
   Saveable as Saveable,
   Offloadable as Offloadable,
 )
+
+_deprecations = {
+  # Added for v0.8.2
+  "checkpoint": (
+    "jax.ad_checkpoint.checkpoint is deprecated; use jax.checkpoint instead.",
+    _deprecated_checkpoint
+  ),
+  "remat": (
+    "jax.ad_checkpoint.remat is deprecated; use jax.remat instead.",
+    _deprecated_checkpoint
+  ),
+}
+
+import typing as _typing
+if _typing.TYPE_CHECKING:
+  checkpoint = _deprecated_checkpoint
+  remat = _deprecated_checkpoint
+else:
+  from jax._src.deprecations import deprecation_getattr as _deprecation_getattr
+  __getattr__ = _deprecation_getattr(__name__, _deprecations)
+  del _deprecation_getattr
+del _typing
+del _deprecated_checkpoint

jax/experimental/jax2tf/tests/jax2tf_test.py

@@ -831,7 +831,7 @@ def f(x1):
       x3 = jnp.sin(x2)
       x4 = jnp.sin(x3)
       return x4
-    remat_f = ad_checkpoint.checkpoint(f)
+    remat_f = jax.checkpoint(f)
 
     # The computation of grad_f computes "sin" 5 times, 3 for the forward pass
     # and then to rematerialize "x2" and "x3" in the backward pass.
@@ -844,7 +844,7 @@ def test_remat_free_var(self):
     def f(x):
       y = 2 * x
 
-      @ad_checkpoint.checkpoint
+      @jax.checkpoint
       def g():
         return y