Remove unused _remat_static_argnums import.

google · Sep 26, 2022 · 8bcf358 · 8bcf358
1 parent 28672cc
commit 8bcf358
Showing 1 changed file with 0 additions and 1 deletion.
diff --git a/jax/_src/api.py b/jax/_src/api.py
@@ -73,7 +73,6 @@
                                      process_count, host_id, host_ids,
                                      host_count, default_backend)
 from jax.ad_checkpoint import checkpoint_policies, checkpoint as new_checkpoint
-from jax._src.ad_checkpoint import _remat_static_argnums
 from jax.core import ShapedArray, raise_to_shaped
 from jax.custom_batching import custom_vmap
 from jax.custom_derivatives import (closure_convert, custom_gradient, custom_jvp,