{"record":{"id":"78c6fc2b23e883d2","repo":"jax-ml/jax","slug":"name-is-a-unreduced-invariant-collective-thi","errorCode":null,"errorMessage":"{name} is a Unreduced -> Invariant collective. This means that the {axes=} passed to `{name}` must be present in jax.typeof(x).mat.unreduced={aval.mat.unreduced}","messagePattern":"(.+?) is a Unreduced -> Invariant collective\\. This means that the (.+?) passed to `(.+?)` must be present in jax\\.typeof\\(x\\)\\.mat\\.unreduced=(.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/parallel.py","lineNumber":2762,"sourceCode":"# Unreduced -> Invariant collective\ndef unreduced_psum(x, axis_name):\n  if not isinstance(axis_name, (tuple, list)):\n    axis_name = (axis_name,)\n  if not axis_name:\n    return x\n  return tree_util.tree_map(\n      lambda leaf: unreduced_psum_p.bind(leaf, axes=tuple(axis_name)), x)\n\nunreduced_psum_p = core.Primitive('unreduced_psum')\n\ndef _unreduced_psum_pmax_pmin_abstract_eval(name, out_u_kind, aval, *, axes):\n  _check_axis_names(axes, name)\n  if not aval.mat.unreduced:\n    raise ValueError(f'{name} only accepts inputs that are'\n                     f' unreduced. Got {aval.str_short(True)}')\n  # If intersection between x.unreduced & axis_name is empty, error\n  if not (aval.mat.unreduced & frozenset(axes)):\n    raise ValueError(\n        f\"{name} is a Unreduced -> Invariant collective. This\"\n        f\" means that the {axes=} passed to `{name}` must\"\n        \" be present in\"\n        f\" jax.typeof(x).mat.unreduced={aval.mat.unreduced}\")\n  if aval.mat.varying & set(axes):\n    raise ValueError(\n        f\"{name}'s input cannot be varying across the \"\n        f\" axis_name provided. Got x={aval.str_short(True)} and {axes=}\")\n\n  if any(isinstance(a, int) for a in axes):\n    raise ValueError(f'{name} does not accept integer axis_name.'\n                     f' Got axis_name={axes}')\n\n  core.check_avals_context_mesh([aval], name)\n  check_unreduced_kind(name, aval.mat, out_u_kind)\n  out_u = frozenset(u for u in aval.mat.unreduced if u not in axes)\n  kind = aval.mat.unreduced_kind if out_u else None\n  out_mat = aval.mat.update(unreduced=out_u, unreduced_kind=kind)","sourceCodeStart":2744,"sourceCodeEnd":2780,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/parallel.py#L2744-L2780","documentation":"`unreduced_psum`/`unreduced_pmax`/`unreduced_pmin` are Unreduced -> Invariant collectives: the axis names passed to them must already appear in the input's `mat.unreduced` set. This ValueError fires when the intersection of `aval.mat.unreduced` and the requested axes is empty — the value is unreduced along some axes, but not the ones you asked to reduce.","triggerScenarios":"Calling `unreduced_psum(x, 'data')` when `jax.typeof(x).mat.unreduced` is e.g. `{'batch'}` but not `'data'`; passing an axis_name not present in the input's unreduced set (even though the set itself is non-empty).","commonSituations":"Typos or renames of mesh axis names between model definition and collective call; using the wrong axis from a multi-axis mesh; refactors that change which axis a tensor is unreduced along.","solutions":["Print `jax.typeof(x).mat.unreduced` and pass one of those axis names to the collective","Fix the axis name typo / use the correct mesh axis the value is unreduced along","Adjust upstream code so the value is produced unreduced along the intended axis"],"exampleFix":"// before\njax.lax.unreduced_psum(x, 'data')  # x unreduced along 'batch'\n// after\njax.lax.unreduced_psum(x, 'batch')","handlingStrategy":"validation","validationCode":"import jax\naxes = {'data'}\nassert axes & set(jax.typeof(x).mat.unreduced), 'axis not in x.mat.unreduced'","typeGuard":"def axis_is_unreduced(x, axis: str) -> bool:\n    return axis in jax.typeof(x).mat.unreduced","tryCatchPattern":null,"preventionTips":["Always derive axis names from jax.typeof(x).mat.unreduced","Validate mesh axis names at config load time","Use constants for mesh axis names to avoid typos"],"tags":["jax","mesh","axis-name","collectives","sharding"],"backgroundTag":"sharding-annotation-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}