{"record":{"id":"0c018f963fdb8343","repo":"jax-ml/jax","slug":"in-specs-containing-unreduced-spec-passed-to-sha","errorCode":null,"errorMessage":"in_specs containing unreduced {spec} passed to shard_map should be equal to the unreduced present on the in_aval {aval.str_short(True)}","messagePattern":"in_specs containing unreduced (.+?) passed to shard_map should be equal to the unreduced present on the in_aval (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/shard_map.py","lineNumber":855,"sourceCode":"    out_avals_lo = out_avals\n  out = trace.emit_eqn([*const_tracers, *in_tracers], out_avals_lo, prim, params,\n                       effs, source_info)\n  if trace.requires_low:\n    out = pe.raise_lo_outs(out_avals, out)\n  return out_avals_ft.update(out)\npe.DynamicJaxprTrace.process_shard_map = _shard_map_staging\n\n# TODO add underscore version, for direct-linearize to consume\n\ndef _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,","sourceCodeStart":837,"sourceCodeEnd":873,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/shard_map.py#L837-L873","documentation":"When an input array carries a sharding whose spec includes `unreduced` names, the in_specs passed to shard_map must declare exactly the same unreduced set. A mismatch means the caller is trying to reinterpret replication semantics inconsistently.","triggerScenarios":"Passing an array with NamedSharding spec P(..., unreduced=('r',)) into shard_map with in_specs=P(...) whose unreduced field is () or a different tuple.","commonSituations":"Chaining shard_map calls or jit->shard_map pipelines where the input was produced with unreduced semantics; hand-constructing specs that drift from the array's stored sharding.","solutions":["Make in_specs' unreduced tuple equal to the input array's aval sharding unreduced tuple","Re-create the input with the intended unreduced sharding before the call","Inspect aval.str_short(True) / arr.sharding to see the expected unreduced names"],"exampleFix":"// before\ny = shard_map(f, mesh, x_unreduced_on_r, in_specs=P('d'))\n// after\ny = shard_map(f, mesh, x_unreduced_on_r, in_specs=P('d', unreduced=('r',)))","handlingStrategy":"validation","validationCode":"def unreduced_matches(x, spec):\n    return getattr(getattr(x, 'sharding', None), 'spec', None) is None or spec.unreduced == x.sharding.spec.unreduced","typeGuard":"def in_spec_consistent(x, spec) -> bool:\n    s = getattr(getattr(x, 'sharding', None), 'spec', None)\n    return s is None or spec.unreduced == s.unreduced","tryCatchPattern":"try: shard_map(...) except ValueError as e: if 'unreduced' in str(e) and 'in_specs' in str(e): read aval via jax.core.shaped_abstractify and copy its unreduced into spec; else: raise","preventionTips":["Propagate sharding specs from input arrays when building in_specs","Print x.sharding.spec before crafting specs"],"tags":["shard-map","unreduced","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"}