Avoid use of deprecated device_buffer attriutes of jax.Array

These have been deprecated as of JAX v0.4.22 PiperOrigin-RevId: 589237893
google-deepmind · Dec 9, 2023 · 226f810 · 226f810
1 parent cbac0f9
commit 226f810
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/chex/_src/asserts.py b/chex/_src/asserts.py
@@ -1226,7 +1226,7 @@ def _assert_fn(path, leaf):
       # Check that the leaf is a ShardedArray.
       if isinstance(leaf, jax.Array):
         if _check_sharding(leaf):
-          shards = tuple(buf.device() for buf in leaf.device_buffers)
+          shards = tuple(shard.device for shard in leaf.addressable_shards)
           if shards != devices:
             errors.append(
                 f"Tree leaf '{_ai.format_tree_path(path)}' is sharded "