Make addressable_shards a property like local_shards

PiperOrigin-RevId: 477358276
google · Sep 28, 2022 · 96a85bd · 96a85bd
1 parent 9489068
commit 96a85bd
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/jax/experimental/global_device_array.py b/jax/experimental/global_device_array.py
@@ -394,6 +394,7 @@ def _create_local_shards(self) -> Sequence[Shard]:
   def local_shards(self) -> Sequence[Shard]:
     return self._create_local_shards()
 
+  @pxla.maybe_cached_property
   def addressable_shards(self) -> Sequence[Shard]:
     return self.local_shards