Repository: jax-ml/jax
Pull request #27730 (2025-04-04 05:35:59)

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


Patch:
diff --git a/jax/_src/pallas/mosaic/lowering.py b/jax/_src/pallas/mosaic/lowering.py
index 0ca298c88..6c8b3c646 100644
--- a/jax/_src/pallas/mosaic/lowering.py
+++ b/jax/_src/pallas/mosaic/lowering.py
@@ -3531,14 +3531,8 @@ def _semaphore_wait_lowering_rule(ctx: LoweringRuleContext, *args, args_tree):
   return []
 lowering_rules[primitives.semaphore_wait_p] = _semaphore_wait_lowering_rule
 
-
-def _dma_start_lowering_rule(
-    ctx: LoweringRuleContext,
-    *args,
-    tree,
-    device_id_type: primitives.DeviceIdType,
-    priority: int,
-):
+def _dma_start_lowering_rule(ctx: LoweringRuleContext, *args, tree,
+                             device_id_type: primitives.DeviceIdType):
   (
       src_ref,
       src_transforms,
@@ -3570,17 +3564,10 @@ def _dma_start_lowering_rule(
   sem, _ = _transform_ref(sem, sem_aval.dtype, sem_aval.shape, sem_transforms)
   if device_id is not None:
     device_id = _device_id_to_logical(ctx, device_id, device_id_type)
-  tpu.enqueue_dma(
-      src_ref,
-      dst_ref,
-      sem,
-      source_semaphore=src_sem,
-      device_id=device_id,
-      priority=priority,
-  )
-  return []
-
+  tpu.enqueue_dma(src_ref, dst_ref, sem, source_semaphore=src_sem,
+                  device_id=device_id)
 
+  return []
 lowering_rules[tpu_primitives.dma_start_p] = _dma_start_lowering_rule
 
 
diff --git a/jax/_src/pallas/mosaic/primitives.py b/jax/_src/pallas/mosaic/primitives.py
index 59856c0ca..c50a21218 100644
--- a/jax/_src/pallas/mosaic/primitives.py
+++ b/jax/_src/pallas/mosaic/primitives.py
@@ -208,14 +208,9 @@ class AsyncCopyDescriptor:
           self.device_id,
       ))
 
-  def start(self, priority: int = 0):
+  def start(self):
     flat_args, tree = self._get_args_and_tree()
-    dma_start_p.bind(
-        *flat_args,
-        tree=tree,
-        device_id_type=self.device_id_type,
-        priority=priority,
-    )
+    dma_start_p.bind(*flat_args, tree=tree, device_id_type=self.device_id_type)
 
   def wait(self):
     if self.is_remote:
@@ -244,9 +239,7 @@ dma_start_p = jax_core.Primitive('dma_start')
 dma_start_p.multiple_results = True
 
 @dma_start_p.def_effectful_abstract_eval
