Merge pull request #23853 from zhenying-liu:remat-scan

PiperOrigin-RevId: 679365040

tests/memories_test.py

@@ -35,6 +35,7 @@
                                      TransferToMemoryKind, PartitionSpec as P)
 from jax.experimental.compute_on import compute_on
 from jax.experimental.shard_map import shard_map
+from jax._src.lib import xla_extension_version
 import numpy as np
 
 config.parse_flags_with_absl()
@@ -1567,8 +1568,8 @@ def g(ys, _):
       self.assertGreater(compiled_stats.host_temp_size_in_bytes, 0)
 
   def test_remat_scan_layout_change_offloadable(self):
-    if not jtu.test_device_matches(["tpu"]):
-      self.skipTest("Remat scan does not work on GPU backend.")
+    if jtu.test_device_matches(["gpu"]) and xla_extension_version < 289:
+      self.skipTest("Requires xla_extension_version >= 289")
     mesh = jtu.create_mesh((2,), ("x",))
     shape = (256, 128)
     np_inp = np.arange(math.prod(shape), dtype=np.float32).reshape(shape)
@@ -1602,6 +1603,10 @@ def g(ys, _):
       self.assertIn('S(5)', compiled_text)
       self.assertNotRegex(compiled_text, r"copy-start.*S\(5\)")
       self.assertNotRegex(compiled_text, r"copy-done.*S\(5\)")
+      self.assertRegex(compiled_text, r"dynamic-update-slice-start.*S\(5\)")
+      self.assertRegex(compiled_text, r"dynamic-update-slice-done.*S\(5\)")
+      self.assertRegex(compiled_text, r"dynamic-slice-start.*S\(5\)")
+      self.assertRegex(compiled_text, r"dynamic-slice-done.*S\(5\)")
 
     compiled_stats = compiled_f.memory_analysis()
     if compiled_stats is not None: