Merge pull request #33634 from jakevdp:ad-checkpoint-doc

PiperOrigin-RevId: 839347694

Author: Google-ML-Automation

docs/jax.ad_checkpoint.rst

@@ -0,0 +1,11 @@
+``jax.ad_checkpoint`` module
+============================
+
+.. currentmodule:: jax.ad_checkpoint
+
+.. automodule:: jax.ad_checkpoint
+
+.. autosummary::
+    :toctree: _autosummary
+
+    checkpoint_name

docs/jax.rst

@@ -14,6 +14,7 @@ Subpackages
    jax.lax
    jax.random
    jax.sharding
+   jax.ad_checkpoint
    jax.debug
    jax.dlpack
    jax.distributed

tests/documentation_coverage_test.py

@@ -52,6 +52,7 @@ def jax_docs_dir() -> str:
 
 UNDOCUMENTED_APIS = {
   'jax': ['NamedSharding', 'P', 'Ref', 'Shard', 'ad_checkpoint', 'api_util', 'checkpoint_policies', 'core', 'custom_derivatives', 'custom_transpose', 'debug_key_reuse', 'device_put_replicated', 'device_put_sharded', 'effects_barrier', 'example_libraries', 'explain_cache_misses', 'experimental', 'extend', 'float0', 'freeze', 'fwd_and_bwd', 'host_count', 'host_id', 'host_ids', 'interpreters', 'jax', 'jax2tf_associative_scan_reductions', 'legacy_prng_key', 'lib', 'make_user_context', 'new_ref', 'no_execution', 'numpy_dtype_promotion', 'remat', 'remove_size_one_mesh_axis_from_type', 'softmax_custom_jvp', 'threefry_partitionable', 'tools', 'transfer_guard_device_to_device', 'transfer_guard_device_to_host', 'transfer_guard_host_to_device', 'version'],
+  'jax.ad_checkpoint': ['checkpoint', 'checkpoint_policies', 'print_saved_residuals', 'remat', 'Offloadable', 'Recompute', 'Saveable'],
   'jax.custom_batching': ['custom_vmap', 'sequential_vmap'],
   'jax.custom_derivatives': ['CustomVJPPrimal', 'SymbolicZero', 'closure_convert', 'custom_gradient', 'custom_jvp', 'custom_jvp_call_p', 'custom_vjp', 'custom_vjp_call_p', 'custom_vjp_primal_tree_values', 'linear_call', 'remat_opt_p', 'zero_from_primal'],
   'jax.custom_transpose': ['custom_transpose'],
@@ -75,7 +76,6 @@ def jax_docs_dir() -> str:
 # A list of modules to skip entirely, either because they cannot be imported
 # or because they are not expected to be documented.
 MODULES_TO_SKIP = [
-  "jax.ad_checkpoint",  # internal tools, not documented.
   "jax.api_util",  # internal tools, not documented.
   "jax.cloud_tpu_init",  # deprecated in JAX v0.8.1
   "jax.collect_profile",  # fails when xprof is not available.