[Pallas:MGPU] Expose fragmented_array.Replicated as part of the public API.

PiperOrigin-RevId: 842675493

Author: allanrenucci

jax/experimental/pallas/mosaic_gpu.py

@@ -87,6 +87,7 @@
 from jax._src.pallas.mosaic_gpu.primitives import wgmma_wait as wgmma_wait
 from jax._src.pallas.mosaic_gpu.torch import as_torch_kernel as as_torch_kernel
 from jax.experimental.mosaic.gpu.core import LoweringSemantics as LoweringSemantics
+from jax.experimental.mosaic.gpu.fragmented_array import Replicated as Replicated
 from jax.experimental.mosaic.gpu.fragmented_array import Tiling as Tiling
 
 

tests/pallas/mosaic_gpu_test.py

@@ -2563,6 +2563,23 @@ def kernel(dst, collective_barrier):
     )()
     np.testing.assert_array_equal(y, np.ones((), dtype=np.int32))
 
+  def test_replicated_layout(self):
+    shape = (32,)
+    @functools.partial(
+        self.pallas_call,
+        out_shape=jax.ShapeDtypeStruct(shape, jnp.float32),
+    )
+    def kernel(src_ref, dst_ref):
+      layout = plgpu.Layout.TILED(
+        plgpu.Tiling(((32,), (1,))),
+        warp_dims=(plgpu.Replicated(4),),
+        lane_dims=(-2,),
+        vector_dim=-1,
+      )
+      dst_ref[...] = plgpu.load(src_ref, (), layout=layout)
+    src = jnp.arange(shape[0], dtype=jnp.float32)
+    np.testing.assert_array_equal(kernel(src), src)
+
 
 class PallasCallWarpPrimitiveSemanticsTest(PallasTest):
   def setUp(self):