Avoid use of deprecated device_buffer attriutes of jax.Array

These have been deprecated as of JAX v0.4.22 PiperOrigin-RevId: 589237955
google-research · Dec 11, 2023 · 23d2f85 · 23d2f85
1 parent 5f7aa29
commit 23d2f85
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/t5x/trainer_test.py b/t5x/trainer_test.py
@@ -175,7 +175,7 @@ def test_timer_blocking_on_donated_buffer(self):
     mm._duration_timer._start_future.result()
 
     # Deleted/donated.
-    x.device_buffer.delete()
+    x.addressable_data(0).delete()
     mm.start_duration_timer(block_on=x)
     mm._duration_timer._start_future.result()