{"record":{"id":"0d3d4098ace6a789","repo":"jax-ml/jax","slug":"name-s-input-cannot-be-varying-across-the-axis","errorCode":null,"errorMessage":"{name}'s input cannot be varying across the  axis_name provided. Got x={aval.str_short(True)} and {axes=}","messagePattern":"(.+?)'s input cannot be varying across the  axis_name provided\\. Got x=(.+?) and (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/parallel.py","lineNumber":2768,"sourceCode":"  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)\n  out_aval = aval.update(manual_axis_type=out_mat)\n  return out_aval, {core.NamedAxisEffect(axis) for axis in axes}\n\ndef _unreduced_psum_abstract_eval(aval, *, axes):\n  return _unreduced_psum_pmax_pmin_abstract_eval(\n      'unreduced_psum', UnreducedKind.sum, aval, axes=axes)","sourceCodeStart":2750,"sourceCodeEnd":2786,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/parallel.py#L2750-L2786","documentation":"The unreduced collectives reject inputs that are already marked as 'varying' across the requested axis. Since the collective would reduce across that axis, JAX requires the input not vary along it; this ValueError reports the input aval and the offending axes.","triggerScenarios":"Calling `unreduced_psum`/`unreduced_pmax`/`unreduced_pmin` with an axis_name that appears in `jax.typeof(x).mat.varying`.","commonSituations":"Passing a per-device-different (sharded/varying) tensor to a collective expecting a replicated-but-unreduced value; inconsistent sharding annotations after mesh or sharding spec changes.","solutions":["Remove the axis from the input's varying set (annotate/produce the value as not varying along it) before calling the collective","Use a different collective appropriate for varying data (e.g. all_gather plus reduce) ","Verify `jax.typeof(x).mat.varying` and reconcile it with the axis_name argument"],"exampleFix":"// before\njax.lax.unreduced_psum(x, 'data')  # x.varying includes 'data'\n// after\njax.lax.unreduced_psum(x, 'data')  # after ensuring x is not varying along 'data', e.g. reshard/annotate accordingly","handlingStrategy":"validation","validationCode":"import jax\nassert not (set(jax.typeof(x).mat.varying) & {axis_name}), 'input varies along axis'","typeGuard":"def is_invariant_along(x, axis: str) -> bool:\n    return axis not in jax.typeof(x).mat.varying","tryCatchPattern":null,"preventionTips":["Validate mat.varying before unreduced collectives","Keep sharding specs consistent with collective axis arguments","Add unit tests asserting sharding annotations on mesh axes"],"tags":["jax","sharding","collectives","varying","mesh"],"backgroundTag":"sharding-annotation-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}