Document all checkpoint policies in one place, on the JAX public API page.

Author: carlosgmartin

docs/gradient-checkpointing.md

@@ -411,22 +411,7 @@ The code defines a function `f` that which applies checkpointing with a custom p
 
 #### List of policies
 
-The policies are:
-* `everything_saveable` (the default strategy, as if `jax.checkpoint` were not being used at all)
-* `nothing_saveable` (i.e. rematerialize everything, as if a custom policy were not being used at all)
-* `dots_saveable` or its alias `checkpoint_dots`
-* `dots_with_no_batch_dims_saveable` or its alias `checkpoint_dots_with_no_batch_dims`
-* `save_anything_but_these_names` (save any values except for the output of
-  `checkpoint_name` with any of the names given)
-* `save_any_names_but_these` (save only named values, i.e. any outputs of
-  `checkpoint_name`, except for those with the names given)
-* `save_only_these_names` (save only named values, and only among the names
-  given)
-* `offload_dot_with_no_batch_dims` same as `dots_with_no_batch_dims_saveable`,
-  but offload to CPU memory instead of recomputing.
-* `save_and_offload_only_these_names` same as `save_only_these_names`, but
-  offload to CPU memory instead of recomputing.
-* `save_from_both_policies(policy_1, policy_2)` (like a logical `or`, so that a residual is saveable if it is saveable according to `policy_1` _or_ `policy_2`)
+The policies can be found [here](https://docs.jax.dev/en/latest/jax.html#checkpoint-policies).
 
 Policies only indicate what is saveable; a value is only saved if it's actually needed by the backward pass.
 

docs/jax.rst

@@ -263,3 +263,21 @@ Miscellaneous
     print_environment_info
     live_arrays
     clear_caches
+
+Checkpoint policies
+-------------------
+
+.. autosummary::
+  :toctree: _autosummary
+
+    checkpoint_policies.everything_saveable
+    checkpoint_policies.nothing_saveable
+    checkpoint_policies.dots_saveable
+    checkpoint_policies.checkpoint_dots
+    checkpoint_policies.dots_with_no_batch_dims_saveable
+    checkpoint_policies.checkpoint_dots_with_no_batch_dims
+    checkpoint_policies.save_any_names_but_these
+    checkpoint_policies.save_only_these_names
+    checkpoint_policies.offload_dot_with_no_batch_dims
+    checkpoint_policies.save_and_offload_only_these_names
+    checkpoint_policies.save_from_both_policies

docs/notebooks/autodiff_remat.ipynb

@@ -812,17 +812,7 @@
    "source": [
     "Another policy which refers to names is `jax.checkpoint_policies.save_only_these_names`.\n",
     "\n",
-    "Some of the policies are:\n",
-    "* `everything_saveable` (the default strategy, as if `jax.checkpoint` were not being used at all)\n",
-    "* `nothing_saveable` (i.e. rematerialize everything, as if a custom policy were not being used at all)\n",
-    "* `dots_saveable` or its alias `checkpoint_dots`\n",
-    "* `dots_with_no_batch_dims_saveable` or its alias `checkpoint_dots_with_no_batch_dims`\n",
-    "* `save_anything_but_these_names` (save any values except for the output of\n",
-    "  `checkpoint_name` with any of the names given)\n",
-    "* `save_any_names_but_these` (save only named values, i.e. any outputs of\n",
-    "  `checkpoint_name`, except for those with the names given)\n",
-    "* `save_only_these_names` (save only named values, and only among the names\n",
-    "  given)\n",
+    "A list of policies can be found [here](https://docs.jax.dev/en/latest/jax.html#checkpoint-policies).\n",
     "\n",
     "Policies only indicate what is saveable; a value is only saved if it's actually needed by the backward pass."
    ]

docs/notebooks/autodiff_remat.md

@@ -396,18 +396,6 @@ print_saved_residuals(loss_checkpoint2, params, x, y)
 
 Another policy which refers to names is `jax.checkpoint_policies.save_only_these_names`.
 
-Some of the policies are:
-* `everything_saveable` (the default strategy, as if `jax.checkpoint` were not being used at all)
-* `nothing_saveable` (i.e. rematerialize everything, as if a custom policy were not being used at all)
-* `dots_saveable` or its alias `checkpoint_dots`
-* `dots_with_no_batch_dims_saveable` or its alias `checkpoint_dots_with_no_batch_dims`
-* `save_anything_but_these_names` (save any values except for the output of
-  `checkpoint_name` with any of the names given)
-* `save_any_names_but_these` (save only named values, i.e. any outputs of
-  `checkpoint_name`, except for those with the names given)
-* `save_only_these_names` (save only named values, and only among the names
-  given)
-
 Policies only indicate what is saveable; a value is only saved if it's actually needed by the backward pass.
 
 +++ {"id": "lixGsLNwxQo7"}

jax/_src/ad_checkpoint.py

@@ -56,30 +56,37 @@
 ### Policies
 
 def everything_saveable(*_, **__) -> bool:
-  # This is the effective policy without any use of jax.remat.
+  """The default strategy, as if ``jax.checkpoint`` were not being used at all.
+
+  This is the effective policy without any use of jax.remat."""
   return True
 
 def nothing_saveable(*_, **__) -> bool:
-  # This is the effective policy when using jax.remat without explicit policy.
+  """Rematerialize everything, as if a custom policy were not being used at all.
+
+  This is the effective policy when using jax.remat without explicit policy."""
   return False
 
 def dots_saveable(prim, *_, **__) -> bool:
-  # Matrix multiplies are expensive, so let's save them (and nothing else).
+  """Matrix multiplies are expensive, so let's save them (and nothing else)."""
   return prim in {lax_internal.dot_general_p,
                   lax_convolution.conv_general_dilated_p}
 checkpoint_dots = dots_saveable
 
 def dots_with_no_batch_dims_saveable(prim, *_, **params) -> bool:
-  # This is a useful heuristic for transformers.
+  """This is a useful heuristic for transformers."""
   if prim is lax_internal.dot_general_p:
     (_, _), (lhs_b, rhs_b) = params['dimension_numbers']
     if not lhs_b and not rhs_b:
       return True
   return False
 
 def offload_dot_with_no_batch_dims(offload_src, offload_dst):
+  """Same as ``dots_with_no_batch_dims_saveable``, but offload to CPU memory
+  instead of recomputing.
+
+  This is a useful heuristic for transformers."""
   def policy(prim, *_, **params):
-    # This is a useful heuristic for transformers.
     if prim is lax_internal.dot_general_p:
       (_, _), (lhs_b, rhs_b) = params['dimension_numbers']
       if not lhs_b and not rhs_b:
@@ -100,7 +107,8 @@ def policy(prim, *_, **params):
   return policy
 
 def save_any_names_but_these(*names_not_to_save):
-  """Save only named values, excluding the names given."""
+  """Save only named values, i.e. any outputs of `checkpoint_name`, excluding
+  the names given."""
   names_not_to_save = frozenset(names_not_to_save)
   def policy(prim, *_, **params):
     if prim is name_p:
@@ -120,6 +128,8 @@ def policy(prim, *_, **params):
 def save_and_offload_only_these_names(
     *, names_which_can_be_saved, names_which_can_be_offloaded,
     offload_src, offload_dst):
+  """Same as ``save_only_these_names``, but offload to CPU memory instead of
+  recomputing."""
   names_which_can_be_saved = set(names_which_can_be_saved)
   names_which_can_be_offloaded = set(names_which_can_be_offloaded)
   intersection = names_which_can_be_saved.intersection(names_which_can_be_offloaded)
@@ -140,7 +150,9 @@ def policy(prim, *_, **params):
 
 
 def save_from_both_policies(policy_1, policy_2):
+  """Logical OR of the given policies.
 
+  A residual is saveable iff it is saveable according to either policy."""
   def policy(prim, *args, **params):
     out1 = policy_1(prim, *args, **params)
     out2 = policy_2(prim, *args, **params)