Repository: jax-ml/jax
Pull request #28597 (2025-05-07 23:06:09)

Add basic support of `unreduced` to sharding-in-types! We cannot lower it right now, but it atleast shows up in types.
Add basic support of `unreduced` to sharding-in-types! We cannot lower it right now, but it atleast shows up in types.

The API to specify unreduced is via `PartitionSpec`. For example: `PartitionSpec('x', 'y', None, unreduced='z')` or `PartitionSpec('x', unreduced=('y', 'z'))`.

In types/jaxpr, unreduced will show up as: `f32[8@x,2]{U:y}`

But we support unreduced only in dot_general and nary ops (add, mul, etc) as of this change: (the support will be expanded in following changes)

* **dot general** only allows unreduced when contracting dims are sharded. And the unreduced axes specified by the user needs to match the sharding of the contracting dims. In all other cases, an error is raised. An example of how unreduced can be specified: `jnp.einsum('xy,yz->xz', x, y, out_sharding=P('x', unreduced='y'))`

* **nary ops** can propagate unreduced (add, mul, etc). If all ops aren't unreduced across the same mesh axes, an error is raised.



Patch:
diff --git a/jax/_src/core.py b/jax/_src/core.py
index 780d2c487..da9efdb0e 100644
--- a/jax/_src/core.py
+++ b/jax/_src/core.py
@@ -1825,11 +1825,12 @@ def get_cur_mesh_sharding(spec=None):
   return NamedSharding(mesh_lib.get_abstract_mesh(), spec)
 
 def _make_lengths_same(sharding, ndim):
-  if ndim > len(sharding.spec):
-    return sharding.with_spec(sharding.spec._normalized_spec_for_aval(ndim))
-  if ndim < len(sharding.spec):
-    assert all(s is None for s in sharding.spec[ndim:]), (ndim, sharding.spec)
-    return sharding.with_spec(sharding.spec[:ndim])
+  pspec = sharding.spec
+  if ndim > len(pspec):
+    return sharding.with_spec(pspec._normalized_spec_for_aval(ndim))
+  if ndim < len(pspec):
+    assert all(s is None for s in pspec[ndim:]), (ndim, pspec)
+    return sharding.with_spec(P(*pspec[:ndim], unreduced=pspec.unreduced))
   assert False, "unreachable"
 
 # TODO(yashkatariya): Only works with User/Auto. Generalize it to work with
@@ -1841,11 +1842,11 @@ def modify_spec_for_auto_manual(spec, mesh) -> P:
       new_spec.append(s)
     else:
       temp_s = s[0] if isinstance(s, tuple) else s
-      new_spec.append(
-          None
-          if mesh._name_to_type[temp_s] in (AxisType.Auto, AxisType.Manual)
-          else s)
-  return P(*new_spec)
+      new_spec.append(s if mesh._name_to_type[temp_s] == AxisType.Explicit
+                      else None)
+  new_unreduced = tuple(u for u in spec.unreduced
+                        if mesh._name_to_type[u] == AxisType.Explicit)
+  return P(*new_spec, unreduced=new_unreduced)
 
 def _maybe_modify_sharding(sharding, ndim):
   if len(sharding.spec) == 0 or all(s is None for s in sharding.spec):
