{"record":{"id":"73d669606a6a7b0e","repo":"jax-ml/jax","slug":"out-specs-passed-to-shard-map-should-be-equal-to-t-73d669","errorCode":null,"errorMessage":"out_specs passed to shard_map should be equal to the unreduced_kind present on the out_aval. Got out_specs={spec} and out_aval={aval.str_short(True)}","messagePattern":"out_specs passed to shard_map should be equal to the unreduced_kind present on the out_aval\\. Got out_specs=(.+?) and out_aval=(.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/shard_map.py","lineNumber":894,"sourceCode":"  unreduced = aval.sharding.spec.unreduced if check_vma else frozenset()\n  reduced = aval.sharding.spec.reduced if check_vma else frozenset()\n  u_kind = aval.sharding.spec.unreduced_kind if check_vma else None\n  mat = core.ManualAxisType(varying=vma, unreduced=unreduced, reduced=reduced,\n                            unreduced_kind=u_kind)\n  return aval.update(shape=new_shape, sharding=new_sharding,\n                     manual_axis_type=mat)\ncore.shard_aval_handlers[core.ShapedArray] = _shard_shaped_array\n\ndef _unshard_shaped_array(mesh: Mesh, check_vma, spec, aval: core.ShapedArray\n                          ) -> core.ShapedArray:\n  assert isinstance(aval, core.ShapedArray)\n  if check_vma and spec.unreduced != aval.mat.unreduced:\n    raise ValueError(\n        \"out_specs passed to shard_map should be equal to the unreduced\"\n        f\" present on the out_aval. Got out_specs={spec} and\"\n        f\" out_aval={aval.str_short(True)}\")\n  if check_vma and spec.unreduced_kind is not aval.mat.unreduced_kind:\n    raise ValueError(\n        \"out_specs passed to shard_map should be equal to the unreduced_kind\"\n        f\" present on the out_aval. Got out_specs={spec} and\"\n        f\" out_aval={aval.str_short(True)}\")\n  if check_vma and spec.reduced != aval.mat.reduced:\n    raise ValueError(\n        \"out_specs passed to shard_map should be equal to the reduced present\"\n        f\" on the out_aval. Got out_specs={spec} and\"\n        f\" out_aval={aval.str_short(True)}\")\n  names = _spec_to_names(spec)\n  new_shape = tuple(sz * prod(mesh.shape[n] for n in names.get(i, ()))\n                    for i, sz in enumerate(aval.shape))\n  names_spec = spec._normalized_spec_for_aval(aval.ndim).partitions\n  if aval.ndim == 0:\n    out_spec = P(unreduced=spec.unreduced, reduced=spec.reduced,\n                 unreduced_kind=spec.unreduced_kind)\n  else:\n    out_spec = []\n    for name_s, aval_s in zip(names_spec, aval.sharding.spec.partitions):","sourceCodeStart":876,"sourceCodeEnd":912,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/shard_map.py#L876-L912","documentation":"When checking is on, shard_map requires out_specs' unreduced_kind to equal the output aval's unreduced_kind. The combination rule declared for replicated partial outputs must match what the computation actually produced.","triggerScenarios":"Body returns arrays with unreduced_kind 'add' but out_specs declares 'mul' (or defaults differ).","commonSituations":"Switching reduction kind when porting code between JAX versions whose defaults for unreduced_kind changed.","solutions":["Match unreduced_kind in out_specs to the output aval's value","Regenerate outputs with the intended kind by fixing upstream shardings/ops","Inspect aval.mat.unreduced_kind to confirm the expected value"],"exampleFix":"// before\nout = shard_map(f, mesh, x, out_specs=P(unreduced=('r',), unreduced_kind='mul'))\n// after\nout = shard_map(f, mesh, x, out_specs=P(unreduced=('r',), unreduced_kind='add'))","handlingStrategy":"validation","validationCode":null,"typeGuard":null,"tryCatchPattern":"try: shard_map(...) except ValueError as e: if 'unreduced_kind' in str(e): align kind in out_specs; else: raise","preventionTips":["Use one constant for unreduced_kind across the codebase"],"tags":["shard-map","unreduced-kind","out-specs","sharding-mismatch","jax"],"backgroundTag":"sharding-spec-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}