{"record":{"id":"e59d0f4ced05b0b2","repo":"jax-ml/jax","slug":"all-gather-reduced-only-accepts-inputs-that-are-va","errorCode":null,"errorMessage":"all_gather_reduced only accepts inputs that are varying. Got {x_aval.str_short(True)}","messagePattern":"all_gather_reduced only accepts inputs that are varying\\. Got (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/parallel.py","lineNumber":2587,"sourceCode":"    return x\n  axis_size = _axis_size(axis_name, None)\n  def bind(leaf):\n    prim = all_gather_reduced_start_p if is_async else all_gather_reduced_p\n    return prim.bind(\n        leaf,\n        all_gather_dimension=canonicalize_axis(\n            axis, np.ndim(leaf) if tiled else np.ndim(leaf) + 1),\n        axis_name=axis_name, axis_size=axis_size, tiled=tiled)\n  return tree_util.tree_map(bind, x)\n\nall_gather_reduced_p = core.Primitive('all_gather_reduced')\n\ndef _all_gather_reduced_effectful_abstract_eval(\n    x_aval, *, all_gather_dimension, axis_name, axis_size, tiled\n):\n  _check_axis_names(axis_name, 'all_gather_reduced')\n  if not x_aval.mat.varying:\n    raise ValueError('all_gather_reduced only accepts inputs that are'\n                     f' varying. Got {x_aval.str_short(True)}')\n  # If the intersection between x.mat.varying and axis_name is empty, error\n  if not (x_aval.mat.varying & set(axis_name)):\n    raise ValueError(\n        'all_gather_reduced is a Varying -> Reduced collective. This means '\n        f'that the {axis_name=} passed to `all_gather_reduced` must be present '\n        f'in jax.typeof(x).mat.varying={x_aval.mat.varying}')\n  if x_aval.mat.reduced & set(axis_name):\n    raise ValueError(\n        \"all_gather_reduced's input cannot be reduced across the axis_name\"\n        f\" provided. Got x={x_aval.str_short(True)} and {axis_name=}\")\n\n  new_shape = list(x_aval.shape)\n  if tiled:\n    new_shape[all_gather_dimension] *= axis_size\n  else:\n    new_shape.insert(all_gather_dimension, axis_size)\n","sourceCodeStart":2569,"sourceCodeEnd":2605,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/parallel.py#L2569-L2605","documentation":"`all_gather_reduced` gathers an array that is varying (replicated-but-named) along the given axis, so its input must have a non-empty `mat.varying` set. The abstract eval rejects inputs with no varying axes because there is nothing to gather along. This almost always means the producer of the input never marked it varying (missing `pvary`).","triggerScenarios":"Calling `lax.all_gather_reduced(x, axis_name='i', ...)` on an input whose abstract value has `mat.varying == frozenset()`, e.g. a plain array or output of a fully reduced collective without a preceding pvary.","commonSituations":"Migrating psum/pmap code to named-axis SPMD style and forgetting the `pvary` annotation; reordering collectives so the varying-producing op is dropped or optimized away.","solutions":["Apply `lax.pvary(x, 'i')` before `all_gather_reduced`","Verify the producer of x actually varies the axis (check `jax.typeof(x).mat.varying`)","Ensure you meant all_gather_reduced rather than a plain all_gather on a sharded array"],"exampleFix":"// before\ny = lax.all_gather_reduced(x, axis_name='i', all_gather_dimension=0, axis_size=8, tiled=False)\n// after\nx = lax.pvary(x, 'i')\ny = lax.all_gather_reduced(x, axis_name='i', all_gather_dimension=0, axis_size=8, tiled=False)","handlingStrategy":"validation","validationCode":"if not jax.typeof(x).mat.varying:\n    x = lax.pvary(x, axis_name)\ny = lax.all_gather_reduced(x, axis_name=axis_name, ...)","typeGuard":"def is_varying(x) -> bool:\n    return bool(jax.typeof(x).mat.varying)","tryCatchPattern":"try:\n    y = lax.all_gather_reduced(x, axis_name=axis_name, ...)\nexcept ValueError as e:\n    if 'only accepts inputs that are varying' in str(e):\n        y = lax.all_gather_reduced(lax.pvary(x, axis_name), axis_name=axis_name, ...)\n    else:\n        raise","preventionTips":["Establish a convention: every mesh-annotated value is pvary'd at creation","Assert mat.varying is non-empty before gather collectives in debug builds"],"tags":["jax","collectives","mesh","spmd","all-gather"],"backgroundTag":"jax-mat-axis-state-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}