@@ -1905,8 +1906,8 @@ def str_short_aval(shape, dtype, mesh, spec, vma,
   dt_str = dt_str.replace('void', 'float0')
   shapestr = _get_shape_sharding_str(shape, spec)
   mesh_axes = f'({mesh._axis_types_dict})' if mesh_axis_types else ''
-  vma = f"{{{','.join(i for i in vma)}}}" if vma else ''
-  return f'{dt_str}[{shapestr}]{vma}{mesh_axes}'
+  vma_ur = _vma_ur_str(vma, spec.unreduced)
+  return f'{dt_str}[{shapestr}]{vma_ur}{mesh_axes}'
 
 def get_vma(vma, mesh):
   if mesh.empty:
@@ -2000,6 +2001,18 @@ def _get_shape_sharding_str(shape, spec):
       out.append(f"{s1}@{s2}")
   return ','.join(out)
 
+def _create_str(x, prefix):
+  x_str = f"{','.join(i for i in x)}"
+  x_str = x_str if len(x) == 1 else f"({x_str})"
+  return f"{prefix}:{x_str}"
+
+def _vma_ur_str(vma, unreduced):
+  if not vma and not unreduced:
+    return ''
+  vma_str = _create_str(vma, 'V') if vma else ''
+  ur_str = _create_str(unreduced, 'U') if unreduced else ''
+  sep = ', ' if vma and unreduced else ''
+  return f"{{{vma_str}{sep}{ur_str}}}"
 
 def primal_dtype_to_tangent_dtype(primal_dtype):
   if isinstance(primal_dtype, dtypes.ExtendedDType):
diff --git a/jax/_src/interpreters/batching.py b/jax/_src/interpreters/batching.py
index 1c6e00861..0fbe54a30 100644
--- a/jax/_src/interpreters/batching.py
+++ b/jax/_src/interpreters/batching.py
@@ -458,6 +458,8 @@ class AxisData:
 
 def get_sharding_for_vmap(axis_data, orig_sharding, axis):
   val = axis_data.explicit_mesh_axis
+  # TODO(yashkatariya): Preserve unreduced here using
+  # `orig_sharding.spec.with_partitions`
   new_spec = P(*tuple_insert(orig_sharding.spec, axis, val))
   return NamedSharding(orig_sharding.mesh, new_spec)
 
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
index 6c2485651..f6b4c1be1 100644
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -3954,7 +3954,7 @@ def broadcasting_sharding_rule(name, *avals):
 
   result_specs = [None] * len(shapes[0])
   for i, (ss, ds) in enumerate(zip(zip(*specs), zip(*shapes))):
-    if all(s == ss[0] for s in ss[1:]):
+    if all(ss[0] == s for s in ss[1:]):
       # if all dimension shardings are same, the resulting dimension sharding is
       # the same.
       result_specs[i] = ss[0]
@@ -3974,7 +3974,20 @@ def broadcasting_sharding_rule(name, *avals):
             raise core.ShardingTypeError(
                 f'{name} got incompatible shardings for broadcasting: '
                 f'{", ".join(map(str, map(tuple, specs)))}.')
-  return NamedSharding(mesh, P(*result_specs))
+
+  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:]):
+      raise core.ShardingTypeError(
+          'All arrays must be unreduced along the same mesh axes. Got'
+          f' {", ".join(map(str, map(tuple, unreduced)))}')
+    result_unreduced = unreduced[0]
+  else:
+    result_unreduced = None
+
+  return NamedSharding(mesh, P(*result_specs, unreduced=result_unreduced))
 
 def naryop(result_dtype, accepted_dtypes, name, allow_extended_dtype=False,
            require_same_dtypes=True):
@@ -5190,11 +5203,26 @@ def _dot_general_sharding_rule(lhs, rhs, *, dimension_numbers, precision,
         'Mesh of both lhs and rhs should match. Got lhs:'
         f' {lhs.sharding.mesh} and rhs: {rhs.sharding.mesh}')
 
+  (lhs_contracting, rhs_contracting), (lhs_batch, rhs_batch) = dimension_numbers
+  lhs_contracting_spec = tuple(lhs.sharding.spec[i] for i in lhs_contracting)
+  rhs_contracting_spec = tuple(rhs.sharding.spec[i] for i in rhs_contracting)
+
   if out_sharding is not None:
     assert isinstance(out_sharding, NamedSharding)
