Update jax/_src/ad_checkpoint.py

Co-authored-by: Jake Vanderplas <jakevdp@gmail.com>
jax-ml · Oct 15, 2024 · e954dfd · e954dfd
1 parent cc0b22a
commit e954dfd
Showing 1 changed file with 0 additions and 1 deletion.
diff --git a/jax/_src/ad_checkpoint.py b/jax/_src/ad_checkpoint.py
@@ -330,7 +330,6 @@ def remat(fun: Callable, *, prevent_cse: bool = True,
   return checkpoint(fun, prevent_cse=prevent_cse, policy=policy,
                     static_argnums=static_argnums)
 
-remat = checkpoint  # alias
 
 # This function is similar to api_util.argnums_partial, except the error
 # messages are specific to jax.remat (and thus more actionable), the