Repository: jax-ml/jax
Pull request #26981 (2025-03-07 05:15:34)

Add a divisibility check so that we make sure that sharding evenly divides the shape (until this restriction is lifted) to make sure we don't create bad shardings.
Add a divisibility check so that we make sure that sharding evenly divides the shape (until this restriction is lifted) to make sure we don't create bad shardings.

Also improve dynamic_update_slice sharding error by printing `aval.str_short()` instead of full sharding because it's concise and gives more info than the current error (i.e. it adds shape too to the error message)

Also make some formatting changes in scan lowering to make it easier to debug.



Patch:
diff --git a/jax/_src/core.py b/jax/_src/core.py
index 17b842359..767c61089 100644
--- a/jax/_src/core.py
+++ b/jax/_src/core.py
@@ -1852,9 +1852,22 @@ def _maybe_modify_sharding(sharding, ndim):
     out = _make_lengths_same(out, ndim)
   return out
 
+def _check_divisibility(sharding, shape):
+  mesh = sharding.mesh
+  for dim, (spec, sh) in enumerate(zip(sharding.spec, shape)):
+    if spec is None:
+      continue
+    spec = spec if isinstance(spec, tuple) else (spec,)
+    size = math.prod(mesh.shape[s] for s in spec)
+    _, remainder = divmod(sh, size)
+    if remainder != 0:
+      raise ValueError(
+          f"Sharding spec {spec} implies that array axis {dim} is partitioned"
+          f" {size} times, but does not evenly divide the dimension size {sh}."
+          f" Got shape: {shape} and sharding {sharding}")
 
 @cache(max_size=4096, trace_context_in_key=True)
-def get_sharding(sharding, ndim):
+def get_sharding(sharding, shape):
   """Modifies and checks the sharding.
 
   Some modifications/checks include:
@@ -1863,6 +1876,7 @@ def get_sharding(sharding, ndim):
     * Checking for len(spec)-ndim match
     * Checking if the mesh is an AbstractMesh.
   """
+  ndim = len(shape)
   if sharding is None:
     return NamedSharding(mesh_lib.empty_abstract_mesh, P(*[None] * ndim))
 
@@ -1874,6 +1888,7 @@ def get_sharding(sharding, ndim):
   if not isinstance(out_s.mesh, mesh_lib.AbstractMesh):
     raise ValueError("Mesh of an aval must be an AbstractMesh. "
                      f"Got {out_s.mesh} of type {type(out_s.mesh)}")
+  _check_divisibility(out_s, shape)
   return out_s
 
 
@@ -1885,7 +1900,7 @@ class ShapedArray(UnshapedArray):
     self.shape = canonicalize_shape(shape)
     self.dtype = _dtype_object(dtype)
     self.weak_type = weak_type
-    self.sharding = get_sharding(sharding, len(self.shape))
+    self.sharding = get_sharding(sharding, self.shape)
 
   def update(self, shape=None, dtype=None, weak_type=None, **kwargs):
     if shape is None:
