Merge pull request #12764 from midjourney:gda_sharding_property

PiperOrigin-RevId: 480628909

jax/experimental/global_device_array.py

@@ -18,6 +18,7 @@
 import numpy as np
 from typing import Callable, Sequence, Tuple, Union, Mapping, Optional, List, Dict, NamedTuple
 
+import jax
 from jax import core
 from jax._src import api_util
 from jax._src.lib import xla_bridge as xb
@@ -454,9 +455,13 @@ def block_until_ready(self):
       for db in self._device_buffers:
         db.block_until_ready()
     else:
-      self._sharded_buffer.block_until_ready() # type: ignore
+      self._sharded_buffer.block_until_ready()  # type: ignore
     return self
 
+  @property
+  def sharding(self):
+    return jax.sharding.MeshPspecSharding(self._global_mesh, self.mesh_axes)
+
   @classmethod
   def from_callback(cls, global_shape: Shape, global_mesh: pxla.Mesh,
                     mesh_axes: MeshAxes, data_callback: Callable[[Index],

tests/global_device_array_test.py

@@ -88,6 +88,7 @@ def cb(index):
     self.assertArraysEqual(gda.local_data(0),
                            global_input_data[expected_index[0]])
     self.assertEqual(gda.local_shards[1].index, expected_index[1])
+    self.assertIsInstance(gda.sharding, jax.sharding.MeshPspecSharding)
     self.assertArraysEqual(gda.local_data(1),
                            global_input_data[expected_index[1]])
     self.assertEqual(gda.local_data(0).shape, expected_shard_shape)