From dc2677ea70c367f3047572011a639145ff46d8a7 Mon Sep 17 00:00:00 2001 From: Jake VanderPlas Date: Wed, 6 Dec 2023 15:01:57 -0800 Subject: [PATCH] Replace references to deprecated device_buffer attributes `jax.Array.device_buffer` and `jax.Array.device_buffers` will be deprecated as of jax version 0.4.22; see https://github.com/google/jax/pull/18844. PiperOrigin-RevId: 588553845 --- t5x/trainer_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/t5x/trainer_test.py b/t5x/trainer_test.py index 0cd576b0b..f1a048e8d 100644 --- 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()