{"record":{"id":"91ae2cb57d432be9","repo":"jax-ml/jax","slug":"name-only-accepts-inputs-that-are-unreduced-got","errorCode":null,"errorMessage":"{name} only accepts inputs that are unreduced. Got {aval.str_short(True)}","messagePattern":"(.+?) only accepts inputs that are unreduced\\. Got (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/parallel.py","lineNumber":2758,"sourceCode":"                       partial(_unreduced_reduce_scatter_lowering, lax.add_p))\n\n############################## unreduced_psum ###########################\n\n# 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)","sourceCodeStart":2740,"sourceCodeEnd":2776,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/parallel.py#L2740-L2776","documentation":"JAX's unreduced collectives (`unreduced_psum`, `unreduced_pmax`, `unreduced_pmin`) require their input to actually be marked as 'unreduced' along the corresponding mesh axes in the sharding annotation (the `mat.unreduced` set of the aval). This ValueError fires when the abstract-eval sees an input whose unreduced set is empty, i.e. the value was never produced/constrained to be unreduced.","triggerScenarios":"Calling `jax.lax.psum(x, axis_name)` on a plain (fully materialized / invariant) array under a named axis context where the API routes to the unreduced variant; passing a value to `unreduced_psum`/`unreduced_pmax`/`unreduced_pmin` that has `mat.unreduced` empty.","commonSituations":"Mixing the newer unreduced-collective API with code that assumes values are replicated/invariant; sharding-annotation mismatch after upgrading JAX to a version with mat (materialization) tracking; forgetting to produce the input via an operation that marks it unreduced.","solutions":["Ensure the input was produced by a pipeline that marks it unreduced (e.g. the output of an operation returning unreduced values along that axis)","Use the standard collective (`psum`, `pmax`, `pmin`) instead of the unreduced variant if your input is invariant","Check that the axis_name matches the mesh axis along which the value is actually unreduced","Inspect `jax.typeof(x).mat` before the call to confirm the unreduced set"],"exampleFix":"// before\ny = jax.lax.unreduced_psum(x, 'data')  # x is invariant\n// after\ny = jax.lax.psum(x, 'data')  # x is not unreduced; use standard psum","handlingStrategy":"validation","validationCode":"import jax\nt = jax.typeof(x)\nassert t.mat.unreduced, f'input not unreduced: {t.mat}'","typeGuard":"def is_unreduced(x) -> bool:\n    return bool(jax.typeof(x).mat.unreduced)","tryCatchPattern":null,"preventionTips":["Check jax.typeof(x).mat before calling unreduced collectives","Use standard psum/pmax/pmin for invariant inputs","Document which producers yield unreduced values"],"tags":["jax","sharding","collectives","validation","mesh"],"backgroundTag":"sharding-annotation-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}