-def _dma_start_abstract_eval(*args, tree, device_id_type, priority):
-  if priority < 0:
-    raise ValueError(f"DMA start priority must be non-negative: {priority}")
+def _dma_start_abstract_eval(*args, tree, device_id_type):
   (
       src_ref_aval,
       src_transforms_avals,
@@ -281,7 +274,6 @@ def _dma_start_pp_eqn(eqn: jax_core.JaxprEqn,
                       settings: jax_core.JaxprPpSettings):
   invars = eqn.invars
   tree = eqn.params["tree"]
-  priority = eqn.params["priority"]
   (
       src_ref,
       src_transforms,
@@ -298,7 +290,7 @@ def _dma_start_pp_eqn(eqn: jax_core.JaxprEqn,
   if src_sem or device_id:
     return jax_core._pp_eqn(eqn, context, settings)
   return pp.concat([
-      pp.text(f"dma_start(p{priority})"),
+      pp.text("dma_start"),
       pp.text(" "),
       sp.pp_ref_transforms(context, src_ref, src_transforms),
       pp.text(" -> "),
@@ -309,12 +301,8 @@ def _dma_start_pp_eqn(eqn: jax_core.JaxprEqn,
 
 jax_core.pp_eqn_rules[dma_start_p] = _dma_start_pp_eqn
 
-
-def dma_start_partial_discharge_rule(
-    should_discharge, in_avals, out_avals, *args, tree, device_id_type, priority
-):
-  # Note: we ignore the DMA priority in discharge rules.
-  del priority
+def dma_start_partial_discharge_rule(should_discharge, in_avals, out_avals,
+                                     *args, tree, device_id_type):
   (
       src_ref,
       src_transforms,
@@ -473,7 +461,6 @@ def dma_start_partial_discharge_rule(
 
   return new_vals, []
 
-
 state_discharge.register_partial_discharge_rule(dma_start_p)(dma_start_partial_discharge_rule)
 
 
@@ -563,7 +550,6 @@ def _get_ref_and_transforms(ref):
     return ref.ref, ref.transforms
   return ref, ()
 
-
 def make_async_copy(src_ref, dst_ref, sem):
   """Issues a DMA copying from src_ref to dst_ref."""
   src_ref, src_transforms = _get_ref_and_transforms(src_ref)
@@ -582,14 +568,12 @@ def make_async_copy(src_ref, dst_ref, sem):
       primitives.DeviceIdType.MESH,
   )
 
-
-def async_copy(src_ref, dst_ref, sem, *, priority: int = 0):
+def async_copy(src_ref, dst_ref, sem):
   """Issues a DMA copying from src_ref to dst_ref."""
   copy_descriptor = make_async_copy(src_ref, dst_ref, sem)
-  copy_descriptor.start(priority=priority)
+  copy_descriptor.start()
   return copy_descriptor
 
-
 def make_async_remote_copy(src_ref, dst_ref, send_sem, recv_sem, device_id,
                            device_id_type: primitives.DeviceIdType = primitives.DeviceIdType.MESH):
   """Creates a description of a remote copy operation.
diff --git a/jax/_src/random.py b/jax/_src/random.py
index 5cbd966e7..e632d4a9a 100644
--- a/jax/_src/random.py
+++ b/jax/_src/random.py
@@ -851,7 +851,8 @@ def truncated_normal(key: ArrayLike,
                      lower: RealArray,
                      upper: RealArray,
                      shape: Shape | None = None,
-                     dtype: DTypeLikeFloat = float) -> Array:
+                     dtype: DTypeLikeFloat = float,
+                     *, out_sharding=None) -> Array:
   r"""Sample truncated standard normal random values with given shape and dtype.
 
   The values are returned according to the probability density function:
@@ -882,12 +883,14 @@ def truncated_normal(key: ArrayLike,
   if shape is not None:
     shape = core.canonicalize_shape(shape)
   key, _ = _check_prng_key("truncated_normal", key)
+  out_sharding = canonicalize_sharding(out_sharding, "truncated_normal")
   dtypes.check_user_dtype_supported(dtype)
   if not dtypes.issubdtype(dtype, np.floating):
     raise ValueError(f"dtype argument to `truncated_normal` must be a float "
                      f"dtype, got {dtype}")
   dtype = dtypes.canonicalize_dtype(dtype)
-  return _truncated_normal(key, lower, upper, shape, dtype)
+  return maybe_auto_axes(_truncated_normal, out_sharding,
+                         shape=shape, dtype=dtype)(key, lower, upper)
 
 @partial(jit, static_argnums=(3, 4))
 def _truncated_normal(key, lower, upper, shape, dtype) -> Array:


Test patch:
diff --git a/tests/pallas/mosaic_gpu_test.py b/tests/pallas/mosaic_gpu_test.py
index 0cfe9197d..d35446359 100644
--- a/tests/pallas/mosaic_gpu_test.py
+++ b/tests/pallas/mosaic_gpu_test.py
@@ -644,6 +644,8 @@ class PallasCallTest(PallasTest):
 
     in_spec = pl.BlockSpec(memory_space=plgpu.GMEM)
     out_spec = plgpu.GPUBlockSpec(
+        (128, 128),
+        lambda: (0, 0),
         transforms=(
             plgpu.TilingTransform((8, 32)),
             plgpu.SwizzleTransform(128),
@@ -674,7 +676,9 @@ class PallasCallTest(PallasTest):
       pl.run_scoped(body, plgpu.SMEM((128, 128), jnp.float32, transforms=ts))
 
     in_spec = pl.BlockSpec(memory_space=plgpu.GMEM)
-    out_spec = plgpu.GPUBlockSpec(transforms=ts, memory_space=plgpu.SMEM)
+    out_spec = plgpu.GPUBlockSpec(
+        (128, 128), lambda: (0, 0), transforms=ts, memory_space=plgpu.SMEM,
+    )
     f = self.pallas_call(
         kernel,
         out_shape=jax.ShapeDtypeStruct([128, 128], jnp.float32),
@@ -715,6 +719,8 @@ class PallasCallTest(PallasTest):
 
     in_spec = pl.BlockSpec(memory_space=plgpu.GMEM)
     out_spec = plgpu.GPUBlockSpec(
+        (2, 128, 128),
+        lambda: (0, 0, 0),
         transforms=(
             plgpu.TilingTransform((8, 32)),
             plgpu.TransposeTransform((0, 2, 1, 3, 4)),
@@ -744,7 +750,11 @@ class PallasCallTest(PallasTest):
         self.pallas_call,
         out_shape=jax.ShapeDtypeStruct([2, 128], jnp.float32),
         in_specs=[pl.BlockSpec(memory_space=src_memory_space)],
-        out_specs=plgpu.GPUBlockSpec(memory_space=plgpu.SMEM),
+        out_specs=plgpu.GPUBlockSpec(
+            (2, 128),
+            lambda: (0, 0),
+            memory_space=plgpu.SMEM,
+        ),
     )
     def kernel(x_ref, o_ref):
       for i in range(2):
@@ -766,7 +776,11 @@ class PallasCallTest(PallasTest):
         self.pallas_call,
         out_shape=jax.ShapeDtypeStruct([2, m], jnp.float32),
         in_specs=[pl.BlockSpec(memory_space=src_memory_space)],
-        out_specs=plgpu.GPUBlockSpec(memory_space=plgpu.SMEM),
+        out_specs=plgpu.GPUBlockSpec(
+            (2, m),
+            lambda: (0, 0),
+            memory_space=plgpu.SMEM,
+        ),
     )
     def kernel(x_ref, o_ref):
       for i in range(2):
@@ -805,19 +819,24 @@ class PallasCallTest(PallasTest):
 
       out = pl.run_scoped(compute, plgpu.ACC((m, n), jnp.float32))
       o_ref[...] = out
+
+    out_spec = plgpu.GPUBlockSpec(
+        (m, n), lambda: (0, 0), memory_space=plgpu.SMEM,
+    )
     f = self.pallas_call(
         kernel,
         out_shape=jax.ShapeDtypeStruct([m, n], jnp.float32),
         in_specs=(
             pl.BlockSpec(memory_space=src_memory_space),
             plgpu.GPUBlockSpec(
+                (k, n),
+                lambda: (0, 0),
                 transforms=(
-                    plgpu.TilingTransform((8, 64)),
-                    plgpu.SwizzleTransform(128),
+                    plgpu.TilingTransform((8, 64)), plgpu.SwizzleTransform(128),
                 ),
             ),
         ),
-        out_specs=plgpu.GPUBlockSpec(memory_space=plgpu.SMEM),
+        out_specs=out_spec,
     )
 
     out_ref = (
@@ -836,7 +855,9 @@ class PallasCallTest(PallasTest):
         plgpu.barrier_wait(barrier_ref)
 
     in_spec = pl.BlockSpec(memory_space=plgpu.GMEM)
-    out_spec = plgpu.GPUBlockSpec(memory_space=plgpu.SMEM)
+    out_spec = plgpu.GPUBlockSpec(
+        (2, 64, 2, 128), lambda: (0, 0, 0, 0), memory_space=plgpu.SMEM,
+    )
     f = self.pallas_call(
         kernel,
         out_shape=jax.ShapeDtypeStruct([2, 64, 2, 128], jnp.float32),
@@ -939,10 +960,11 @@ class PallasCallTest(PallasTest):
         out_shape=jax.ShapeDtypeStruct(shape, jnp.float32),
         in_specs=[
             plgpu.GPUBlockSpec(
+                shape,
+                lambda: (0, 0),
                 transforms=(
-                    plgpu.TilingTransform((8, 32)),
-                    plgpu.SwizzleTransform(128),
-                )
+                    plgpu.TilingTransform((8, 32)), plgpu.SwizzleTransform(128),
+                ),
             )
         ],
     )
@@ -1121,8 +1143,7 @@ class PallasCallTest(PallasTest):
         (128, 64),
         lambda *i: i,
         transforms=(
-            plgpu.TilingTransform((8, 64)),
-            plgpu.SwizzleTransform(128),
+            plgpu.TilingTransform((8, 64)), plgpu.SwizzleTransform(128),
         ),
     )
     @functools.partial(
@@ -1306,7 +1327,9 @@ class PallasCallTest(PallasTest):
 
     shape = (256, 128)
     block_spec = plgpu.GPUBlockSpec(
-        transforms=(plgpu.TilingTransform((8, 64)), plgpu.SwizzleTransform(128))
+        transforms=(
+            plgpu.TilingTransform((8, 64)), plgpu.SwizzleTransform(128),
+        )
     )
     @functools.partial(
         self.pallas_call,
@@ -1357,7 +1380,12 @@ class PallasCallTest(PallasTest):
 
     x = jnp.arange(128 * 128).astype(jnp.float16).reshape(128, 128)
     spec = plgpu.GPUBlockSpec(
-        transforms=(plgpu.TilingTransform((8, 64)), plgpu.SwizzleTransform(128))
+        (128, 128),
+        lambda: (0, 0),
+        transforms=(
+            plgpu.TilingTransform((8, 64)),
+            plgpu.SwizzleTransform(128),
+        ),
     )
     f = self.pallas_call(rotate, out_shape=x, in_specs=[spec], out_specs=spec)
     expected = np.empty_like(x)
@@ -1532,9 +1560,11 @@ class PallasCallSm90ATest(PallasSm90ATest):
     transforms = (plgpu.TilingTransform((8, 64)), plgpu.SwizzleTransform(128))
     @functools.partial(
         self.pallas_call,
-        in_specs=[plgpu.GPUBlockSpec((64, 64), transforms=transforms)],
+        in_specs=[
+            plgpu.GPUBlockSpec((64, 64), lambda: (0, 0), transforms=transforms)
+        ],
         out_shape=jax.ShapeDtypeStruct((64, 64), jnp.float16),
-        out_specs=plgpu.GPUBlockSpec((64, 64)),
+        out_specs=plgpu.GPUBlockSpec((64, 64), lambda: (0, 0)),
     )
     def kernel(i_ref, o_ref):
       def scope(acc_ref):
@@ -1583,28 +1613,25 @@ class PallasCallSm90ATest(PallasSm90ATest):
 
     if self.THREAD_SEMANTICS == plgpu.ThreadSemantics.Lane:
       lhs_spec = plgpu.GPUBlockSpec(
-          lhs_spec.block_shape,
-          lhs_spec.index_map,
+          lhs_spec.block_shape, lhs_spec.index_map,
           transforms=(
               plgpu.TilingTransform((8, elems_128b)),
               plgpu.SwizzleTransform(128),
-          ),
+          )
       )
       rhs_spec = plgpu.GPUBlockSpec(
-          rhs_spec.block_shape,
-          rhs_spec.index_map,
+          rhs_spec.block_shape, rhs_spec.index_map,
           transforms=(
               plgpu.TilingTransform((8, elems_128b)),
               plgpu.SwizzleTransform(128),
-          ),
+          )
       )
       out_spec = plgpu.GPUBlockSpec(
-          out_spec.block_shape,
-          out_spec.index_map,
+          out_spec.block_shape, out_spec.index_map,
           transforms=(
               plgpu.TilingTransform((8, elems_128b)),
               plgpu.SwizzleTransform(128),
-          ),
+          )
       )
 
     res = self.pallas_call(
@@ -1690,9 +1717,14 @@ class PallasCallSm90ATest(PallasSm90ATest):
     res = self.pallas_call(
         kernel,
         in_specs=[
-            plgpu.GPUBlockSpec(transforms=transforms),
-            plgpu.GPUBlockSpec(transforms=transforms),
+            plgpu.GPUBlockSpec(
+                (64, 128), lambda: (0, 0), transforms=transforms
+            ),
+            plgpu.GPUBlockSpec(
+                (128, 192), lambda: (0, 0), transforms=transforms
+            ),
         ],
+        out_specs=plgpu.GPUBlockSpec((64, 192), lambda: (0, 0)),
         out_shape=jax.ShapeDtypeStruct((64, 192), jnp.float32),
     )(a, b)
     np.testing.assert_allclose(res, a @ b, rtol=1e-3)
@@ -1715,10 +1747,17 @@ class PallasCallSm90ATest(PallasSm90ATest):
     res = self.pallas_call(
         kernel,
         in_specs=[
-            plgpu.GPUBlockSpec(transforms=transforms),
-            plgpu.GPUBlockSpec(transforms=transforms),
-            plgpu.GPUBlockSpec(transforms=transforms),
+            plgpu.GPUBlockSpec(
+                (64, 128), lambda: (0, 0), transforms=transforms
+            ),
+            plgpu.GPUBlockSpec(
+                (128, 192), lambda: (0, 0), transforms=transforms
+            ),
+            plgpu.GPUBlockSpec(
+                (64, 192), lambda: (0, 0), transforms=transforms
+            ),
         ],
+        out_specs=plgpu.GPUBlockSpec((64, 192), lambda: (0, 0)),
         out_shape=jax.ShapeDtypeStruct((64, 192), jnp.float16),
     )(a, b, i)
     np.testing.assert_allclose(res, i + a @ b, rtol=2e-3)
@@ -1744,9 +1783,14 @@ class PallasCallSm90ATest(PallasSm90ATest):
     res = self.pallas_call(
         kernel,
         in_specs=[
-            plgpu.GPUBlockSpec(transforms=transforms),
-            plgpu.GPUBlockSpec(transforms=transforms),
+            plgpu.GPUBlockSpec(
+                (2, 64, 128), lambda: (0, 0, 0), transforms=transforms
+            ),
+            plgpu.GPUBlockSpec(
+                (2, 128, 192), lambda: (0, 0, 0), transforms=transforms
+            ),
         ],
+        out_specs=plgpu.GPUBlockSpec((64, 192), lambda: (0, 0)),
         out_shape=jax.ShapeDtypeStruct((64, 192), jnp.float32),
     )(a, b)
     np.testing.assert_allclose(res, a[0] @ b[0], rtol=1e-3)
diff --git a/tests/pallas/tpu_pallas_test.py b/tests/pallas/tpu_pallas_test.py
index 2e773b88f..128fe5068 100644
--- a/tests/pallas/tpu_pallas_test.py
+++ b/tests/pallas/tpu_pallas_test.py
@@ -1135,39 +1135,6 @@ class PallasCallDMATest(PallasBaseTest):
     np.testing.assert_array_equal(y, x)
     np.testing.assert_array_equal(sem_val, 0)
 
-  def test_set_dma_priority(self):
-    if not jtu.if_cloud_tpu_at_least(2025, 4, 5):
-      self.skipTest('Needs a newer libTPU')
-    if jtu.get_tpu_version() < 5:
-      self.skipTest('Target does not support DMA prefetch between HBM and VMEM')
-    def kernel(x1, x2, y1, y2, scratch1, scratch2, sem1, sem2):
-      copy1 = pltpu.async_copy(x1, scratch1, sem1, priority=1)
-      copy2 = pltpu.async_copy(x2, scratch2, sem2, priority=0)
-      copy1.wait()
-      copy2.wait()
-      copy1 = pltpu.async_copy(scratch1, y1, sem1, priority=0)
-      copy2 = pltpu.async_copy(scratch2, y2, sem2, priority=1)
-      copy1.wait()
-      copy2.wait()
-
-    shape = (8, 128)
-    dtype = jnp.int32
-    x1 = jnp.arange(np.prod(shape), dtype=dtype).reshape(shape)
-    x2 = x1 + 1
-    y1, y2 = self.pallas_call(
-        kernel,
-        grid_spec=pltpu.PrefetchScalarGridSpec(
-            num_scalar_prefetch=0,
-            in_specs=[pl.BlockSpec(memory_space=pl.ANY)] * 2,
-            scratch_shapes=[pltpu.VMEM(shape, dtype)] * 2
-            + [pltpu.SemaphoreType.DMA] * 2,
-            out_specs=[pl.BlockSpec(memory_space=pl.ANY)] * 2,
-        ),
-        out_shape=[jax.ShapeDtypeStruct(shape, dtype)] * 2,
-    )(x1, x2)
-    np.testing.assert_array_equal(y1, x1)
-    np.testing.assert_array_equal(y2, x2)
-
   def test_hbm_hbm_dma(self):
     def kernel(x_hbm_ref, y_hbm_ref):
       def body(sem):
@@ -2698,19 +2665,19 @@ class PrettyPrintingTest(PallasBaseTest):
   @parameterized.parameters(
       (
           lambda i: (i, pl.ds(0, 8), pl.ds(0, 128)),
-          'dma_start(p0) c[d,:,:] -> e[...] f',
+          'dma_start c[d,:,:] -> e[...] f',
       ),
       (
           lambda i: (0, pl.ds(i, 8), pl.ds(0, 128)),
-          'dma_start(p0) c[0,d:d+8,:] -> e[...] f',
+          'dma_start c[0,d:d+8,:] -> e[...] f',
       ),
       (
           lambda i: (i, pl.ds(2, 4), pl.ds(0, 100)),
-          'dma_start(p0) c[d,2:6,:100] -> e[...] f',
+          'dma_start c[d,2:6,:100] -> e[...] f',
       ),
       (
           lambda i: (i, pl.ds(2, 6), pl.ds(4, 100)),
-          'dma_start(p0) c[d,2:,4:104] -> e[...] f',
+          'dma_start c[d,2:,4:104] -> e[...] f',
       ),
   )
   def test_dma_custom_pretty_print(self, indexer, expected):
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index 580cfcd7a..f2db913af 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -7423,6 +7423,26 @@ class ShardingInTypesTest(jtu.JaxTestCase):
     out = f(jax.device_put(np.arange(8 * 12.).reshape(8, 12), s), key)
     self.assertEqual(out.sharding, NamedSharding(mesh, P('x', 'y')))
 
+  @jtu.with_user_mesh((2, 2), ('x', 'y'))
+  def test_random_truncated_normal(self, mesh):
+    @jax.jit
+    def f(key, lower):
+      out = jax.random.truncated_normal(key, lower, 2., 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, -1.)
+    self.assertEqual(out.sharding, NamedSharding(mesh, P('x', 'y')))
+
+    lowered_text = f.lower(key, -1.).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)
+
   def test_auto_axes_no_context_mesh(self):
     mesh = jtu.create_mesh((2, 2), ('x', 'y'), axis_types=(AxisType.Explicit,) * 2)
     np_inp = np.arange(16.).reshape(8, 2)