{"record":{"id":"1962c972cdb75a4c","repo":"jax-ml/jax","slug":"in-specs-containing-unreduced-kind-spec-passed-t","errorCode":null,"errorMessage":"in_specs containing unreduced_kind {spec} passed to shard_map should be equal to the unreduced_kind present on the in_aval {aval.str_short(True)}","messagePattern":"in_specs containing unreduced_kind (.+?) passed to shard_map should be equal to the unreduced_kind present on the in_aval (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/shard_map.py","lineNumber":860,"sourceCode":"  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,\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","sourceCodeStart":842,"sourceCodeEnd":878,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/shard_map.py#L842-L878","documentation":"Input arrays with unreduced sharding also carry an `unreduced_kind` (e.g. how the replicated partial values combine). shard_map requires the in_specs' unreduced_kind to equal the aval's stored unreduced_kind; a mismatch raises this ValueError.","triggerScenarios":"Passing an array whose sharding was created with a different unreduced_kind (e.g. 'add' vs 'mul') than the unreduced_kind declared in in_specs.","commonSituations":"Mixing APIs that default unreduced_kind differently across JAX versions; constructing NamedSharding manually with mismatched kind parameters.","solutions":["Align unreduced_kind in in_specs with the input's sharding (inspect arr.sharding.spec.unreduced_kind)","Rebuild the input array's sharding with the desired unreduced_kind before shard_map","Avoid hand-building unreduced shardings; use the same helper on both sides"],"exampleFix":"// before\ny = shard_map(f, mesh, x, in_specs=P(unreduced=('r',), unreduced_kind='mul'))  # x is 'add'\n// after\ny = shard_map(f, mesh, x, in_specs=P(unreduced=('r',), unreduced_kind='add'))","handlingStrategy":"validation","validationCode":"def kind_matches(x, spec):\n    s = getattr(getattr(x, 'sharding', None), 'spec', None)\n    return s is None or spec.unreduced_kind is s.unreduced_kind","typeGuard":"def unreduced_kind_consistent(x, spec) -> bool:\n    s = getattr(getattr(x, 'sharding', None), 'spec', None)\n    return s is None or spec.unreduced_kind is s.unreduced_kind","tryCatchPattern":null,"preventionTips":["Create unreduced shardings through a single helper so kind is uniform","Pin JAX version to avoid default unreduced_kind drift"],"tags":["shard-map","unreduced-kind","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"}