+    if out_sharding.spec.unreduced:
+      if lhs_contracting_spec != rhs_contracting_spec:
+        raise core.ShardingTypeError(
+            'lhs and rhs contracting dims should be sharded identically when'
+            ' out_sharding provided to dot_general mentions unreduced_axes.'
+            f' Got {out_sharding=}, {lhs_contracting_spec=},'
+            f' {rhs_contracting_spec=}')
+      if out_sharding.spec.unreduced != lhs_contracting_spec:
+        raise core.ShardingTypeError(
+            "out_sharding's unreduced axes should be equal to the contracting"
+            f' specs. Got unreduced axes={out_sharding.spec.unreduced} and'
+            f' contracting spec={lhs_contracting_spec}')
     return out_sharding
 
-  (lhs_contracting, rhs_contracting), (lhs_batch, rhs_batch) = dimension_numbers
   lhs_batch_spec = tuple(lhs.sharding.spec[i] for i in lhs_batch)
   rhs_batch_spec = tuple(rhs.sharding.spec[i] for i in rhs_batch)
   msg = ("dot_general requires lhs batch dimensions and rhs batch dimensions "
@@ -5202,8 +5230,6 @@ def _dot_general_sharding_rule(lhs, rhs, *, dimension_numbers, precision,
         f"{rhs_batch_spec}.")
   _check_specs_match(lhs_batch_spec, rhs_batch_spec, msg)
 
-  lhs_contracting_spec = tuple(lhs.sharding.spec[i] for i in lhs_contracting)
-  rhs_contracting_spec = tuple(rhs.sharding.spec[i] for i in rhs_contracting)
   msg = ("dot_general requires contracting dimensions to have consistent "
         f"sharding, got {lhs_contracting_spec} and {rhs_contracting_spec}.")
   _check_specs_match(lhs_contracting_spec, rhs_contracting_spec, msg)
diff --git a/jax/_src/named_sharding.py b/jax/_src/named_sharding.py
index d0a49b6b0..3dfcbd29f 100644
--- a/jax/_src/named_sharding.py
+++ b/jax/_src/named_sharding.py
@@ -25,6 +25,7 @@ from jax._src.util import use_cpp_class, cache, use_cpp_method
 from jax._src.lib import xla_client as xc
 from jax._src.lib.mlir.dialects import sdy
 from jax._src import mesh as mesh_lib
+from jax._src.mesh import AxisType
 from jax._src.partition_spec import PartitionSpec
 from jax._src import sharding as JSharding
 import numpy as np
@@ -410,6 +411,7 @@ def array_mapping_to_axis_resources(array_mapping: ArrayMapping):
 def check_pspec(mesh, spec, _manual_axes=frozenset()):
   _check_unique_resources(spec, "NamedSharding spec", mesh)
   _check_mesh_resource_axis(mesh, spec)
+  _check_mesh_unreduced(mesh, spec)
 
 class DuplicateSpecError(Exception):
   def __init__(self, message, mesh, pspec):
@@ -443,14 +445,13 @@ def _check_unique_resources(pspec: PartitionSpec, arg_name: str, mesh=None
             f' for {mesh_lib.show_axes(multiple_uses)}'),
         mesh=mesh, pspec=pspec)
 
-
 def _check_mesh_resource_axis(mesh, pspec):
   for p in pspec:
     if p is PartitionSpec.UNCONSTRAINED or p is None:
       continue
     p = p if isinstance(p, tuple) else (p,)
     for r in p:
-      if r not in mesh.shape:
+      if r not in mesh.axis_names:
         raise ValueError(
             f"Resource axis: {r} of {pspec} "
             f"is not found in mesh: {tuple(mesh.shape.keys())}.")
@@ -459,9 +460,34 @@ def _check_mesh_resource_axis(mesh, pspec):
           'AxisTypes should be the same in a tuple subset of PartitionSpec:'
           f' {pspec}. Got subset {p} with axis'
           f' types: ({", ".join(str(mesh._name_to_type[r]) for r in p)})')
