Replace references to deprecated jax array attributes device_buffer a…

…nd device_buffers PiperOrigin-RevId: 588553845
google-deepmind · Dec 6, 2023 · d4414e2 · d4414e2
1 parent cbac0f9
commit d4414e2
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 "