{"record":{"id":"af034c79faa71381","repo":"jax-ml/jax","slug":"out-sharding-passed-to-dot-general-can-only-co","errorCode":null,"errorMessage":"`out_sharding` passed to `dot_general` can only contain unreduced of kind `sum`. Got {out_sharding=}","messagePattern":"`out_sharding` passed to `dot_general` can only contain unreduced of kind `sum`\\. Got (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":5839,"sourceCode":"    (lhs_contracting, rhs_contracting), _ = dimension_numbers\n    lhs_contracting_spec = tuple(lhs.sharding.spec.partitions[i]\n                                 for i in lhs_contracting)\n    rhs_contracting_spec = tuple(rhs.sharding.spec.partitions[i]\n                                 for i in rhs_contracting)\n    if lhs_contracting_spec != rhs_contracting_spec:\n      raise core.ShardingTypeError(\n          'lhs and rhs contracting dims should be sharded identically when'\n          ' out_sharding provided to dot_general mentions unreduced_axes.'\n          f' Got {lhs_contracting_spec=}, {rhs_contracting_spec=}')\n    flat_spec = [s for s in flatten_spec(lhs_contracting_spec) if s is not None]\n    if out_sharding.spec.unreduced != frozenset(flat_spec):\n      raise core.ShardingTypeError(\n          \"out_sharding's unreduced axes should be equal to the contracting\"\n          f' specs. Got unreduced axes={out_sharding.spec.unreduced} and'\n          f' contracting spec={lhs_contracting_spec}')\n    out_u, out_k = out_sharding.spec.unreduced, out_sharding.spec.unreduced_kind\n    if out_k is not None and out_k is not UnreducedKind.sum:\n      raise ValueError(\n          '`out_sharding` passed to `dot_general` can only contain'\n          f' unreduced of kind `sum`. Got {out_sharding=}')\n    return out_u, out_k\n  return frozenset(), None\n\ndef _dot_general_ur_rule(lhs, rhs, *, dimension_numbers, out_sharding, **kwargs):\n  out_unreduced, kind = _dot_general_unreduced_rule(lhs, rhs, dimension_numbers,\n                                                    out_sharding)\n  # TODO(yashkatariya): Propagate reduced and make checks like nary_reduced_rule\n  return out_unreduced, frozenset(), kind\n\ndef tuple_delete(tup, idx):\n  idx_ = set(idx)\n  return tuple(tup[i] for i in range(len(tup)) if i not in idx_)\n\n\ndef _dot_general_dtype_rule(lhs, rhs, *, dimension_numbers, precision,\n                            preferred_element_type: DTypeLike | None,","sourceCodeStart":5821,"sourceCodeEnd":5857,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L5821-L5857","documentation":"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).","triggerScenarios":"Passing out_sharding=NamedSharding(...) with a GSPMDSharding spec whose unreduced_kind is e.g. min/max/product to jax.lax.dot_general.","commonSituations":"Experimenting with the sharded/SPMD dot API or automatic sharding propagation; constructing custom sharding specs by hand instead of from named axes.","solutions":["Drop the custom out_sharding and let sharding propagation infer it","Set unreduced_kind to UnreducedKind.sum (or None) in the spec","File/check the experimental API docs if you genuinely need other reduction kinds"],"exampleFix":null,"handlingStrategy":"validation","validationCode":"if out_sharding is not None:\n    k = getattr(out_sharding.spec, 'unreduced_kind', None) if hasattr(out_sharding, 'spec') else None\n    assert k is None or k.__class__.__name__ == 'sum' or str(k).endswith('sum'), f'bad unreduced kind: {k}'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Avoid hand-built out_sharding; use NamedSharding from a Mesh","Track experimental sharded-dot API changes between JAX versions"],"tags":["jax","dot-general","sharding","spmd","experimental-api"],"backgroundTag":"invalid-sharding-spec","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}