jax-ml/jax · error · ValueError

`out_sharding` passed to `dot_general` can only contain unre

Error message

`out_sharding` passed to `dot_general` can only contain unreduced of kind `sum`. Got {out_sharding=}

What it means

Raised when an out_sharding passed to dot_general declares unreduced output axes whose kind is not UnreducedKind.sum. For SPMD sharded dot products, the only permissible unreduced-kind on the output is 'sum' (partial sums awaiting reduction).

Source

Thrown at jax/_src/lax/lax.py:5839

    (lhs_contracting, rhs_contracting), _ = dimension_numbers
    lhs_contracting_spec = tuple(lhs.sharding.spec.partitions[i]
                                 for i in lhs_contracting)
    rhs_contracting_spec = tuple(rhs.sharding.spec.partitions[i]
                                 for i in rhs_contracting)
    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 {lhs_contracting_spec=}, {rhs_contracting_spec=}')
    flat_spec = [s for s in flatten_spec(lhs_contracting_spec) if s is not None]
    if out_sharding.spec.unreduced != frozenset(flat_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}')
    out_u, out_k = out_sharding.spec.unreduced, out_sharding.spec.unreduced_kind
    if out_k is not None and out_k is not UnreducedKind.sum:
      raise ValueError(
          '`out_sharding` passed to `dot_general` can only contain'
          f' unreduced of kind `sum`. Got {out_sharding=}')
    return out_u, out_k
  return frozenset(), None

def _dot_general_ur_rule(lhs, rhs, *, dimension_numbers, out_sharding, **kwargs):
  out_unreduced, kind = _dot_general_unreduced_rule(lhs, rhs, dimension_numbers,
                                                    out_sharding)
  # TODO(yashkatariya): Propagate reduced and make checks like nary_reduced_rule
  return out_unreduced, frozenset(), kind

def tuple_delete(tup, idx):
  idx_ = set(idx)
  return tuple(tup[i] for i in range(len(tup)) if i not in idx_)


def _dot_general_dtype_rule(lhs, rhs, *, dimension_numbers, precision,
                            preferred_element_type: DTypeLike | None,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Drop the custom out_sharding and let sharding propagation infer it
2. Set unreduced_kind to UnreducedKind.sum (or None) in the spec
3. File/check the experimental API docs if you genuinely need other reduction kinds

if out_sharding is not None:
    k = getattr(out_sharding.spec, 'unreduced_kind', None) if hasattr(out_sharding, 'spec') else None
    assert k is None or k.__class__.__name__ == 'sum' or str(k).endswith('sum'), f'bad unreduced kind: {k}'

When it happens

Trigger: Passing out_sharding=NamedSharding(...) with a GSPMDSharding spec whose unreduced_kind is e.g. min/max/product to jax.lax.dot_general.

Common situations: Experimenting with the sharded/SPMD dot API or automatic sharding propagation; constructing custom sharding specs by hand instead of from named axes.

Related errors

• callbacks are only supported in spmd computations when all m
• callbacks do not support specifying sharding inside spmd com
• {name} cannot accept args which are unreduced. Got {a.str_sh
• {name} cannot accept args with unreduced_kind={a.mat.unreduc
• {name} cannot accept args which are reduced. Got {a.str_shor

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/af034c79faa71381. Report an issue: GitHub.