Merge pull request #16344 from JoeyTeng:cleanup-test-lax_control_flow…

…_test PiperOrigin-RevId: 540164431
google · Jun 14, 2023 · 9015920 · 9015920
2 parents 4d512c9 + c2c8314
commit 9015920
Showing 1 changed file with 0 additions and 7 deletions.
diff --git a/tests/lax_control_flow_test.py b/tests/lax_control_flow_test.py
@@ -92,13 +92,6 @@ def scan_with_for(f, *args, **kwargs):
 def scan_with_remat_for(f, *args, **kwargs):
   return jax.remat(lambda *args: for_loop.scan(f, *args, **kwargs))(*args)
 
-SCAN_IMPLS = [
-    (lax.scan, 'unroll1'),
-    (partial(lax.scan, unroll=2), 'unroll2'),
-    (scan_with_new_checkpoint , 'new_checkpoint'),
-    (scan_with_new_checkpoint2, 'new_checkpoint2'),
-]
-
 SCAN_IMPLS_WITH_FOR = [
     (lax.scan, 'unroll1'),
     (partial(lax.scan, unroll=2), 'unroll2'),