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 jax-ml/jax#18844. PiperOrigin-RevId: 588553845
google-deepmind · Dec 7, 2023 · a6381cd · a6381cd
1 parent cbac0f9
commit a6381cd
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 "