Repository: jax-ml/jax
Pull request #28629 (2025-05-08 22:34:53)

Simplify add's unreduced rule. Only propagate unreduced if both lhs and rhs are unreduced.
Simplify add's unreduced rule. Only propagate unreduced if both lhs and rhs are unreduced.



Patch:
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
index b41e78899..74297fc57 100644
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -4574,16 +4574,24 @@ def _add_transpose(t, x, y):
   else:
     return [_unbroadcast(x_aval, t), _unbroadcast(y_aval, t)]
 
-def _add_unreduced(out_sharding, *avals):
-  unreduced = [a.sharding.spec.unreduced for a in avals if a.shape]
-  # TODO(yashkatariya): Relax this restriction to allow
-  # `f32[8]{R:x} + f32[8]{U:x} -> f32[8]{U:x}` for example and maybe more cases.
-  if unreduced:
-    if not all(unreduced[0] == u for u in unreduced[1:]):
+def _add_unreduced(out_sharding, x, y):
+  x_ur, y_ur = x.sharding.spec.unreduced, y.sharding.spec.unreduced
+  if x_ur and y_ur:
+    if x_ur != y_ur:
       raise core.ShardingTypeError(
-          'All arrays must be unreduced along the same mesh axes. Got'
-          f' {", ".join(map(str, map(tuple, unreduced)))}')
-    res_unreduced = unreduced[0]
+          'lhs and rhs to `add` must be unreduced along the same mesh axes. '
+          f'Got lhs={x_ur}, rhs={y_ur}')
+    res_unreduced = x_ur
+  elif x_ur or y_ur:
+    if x_ur and not y_ur:
+      lhs_str, rhs_str = 'lhs', 'rhs'
+    else:
+      assert not x_ur and y_ur
+      lhs_str, rhs_str = 'rhs', 'lhs'
+    raise core.ShardingTypeError(
+        f'{lhs_str} is unreduced while {rhs_str} is not. `add` operation does'
+        ' not allow this because there will be implicit communication. Please'
+        f' reduce {lhs_str} via `reshard` before calling `add`.')
   else:
     res_unreduced = None
   return out_sharding.with_spec(out_sharding.spec.with_unreduced(res_unreduced))


Test patch:
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index 0c43df76c..e4b91aa17 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -7763,22 +7763,34 @@ class ShardingInTypesTest(jtu.JaxTestCase):
         "unreduced axes should be equal to the contracting specs"):
       h.trace(x, y)
 
-  @jtu.with_explicit_mesh((2, 2), ('x', 'y'))
+  @jtu.with_explicit_mesh((2, 2, 1), ('x', 'y', 'z'))
   def test_add_unreduced_error(self, mesh):
     np_inp = np.arange(16).reshape(8, 2)
     x = jax.device_put(np_inp, P('x', 'y'))
     y = jax.device_put(np_inp.T, P('y', None))
+    a = jax.device_put(np_inp, P('x', 'z'))
+    b = jax.device_put(np_inp.T, P('z', None))
 
     @jax.jit
-    def f(x, y):
+    def f(x, y, a, b):
       m1 = jnp.einsum('xy,yz->xz', x, y, out_sharding=P('x', unreduced='y'))
-      m2 = jnp.einsum('xy,yz->xz', x, y, out_sharding=P('x'))
+      m2 = jnp.einsum('xy,yz->xz', a, b, out_sharding=P('x', unreduced='z'))
       return m1 + m2
 
     with self.assertRaisesRegex(
         core.ShardingTypeError,
-        "arrays must be unreduced along the same mesh axes"):
-      f.trace(x, y)
+        "lhs and rhs to `add` must be unreduced along the same mesh axes"):
+      f.trace(x, y, a, b)
+
+    @jax.jit
+    def g(x, y):
+      m1 = jnp.einsum('xy,yz->xz', x, y, out_sharding=P('x', unreduced='y'))
+      m2 = jnp.einsum('xy,yz->xz', a, b, out_sharding=P('x'))
+      return m1 + m2
+
+    with self.assertRaisesRegex(
+        core.ShardingTypeError, "lhs is unreduced while rhs is not"):
+      g.trace(x, y)
 
 
 @jtu.pytest_mark_if_available('multiaccelerator')