-  if (mesh_lib.AxisType.Auto not in mesh._axis_types_dict and
+  if (AxisType.Auto not in mesh._axis_types_dict and
       PartitionSpec.UNCONSTRAINED in pspec):
     raise ValueError(
         f'{pspec} cannot contain'
         ' `P.UNCONSTRAINED` when no mesh axis_types are `Auto`. Got mesh'
         f' axis_types: {mesh._axis_types_dict}')
+
+def _check_mesh_unreduced(mesh, pspec):
+  counts = {}
+  duplicate = False
+  for u in pspec.unreduced:
+    if u not in mesh.axis_names:
+      raise ValueError(
+          f'Unreduced axes {u} is not found in {mesh.axis_names=}. '
+          f'Got {pspec=}')
+    count = counts.get(u, 0)
+    if count > 0:
+      duplicate = True
+    counts[u] = count + 1
+  if duplicate:
+    multiple_uses = [r for r, c in counts.items() if c > 1]
+    raise ValueError(
+        f'Unreduced axes in {pspec} has duplicate entries which is not allowed.'
+        f' Got {mesh_lib.show_axes(multiple_uses)}')
+
+  for u in pspec.unreduced:
+    if mesh._name_to_type[u] in (AxisType.Auto, AxisType.Manual):
+      raise ValueError(
+          'Unreduced axes can only refer to mesh axes that is of type'
+          f' `Explicit`. Got unreduced axes: {pspec.unreduced} and'
+          f' mesh: {mesh}')
diff --git a/jax/_src/numpy/einsum.py b/jax/_src/numpy/einsum.py
index 21333a9e7..3f657082e 100644
--- a/jax/_src/numpy/einsum.py
+++ b/jax/_src/numpy/einsum.py
@@ -25,7 +25,7 @@ from jax._src.api import jit, named_call
 from jax._src.lax import lax
 from jax._src.lax.lax import PrecisionLike
 from jax._src.numpy import util
-from jax._src.sharding_impls import canonicalize_sharding, NamedSharding, PartitionSpec as P
+from jax._src.sharding_impls import canonicalize_sharding, NamedSharding
 from jax._src.typing import Array, ArrayLike, DTypeLike
 from jax._src.util import partition_list, set_module, unzip2
 
@@ -422,7 +422,8 @@ def _einsum(
         " instances. Please file a bug if this is not enough for your use case.")
   dtypes.check_user_dtype_supported(preferred_element_type, "einsum")
   if preferred_element_type is None:
-    preferred_element_type, output_weak_type = dtypes.result_type(*operands, return_weak_type_flag=True)
+    preferred_element_type, output_weak_type = dtypes.result_type(
+        *operands, return_weak_type_flag=True)
   else:
     output_weak_type = False
 
@@ -557,7 +558,7 @@ def _einsum(
           spec = out_sharding.spec
           inverse_spec = tuple(spec[result_names.index(name)] for name in names)
           dot_general_out_sharding = NamedSharding(
-              out_sharding.mesh, P(*inverse_spec))
+              out_sharding.mesh, spec.with_partitions(inverse_spec))
         else:
           dot_general_out_sharding = out_sharding  # type: ignore
         dimension_numbers = ((lhs_cont, rhs_cont), (lhs_batch, rhs_batch))
diff --git a/jax/_src/partition_spec.py b/jax/_src/partition_spec.py
index 158e36148..fcea21934 100644
--- a/jax/_src/partition_spec.py
+++ b/jax/_src/partition_spec.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from __future__ import annotations
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 class UnconstrainedSingleton:
 
@@ -43,6 +43,25 @@ def _canonicalize_partition(partition):
     return tuple(partition)
   return partition
 
+def _check(partitions, unreduced):
+  us = set(unreduced)
+  for p in partitions:
+    p = p if isinstance(p, tuple) else (p,)
+    for r in p:
+      if r in us:
+        raise ValueError(
+            "partitions cannot overlap with unreduced axes passed to"
+            f" PartitionSpec. Got partitions: {partitions} and unreduced axes:"
+            f" {unreduced}")
+  if None in unreduced:
+    raise ValueError(
+        "unreduced cannot contain None. All elements in unreduced should refer"
+        " to the mesh axes.")
+
+def unpicke_pspec(partitions, unreduced):
+  return PartitionSpec(*partitions, unreduced=unreduced)
+
+AxisName = Any
 
 class PartitionSpecImpl:
   """Tuple describing how to partition an array across a mesh of devices.
@@ -53,20 +72,34 @@ class PartitionSpecImpl:
   This class exists so JAX's pytree utilities can distinguish a partition
   specifications from tuples that should be treated as pytrees.
   """
-  __slots__ = ("_partitions",)
+  __slots__ = ("_partitions", "_unreduced")
   __match_args__ = ("_partitions",)
 
   # A sentinel value representing a dim is unconstrained.
   UNCONSTRAINED = _UNCONSTRAINED_PARTITION
 
-  def __init__(self, *partitions):
+  def __init__(self, *partitions,
+               unreduced: tuple[AxisName, ...] | AxisName | None = None):
     self._partitions = tuple(_canonicalize_partition(p) for p in partitions)
+    self._unreduced = (
+        () if unreduced is None else tuple(unreduced)
+        if isinstance(unreduced, (list, tuple)) else (unreduced,))
+    _check(self._partitions, self._unreduced)
+
+  @property
+  def unreduced(self):
+    return self._unreduced
 
   def __repr__(self):
-    return f"PartitionSpec{self._partitions!r}"
+    pr = repr(self._partitions)[1:-1]
+    if not self._unreduced:
+      return f"PartitionSpec({pr})"
+    ur_str = f"unreduced={self._unreduced!r}"
+    pr = '' if not pr else f"{pr} " if pr.endswith(',') else f"{pr}, "
+    return (f"PartitionSpec({pr}{ur_str})")
 
   def __reduce__(self):
-    return (PartitionSpec, self._partitions)
+    return (unpicke_pspec, (self._partitions, self._unreduced))
 
   def __getitem__(self, i):
     return self._partitions[i]
@@ -80,20 +113,42 @@ class PartitionSpecImpl:
   def __eq__(self, other):
     if not isinstance(other, (PartitionSpec, tuple)):
       return False
-    other = tuple(_canonicalize_partition(o) for o in other)
-    return self._partitions == other
+    other_p = tuple(_canonicalize_partition(o) for o in other)
+    if isinstance(other, PartitionSpec):
+      return (self._partitions == other_p and
+              self._unreduced == other._unreduced)
+    else:
+      if self._unreduced:
+        raise TypeError(
+            f"other {other} cannot be of instance `tuple` when self {self} has"
+            " unreduced in `__eq__` of PartitionSpec.")
+      return self._partitions == other_p
 
   def __hash__(self):
-    return hash(self._partitions)
+    return hash((self._partitions, self._unreduced))
 
   def __add__(self, other):
     if not isinstance(other, (tuple, PartitionSpec)):
-      return NotImplemented
-    return PartitionSpec(*self, *other)
+      raise NotImplementedError
+    if isinstance(other, PartitionSpec):
+      return PartitionSpec(
+          *self, *other,
+          unreduced=(*self._unreduced, *other._unreduced))
+    else:
+      if self._unreduced:
+        raise TypeError(
+            f"other {other} cannot be of instance `tuple` when self {self} has"
+            " unreduced in `__add__` of PartitionSpec.")
+      return PartitionSpec(*self, *other)
 
   def __radd__(self, other):
-    if not isinstance(other, (tuple, PartitionSpec)):
-      return NotImplemented
+    if not isinstance(other, tuple):
+      raise NotImplementedError
+    # other will always be a tuple.
+    if self._unreduced:
+      raise TypeError(
+          f"other {other} cannot be of instance `tuple` when self {self} has"
+          " unreduced in `__radd__` of PartitionSpec.")
     return PartitionSpec(*other, *self)
 
   def index(self, value):
@@ -102,12 +157,16 @@ class PartitionSpecImpl:
   def count(self, value):
     return self._partitions.count(_canonicalize_partition(value))
 
+  def with_partitions(self, new_partitions):
+    return PartitionSpec(*new_partitions, unreduced=self._unreduced)
+
   def _normalized_spec_for_aval(self, ndim: int) -> PartitionSpec:
     out = [None if p is _UNCONSTRAINED_PARTITION else p
            for p in self._partitions]
     if len(out) < ndim:
       out.extend([None] * (ndim - len(out)))
-    return PartitionSpec(*out)
+    return self.with_partitions(out)
+
 
 if TYPE_CHECKING:
   class PartitionSpec(PartitionSpecImpl, tuple):  # type: ignore


Test patch:
diff --git a/tests/array_test.py b/tests/array_test.py
index 97d66566c..b951f7f6b 100644
--- a/tests/array_test.py
+++ b/tests/array_test.py
@@ -1403,6 +1403,69 @@ class ShardingTest(jtu.JaxTestCase):
         ValueError, 'Got invalid memory kind'):
       NamedSharding(abstract_mesh, P(), memory_kind='weird_device')
 
