Repository: jax-ml/jax
Pull request #27717 (2025-04-03 23:34:04)

add an `out_sharding` option to `jax.random.permutation`
Drop into `Auto` mode in the implementation.


Patch:
diff --git a/jax/_src/pallas/mosaic_gpu/lowering.py b/jax/_src/pallas/mosaic_gpu/lowering.py
index 80757ef69..d8e2083a8 100644
--- a/jax/_src/pallas/mosaic_gpu/lowering.py
+++ b/jax/_src/pallas/mosaic_gpu/lowering.py
@@ -1166,37 +1166,13 @@ def _swap_lowering_rule(
     raise TypeError(f"Can only store to references (got {x_smem}).")
   x_aval = ctx.avals_in[0]
   transforms = jax.tree.unflatten(tree, leaves)
-  transposed_value = value.layout == mgpu.WGMMA_TRANSPOSED_LAYOUT
-  x_smem, transforms = _handle_transforms(
-      x_smem, transforms, handle_transposes=not transposed_value
-  )
+  x_smem, transforms = _handle_transforms(x_smem, transforms)
   match transforms:
-    case (
-        gpu_core.UnswizzleRef(swizzle),
-        gpu_core.UntileRef(tiling),
-        *maybe_transpose,
-    ):
+    case (gpu_core.UnswizzleRef(swizzle), gpu_core.UntileRef(tiling)):
       if tiling != (8, swizzle // x_aval.dtype.itemsize):
         raise NotImplementedError("Tiling does not fit swizzle")
-
-      if transposed_value != bool(maybe_transpose):
-        raise ValueError(
-            "Either both the ref and the value are transposed or neither is."
-        )
-
-      if maybe_transpose:
-        if maybe_transpose != [gpu_core.TransposeRef((1, 0))]:
-          raise NotImplementedError(
-              f"Unsupported transforms: {transforms} ({maybe_transpose})"
-          )
-
-        x_smem = mgpu.memref_transpose(x_smem, (1, 0, 3, 2))
-
       old_value = mgpu.FragmentedArray.load_tiled(
-          x_smem,
-          is_signed=mgpu_utils.is_signed(x_aval.dtype),
-          swizzle=swizzle,
-          layout=value.layout,
+          x_smem, is_signed=mgpu_utils.is_signed(x_aval.dtype), swizzle=swizzle
       )
       value.store_tiled(x_smem, swizzle=swizzle)
       return old_value
diff --git a/jax/_src/random.py b/jax/_src/random.py
index 1d044ec11..a21cdf89a 100644
--- a/jax/_src/random.py
+++ b/jax/_src/random.py
@@ -556,7 +556,8 @@ def _randint(key, minval, maxval, shape, dtype) -> Array:
 def permutation(key: ArrayLike,
                 x: int | ArrayLike,
                 axis: int = 0,
-                independent: bool = False) -> Array:
+                independent: bool = False,
+                out_sharding=None) -> Array:
   """Returns a randomly permuted array or range.
 
   Args:
@@ -573,11 +574,17 @@ def permutation(key: ArrayLike,
   key, _ = _check_prng_key("permutation", key)
   check_arraylike("permutation", x)
   axis = canonicalize_axis(axis, np.ndim(x) or 1)
+  out_sharding = canonicalize_sharding(out_sharding, "permutation")
   if not np.ndim(x):
     if not np.issubdtype(lax.dtype(x), np.integer):
       raise TypeError("x must be an integer or at least 1-dimensional")
-    r = core.concrete_or_error(int, x, 'argument x of jax.random.permutation()')
-    return _shuffle(key, jnp.arange(r), axis)
+    r = core.concrete_or_error(int, x, "argument x of jax.random.permutation()")
+    return maybe_auto_axes(lambda key: _shuffle(key, jnp.arange(r), axis),
+                           out_sharding)(key)
+  return maybe_auto_axes(
+      _permutation, out_sharding, axis=axis, independent=independent)(key, x)
+
+def _permutation(key, x, axis, independent):
   if independent or np.ndim(x) == 1:
     return _shuffle(key, x, axis)
   ind = _shuffle(key, jnp.arange(x.shape[axis]), 0)  # type: ignore[union-attr]


Test patch:
diff --git a/tests/pallas/mosaic_gpu_test.py b/tests/pallas/mosaic_gpu_test.py
index 754e53255..06dfd453f 100644
--- a/tests/pallas/mosaic_gpu_test.py
+++ b/tests/pallas/mosaic_gpu_test.py
@@ -1372,44 +1372,6 @@ class PallasCallTest(PallasTest):
     x = jnp.full(shape, 42.0, jnp.float32)
     np.testing.assert_array_equal(kernel(), x)
 
-  def test_wgmma_transposed_layout(self):
-    """Tests that the result of wgmma can be store transposed using
-    the WGMMA_TRNASPOSED layout.
-    """
-
-    dtype = jnp.dtype(jnp.float16)
-    swizzle_elems = 128 // dtype.itemsize
-    shape = (128, 128)
-    @functools.partial(
-        pl.pallas_call,
-        out_shape=jax.ShapeDtypeStruct(shape, dtype),
-        out_specs=pl.BlockSpec(memory_space=plgpu.GMEM),
-        scratch_shapes=[
-            plgpu.SMEM(
-                shape, dtype,
-                transforms=(
-                    plgpu.TilingTransform((8, swizzle_elems)),
-                    plgpu.SwizzleTransform(128),
-                ),
-            )
-        ]
-    )
-    def kernel(o_ref, smem):
-      iota = plgpu.broadcasted_iota(
-          dtype, o_ref.shape, 0, layout=plgpu.Layout.WGMMA
-      ) * o_ref.shape[0]
-      iota += plgpu.broadcasted_iota(
-          dtype, o_ref.shape, 1, layout=plgpu.Layout.WGMMA
-      )
-
-      smem_trns = plgpu.transpose_ref(smem, (1, 0))
-      smem_trns[...] = plgpu.layout_cast(iota, plgpu.Layout.WGMMA_TRANSPOSED)
-      plgpu.commit_smem()
-      plgpu.copy_smem_to_gmem(smem, o_ref)
-
-    x = jnp.arange(128 * 128, dtype=dtype).reshape((128, 128)).T
-    np.testing.assert_array_equal(kernel(), x)
-
   def test_profiler(self):
     self.skip_if_wg_semantics()  # Transform inference fails.
 
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index d3d9cab7a..580cfcd7a 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -7332,6 +7332,45 @@ class ShardingInTypesTest(jtu.JaxTestCase):
     else:
       self.assertIn('mhlo.sharding = "{devices=[2,2]<=[4]}"}', lowered_text)
 
+  @jtu.with_user_mesh((4,), ('x',))
+  def test_random_permutation_1d(self, mesh):
+    @jax.jit
+    def f(key):
+      out = jax.random.permutation(key, 8, out_sharding=P('x'))
+      self.assertEqual(out.aval.sharding.spec, P('x'))
+      return out
+
+    key = jax.random.key(1)
+    out = f(key)
+    self.assertEqual(out.sharding, NamedSharding(mesh, P('x')))
+
+    lowered_text = f.lower(key).as_text()
+    if config.use_shardy_partitioner.value:
+      self.assertIn('sdy.sharding_constraint', lowered_text)
+      self.assertIn('<@mesh, [{"x"}]>', lowered_text)
+    else:
+      self.assertIn('mhlo.sharding = "{devices=[4]<=[4]}"}', lowered_text)
+
+  @jtu.with_user_mesh((2, 2), ('x', 'y'))
+  def test_random_permutation_2d(self, mesh):
+    @jax.jit
+    def f(key):
+      out = jax.random.permutation(key, jnp.arange(8 * 12).reshape(8, 12),
+                                   out_sharding=P('x', 'y'))
+      self.assertEqual(out.aval.sharding.spec, P('x', 'y'))
+      return out
+
+    key = jax.random.key(1)
+    out = f(key)
+    self.assertEqual(out.sharding, NamedSharding(mesh, P('x', 'y')))
+
+    lowered_text = f.lower(key).as_text()
+    if config.use_shardy_partitioner.value:
+      self.assertIn('sdy.sharding_constraint', lowered_text)
+      self.assertIn('<@mesh, [{"x"}, {"y"}]>', lowered_text)
+    else:
+      self.assertIn('mhlo.sharding = "{devices=[2,2]<=[4]}"}', lowered_text)
+
   @jtu.with_user_mesh((2, 2), ('x', 'y'))
   def test_random_normal(self, mesh):
     @jax.jit