diff --git a/jax/_src/interpreters/mlir.py b/jax/_src/interpreters/mlir.py
index 7c10c7b8d..e5dabd146 100644
--- a/jax/_src/interpreters/mlir.py
+++ b/jax/_src/interpreters/mlir.py
@@ -939,7 +939,7 @@ def sharded_aval(aval: core.AbstractValue,
     return aval
   if not isinstance(aval, (core.ShapedArray, core.DShapedArray)):
     raise NotImplementedError
-  return aval.update(sharding.shard_shape(aval.shape))  # type: ignore
+  return aval.update(sharding.shard_shape(aval.shape), sharding=None)  # type: ignore
 
 
 def eval_dynamic_shape(ctx: LoweringRuleContext,
diff --git a/jax/_src/lax/control_flow/loops.py b/jax/_src/lax/control_flow/loops.py
index 4c0bc6e6f..caabace86 100644
--- a/jax/_src/lax/control_flow/loops.py
+++ b/jax/_src/lax/control_flow/loops.py
@@ -443,6 +443,7 @@ def _scan_impl(*args, reverse, length, num_consts, num_carry, jaxpr, linear,
   consts, carry, xs_ = split_list(args, [num_consts, num_carry])
   _, y_avals = split_list(jaxpr.out_avals, [num_carry])
   num_trips, remainder = divmod(length, unroll)
+
   if unroll != 1 and num_trips == 1 and remainder == 0:
     # In that case, we explicitly want to fully unroll the loop. Put everything
     # into the remainder block and avoid lowering to a while loop.
@@ -459,26 +460,6 @@ def _scan_impl(*args, reverse, length, num_consts, num_carry, jaxpr, linear,
     xss = [lax.reshape(x, (num_trips, unroll, *x.shape[1:])) for x in xs_]
     yss = _map(partial(_empty_array, (num_trips, unroll), None), y_avals)
 
-  def cond_fun(while_carry):
-    i, _, _ = while_carry
-    return i < num_trips
-  def body_fun(while_carry):
-    i_, carry, yss = while_carry
-    i = num_trips - i_ - 1 if reverse else i_
-    xs = [
-        slicing.dynamic_index_in_dim(
-            xs, i, keepdims=False, allow_negative_indices=False
-        )
-        for xs in xss
-    ]
-    carry, ys = inner(unroll, carry, xs)
-    yss = [
-        slicing.dynamic_update_index_in_dim(
-            ys, upd, i, 0, allow_negative_indices=False
-        )
-        for ys, upd in zip(yss, ys)
-    ]
-    return i_ + 1, carry, yss
   def inner(n, carry, xs):
     ys = []
     if unroll == 1:
@@ -493,6 +474,22 @@ def _scan_impl(*args, reverse, length, num_consts, num_carry, jaxpr, linear,
     ys = list(reversed(ys)) if reverse else ys
     return carry, _map(_stack, zip(*ys))
 
+  def body_fun(while_carry):
+    i_, carry, yss = while_carry
+    i = num_trips - i_ - 1 if reverse else i_
+    xs = [slicing.dynamic_index_in_dim(xs, i, keepdims=False,
+                                       allow_negative_indices=False)
+          for xs in xss]
+    carry, ys = inner(unroll, carry, xs)
+    yss = [slicing.dynamic_update_index_in_dim(y, upd, i, 0,
+                                               allow_negative_indices=False)
+           for y, upd in zip(yss, ys)]
+    return i_ + 1, carry, yss
+
+  def cond_fun(while_carry):
+    i, _, _ = while_carry
+    return i < num_trips
+
   if num_trips:
     i = lax._const(num_trips, 0)
     _, carry, yss = while_loop(cond_fun, body_fun, (i, carry, yss))
diff --git a/jax/_src/lax/slicing.py b/jax/_src/lax/slicing.py
index bb1647d7e..c26de99c7 100644
--- a/jax/_src/lax/slicing.py
+++ b/jax/_src/lax/slicing.py
@@ -1608,8 +1608,9 @@ def _dynamic_update_slice_sharding_rule(operand, update, *start_indices):
   if operand.sharding != update.sharding:
     raise TypeError(
         "dynamic_update_slice update sharding must be equal to operand"
-        f" sharding, got update sharding {update.sharding} for operand sharding"
-        f" {operand.sharding}.")
+        " sharding, got update sharding"
+        f" {update.str_short(mesh_axis_types=True)} for operand sharding"
+        f" {operand.str_short(mesh_axis_types=True)}.")
   return operand.sharding
 
 def _dynamic_update_slice_dtype_rule(operand, update, *start_indices):


Test patch:
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index 4b267196e..b0a59027f 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -7072,6 +7072,14 @@ class ShardingInTypesTest(jtu.JaxTestCase):
     with self.assertRaisesRegex(ValueError, "Context mesh.*cannot be empty"):
       auto_axes(f, out_shardings=s)(arr)
 
+  def test_divisbility_aval_error(self):
+    abstract_mesh = mesh_lib.AbstractMesh(
+        (('x', 2),), axis_types={AxisTypes.Explicit: 'x'})
+    s = NamedSharding(abstract_mesh, P('x'))
+    with self.assertRaisesRegex(
+        ValueError, 'does not evenly divide the dimension size'):
+      core.ShapedArray((5, 2), jnp.int32, sharding=s)
+
 
 @jtu.pytest_mark_if_available('multiaccelerator')
 class PJitErrorTest(jtu.JaxTestCase):