+  def test_pspec_unreduced(self):
+    pspec1 = P('a', 'b', None, unreduced=('c',))
+    self.assertEqual(repr(pspec1),
+                     "PartitionSpec('a', 'b', None, unreduced=('c',))")
+
+    pspec2 = P('a', 'b', None, unreduced=('c',))
+    self.assertEqual(pspec1, pspec2)
+
+    pspec3 = P('a', 'b', None, unreduced=('d',))
+    self.assertNotEqual(pspec1, pspec3)
+
+    out = P('x', unreduced=('z',)) + P('a', unreduced='b')
+    self.assertEqual(out, P('x', 'a', unreduced=('z', 'b')))
+
+    pspec4 = P('x', unreduced='y')
+    self.assertEqual(repr(pspec4),
+                     "PartitionSpec('x', unreduced=('y',))")
+
+    pspec5 = P(None, None, unreduced='x')
+    self.assertEqual(repr(pspec5),
+                     "PartitionSpec(None, None, unreduced=('x',))")
+
+    pspec6 = P(None, unreduced='x')
+    self.assertEqual(repr(pspec6), "PartitionSpec(None, unreduced=('x',))")
+
+    pspec7 = P(unreduced='x')
+    self.assertEqual(repr(pspec7), "PartitionSpec(unreduced=('x',))")
+
+    with self.assertRaisesRegex(
+        TypeError, 'unreduced in `__add__` of PartitionSpec'):
+      P('x', unreduced=('z',)) + (None,) * 2
+
+    with self.assertRaisesRegex(
+        TypeError, "unreduced in `__radd__` of PartitionSpec"):
+      (None,) * 2 + P('x', unreduced='y')
+
+    with self.assertRaisesRegex(
+        ValueError, "partitions cannot overlap with unreduced"):
+      P('x', 'y', unreduced='x')
+
+    with self.assertRaisesRegex(
+        ValueError, "partitions cannot overlap with unreduced"):
+      P('x', None, 'y', unreduced=('z', 'y'))
+
+  def test_named_sharding_unreduced_error(self):
+    mesh = jtu.create_mesh((1, 1, 1), ('x', 'y', 'z'))
+
+    with self.assertRaisesRegex(
+        ValueError, "Unreduced axes.*not found in mesh.*"):
+      NamedSharding(mesh, P('x', unreduced='a'))
+
+    with self.assertRaisesRegex(
+        ValueError, "Unreduced.*has duplicate entries"):
+      NamedSharding(mesh, P('x', unreduced=('y', 'y')))
+
+    with self.assertRaisesRegex(
+        ValueError, "Unreduced axes can only refer to mesh axes.*Explicit"):
+      NamedSharding(mesh, P('x', unreduced=('y', 'z')))
+
+    with self.assertRaisesRegex(
+        ValueError, "unreduced cannot contain None.*"):
+      NamedSharding(mesh, P('x', unreduced=('y', None)))
+
 
 @jtu.with_config(jax_use_shardy_partitioner=True)
 class ShardyShardingTest(jtu.JaxTestCase):
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index e8d8d4645..3215881a2 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -7677,8 +7677,88 @@ class ShardingInTypesTest(jtu.JaxTestCase):
       self.assertEqual(jax.typeof(out).sharding, jax.typeof(x).sharding)
       return out
 
