{"record":{"id":"e1c699830540db0a","repo":"jax-ml/jax","slug":"in-specs-containing-reduced-spec-passed-to-shard","errorCode":null,"errorMessage":"in_specs containing reduced {spec} passed to shard_map should be equal to the reduced present on the in_aval {aval.str_short(True)}","messagePattern":"in_specs containing reduced (.+?) passed to shard_map should be equal to the reduced present on the in_aval (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/shard_map.py","lineNumber":865,"sourceCode":"def _spec_to_names(spec: PartitionSpec):\n  return {i: names if isinstance(names, tuple) else (names,)\n          for i, names in enumerate(spec.partitions) if names is not None}\n\ndef _shard_shaped_array(mesh: Mesh, manual_axes: frozenset, check_vma,\n                        spec, aval: core.ShapedArray) -> core.ShapedArray:\n  assert isinstance(aval, core.ShapedArray)\n  if spec.unreduced != aval.sharding.spec.unreduced:\n    raise ValueError(\n        f\"in_specs containing unreduced {spec} passed to shard_map should be\"\n        \" equal to the unreduced present on the in_aval\"\n        f\" {aval.str_short(True)}\")\n  if spec.unreduced_kind is not aval.sharding.spec.unreduced_kind:\n    raise ValueError(\n        f\"in_specs containing unreduced_kind {spec} passed to shard_map should\"\n        \" be equal to the unreduced_kind present on the in_aval\"\n        f\" {aval.str_short(True)}\")\n  if spec.reduced != aval.sharding.spec.reduced:\n    raise ValueError(\n        f\"in_specs containing reduced {spec} passed to shard_map should be\"\n        f\" equal to the reduced present on the in_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  manual_mesh = _as_manual_mesh(mesh, manual_axes)\n  new_sharding = aval.sharding.update(\n      mesh=manual_mesh,\n      spec=core.modify_spec_for_auto_manual(aval.sharding.spec, manual_mesh))\n  vma = (_spec_to_vma(spec) if check_vma else frozenset()) | aval.mat.varying\n  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","sourceCodeStart":847,"sourceCodeEnd":883,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/shard_map.py#L847-L883","documentation":"When an input array's sharding spec includes `reduced` names, shard_map requires the in_specs' reduced tuple to exactly equal the aval's reduced tuple, so reduction semantics are not silently redefined at the boundary.","triggerScenarios":"Input array sharded with reduced=('y',) but in_specs declares reduced=() or a different axis tuple.","commonSituations":"Pipelines where a previous stage produced partially-reduced arrays (e.g. SP-style gradients) that are fed into another shard_map with fresh specs.","solutions":["Copy the reduced tuple from the input's sharding spec into in_specs","Normalize the input (e.g. fully reduce / reshard) before entering shard_map","Print aval.str_short(True) to confirm what reduced names are expected"],"exampleFix":"// before\ny = shard_map(f, mesh, x_reduced_y, in_specs=P('d'))\n// after\ny = shard_map(f, mesh, x_reduced_y, in_specs=P('d', reduced=('y',)))","handlingStrategy":"validation","validationCode":"def reduced_matches(x, spec):\n    s = getattr(getattr(x, 'sharding', None), 'spec', None)\n    return s is None or spec.reduced == s.reduced","typeGuard":"def reduced_consistent(x, spec) -> bool:\n    s = getattr(getattr(x, 'sharding', None), 'spec', None)\n    return s is None or spec.reduced == s.reduced","tryCatchPattern":null,"preventionTips":["Feed partially-reduced arrays with specs derived from their own sharding"],"tags":["shard-map","reduced","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"}