{"record":{"id":"c4accd0fe7abd87b","repo":"jax-ml/jax","slug":"all-gather-reduced-is-a-varying-reduced-collect","errorCode":null,"errorMessage":"all_gather_reduced is a Varying -> Reduced collective. This means that the {axis_name=} passed to `all_gather_reduced` must be present in jax.typeof(x).mat.varying={x_aval.mat.varying}","messagePattern":"all_gather_reduced is a Varying -> Reduced collective\\. This means that the (.+?) passed to `all_gather_reduced` must be present in jax\\.typeof\\(x\\)\\.mat\\.varying=(.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/parallel.py","lineNumber":2591,"sourceCode":"    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\n  if x_aval.mat.unreduced:\n    check_unreduced_kind('all_gather_reduced', x_aval.mat, UnreducedKind.sum)\n  new_reduced = x_aval.mat.reduced | frozenset(axis_name)\n  out_vma = frozenset(v for v in x_aval.mat.varying if v not in axis_name)","sourceCodeStart":2573,"sourceCodeEnd":2609,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/parallel.py#L2573-L2609","documentation":"`all_gather_reduced` is a Varying -> Reduced collective: the `axis_name` you pass must appear in `jax.typeof(x).mat.varying`. If the input varies along other axes but not the requested one, gathering along it is meaningless, so the abstract eval raises this error. It indicates the array and the collective disagree about which mesh axis to operate on.","triggerScenarios":"Input varied along axis 'i' but calling `all_gather_reduced(x, axis_name='j', ...)`; passing a tuple of axis names none of which intersect `x.mat.varying`.","commonSituations":"Renaming mesh axes or copy-pasting collective calls between pipelines with different axis names; typos in axis_name strings.","solutions":["Make axis_name match an axis in `jax.typeof(x).mat.varying` (print it to confirm)","Add `lax.pvary(x, axis_name)` if the axis was never varied","Check for axis-name typos against the mesh definition"],"exampleFix":"// before\nx = lax.pvary(x, 'i')\ny = lax.all_gather_reduced(x, axis_name='j', ...)\n// after\nx = lax.pvary(x, 'i')\ny = lax.all_gather_reduced(x, axis_name='i', ...)","handlingStrategy":"validation","validationCode":"assert set(jax.typeof(x).mat.varying) & set(axis_name), (\n    f'axis_name {axis_name} not in {jax.typeof(x).mat.varying}')\ny = lax.all_gather_reduced(x, axis_name=axis_name, ...)","typeGuard":"def axis_is_varying(x, axis_name) -> bool:\n    return bool(set(jax.typeof(x).mat.varying) & set(axis_name))","tryCatchPattern":"try:\n    y = lax.all_gather_reduced(x, axis_name=axis_name, ...)\nexcept ValueError as e:\n    if 'Varying -> Reduced' in str(e):\n        correct = next(iter(jax.typeof(x).mat.varying))\n        y = lax.all_gather_reduced(x, axis_name=correct, ...)\n    else:\n        raise","preventionTips":["Derive axis_name from jax.typeof(x).mat.varying instead of hardcoding","Define mesh axis names once as module constants"],"tags":["jax","collectives","mesh","axis-name","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"}