-    f(arr)
-    jax.jit(f)(arr)
+    f(arr)  # doesn't crash
+    jax.jit(f)(arr)  # doesn't crash
+
+  @jtu.with_explicit_mesh((2, 2), ('x', 'y'))
+  def test_unreduced_basic(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', 'y'))
+    b = jax.device_put(np_inp.T, P('y', None))
+
+    @jax.jit
+    def f(x, y, a, b):
+      m1 = jnp.einsum('xy,yz->xz', x, y, out_sharding=P('x', unreduced='y'))
+      self.assertEqual(m1.aval.sharding.spec, P('x', None, unreduced='y'))
+
+      m2 = jnp.einsum('xy,yz->xz', a, b, out_sharding=P('x', unreduced='y'))
+      self.assertEqual(m2.aval.sharding.spec, P('x', None, unreduced='y'))
+
+      s = m1 + m2  # unreduced
+      self.assertEqual(s.aval.sharding.spec, P('x', None, unreduced='y'))
+
+      out = reshard(s, P('x'))  # reduce
+      self.assertEqual(out.aval.sharding.spec, P('x', None))
+      return out
+
+    f.trace(x, y, a, b)  # doesn't crash
+
+  @jtu.with_explicit_mesh((2, 2, 1), ('x', 'y', 'z'))
+  def test_dot_general_unreduced_error(self, mesh):
+    np_inp = np.arange(16).reshape(8, 2)
+    # Case 1
+    x = jax.device_put(np_inp, P('x', 'y'))
+    y = jax.device_put(np_inp.T, P('y', None))
+
+    @jax.jit
+    def f(x, y):
+      return jnp.einsum('xy,yz->xz', x, y, out_sharding=P('x', unreduced='z'))
+    with self.assertRaisesRegex(
+        core.ShardingTypeError,
+        "unreduced axes should be equal to the contracting specs"):
+      f.trace(x, y)
+
+    # Case 2
+    x = jax.device_put(np_inp, P('x', 'y'))
+    y = jax.device_put(np_inp.T, P(None, None))
+    @jax.jit
+    def g(x, y):
+      return jnp.einsum('xy,yz->xz', x, y, out_sharding=P('x', unreduced='y'))
+    with self.assertRaisesRegex(
+        core.ShardingTypeError,
+        "lhs and rhs contracting dims should be sharded identically"):
+      g.trace(x, y)
+
+    # Case 3
+    x = jax.device_put(np_inp, P('x', None))
+    y = jax.device_put(np_inp.T, P(None, None))
+
+    @jax.jit
+    def h(x, y):
+      return jnp.einsum('xy,yz->xz', x, y, out_sharding=P('x', unreduced='y'))
+    with self.assertRaisesRegex(
+        core.ShardingTypeError,
+        "unreduced axes should be equal to the contracting specs"):
+      h.trace(x, y)
+
+  @jtu.with_explicit_mesh((2, 2), ('x', 'y'))
+  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))
+
+    @jax.jit
+    def f(x, y):
+      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'))
+      return m1 + m2
+
+    with self.assertRaisesRegex(
+        core.ShardingTypeError,
+        "arrays must be unreduced along the same mesh axes"):
+      f.trace(x, y)
 
 
 @jtu.pytest_mark_if_available('multiaccelerator')