Repository: jax-ml/jax
Pull request #27447 (2025-03-25 22:31:21)

[sharding_in_types] Add `out_sharding` to `jax.random.uniform`.
[sharding_in_types] Add `out_sharding` to `jax.random.uniform`.

Drop into `Auto` mode inside for implementation.

Co-authored-by: Roy Frostig <frostig@google.com>



Patch:
diff --git a/jax/_src/random.py b/jax/_src/random.py
index c0663dc67..2d315ed0c 100644
--- a/jax/_src/random.py
+++ b/jax/_src/random.py
@@ -38,6 +38,8 @@ from jax._src.api import jit, vmap
 from jax._src.interpreters import ad
 from jax._src.interpreters import batching
 from jax._src.interpreters import mlir
+from jax._src.sharding_impls import canonicalize_sharding
+from jax._src.pjit import auto_axes
 from jax._src.lax import lax as lax_internal
 from jax._src.numpy.lax_numpy import _convert_and_clip_integer
 from jax._src.numpy.util import _arraylike, check_arraylike, promote_dtypes_inexact
@@ -379,7 +381,8 @@ def uniform(key: ArrayLike,
             shape: Shape = (),
             dtype: DTypeLikeFloat = float,
             minval: RealArray = 0.,
-            maxval: RealArray = 1.) -> Array:
+            maxval: RealArray = 1.,
+            out_sharding=None) -> Array:
   """Sample uniform random values in [minval, maxval) with given shape/dtype.
 
   Args:
@@ -397,14 +400,21 @@ def uniform(key: ArrayLike,
   key, _ = _check_prng_key("uniform", key)
   dtypes.check_user_dtype_supported(dtype)
   shape = core.canonicalize_shape(shape)
+  out_sharding = canonicalize_sharding(out_sharding, "uniform")
 
   if not dtypes.issubdtype(dtype, np.floating):
     raise ValueError(f"dtype argument to `uniform` must be a float dtype, "
                      f"got {dtype}")
   dtype = dtypes.canonicalize_dtype(dtype)
-  return _uniform(key, shape, dtype, minval, maxval)
+  return _uniform_auto(key, shape, dtype, minval, maxval, out_sharding)
+
+@partial(jit, static_argnums=(1, 2, 5))
+def _uniform_auto(key, shape, dtype, minval, maxval, out_sharding) -> Array:
+  if out_sharding is None:
+    return _uniform(key, shape, dtype, minval, maxval)
+  def f(key, minval, maxval): return _uniform(key, shape, dtype, minval, maxval)
+  return auto_axes(f, out_shardings=out_sharding)(key, minval, maxval)
 
-@partial(jit, static_argnums=(1, 2))
 def _uniform(key, shape, dtype, minval, maxval) -> Array:
   _check_shape("uniform", shape)
   if not jnp.issubdtype(dtype, np.floating):


Test patch:
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index ed1f9e9b6..528384358 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -7226,6 +7226,25 @@ class ShardingInTypesTest(jtu.JaxTestCase):
     out = f(key)
     self.assertEqual(out.sharding, NamedSharding(mesh, P()))
 
+  @jtu.with_user_mesh((2, 2), ('x', 'y'))
+  def test_random_uniform(self, mesh):
+    @jax.jit
+    def f(key):
+      out = jax.random.uniform(key, shape=(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.pytest_mark_if_available('multiaccelerator')
 class PJitErrorTest(jtu.JaxTestCase):