{"record":{"id":"6ae637af903947b2","repo":"jax-ml/jax","slug":"name-does-not-accept-integer-axis-name-got-axis","errorCode":null,"errorMessage":"{name} does not accept integer axis_name. Got axis_name={axes}","messagePattern":"(.+?) does not accept integer axis_name\\. Got axis_name=(.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/parallel.py","lineNumber":2773,"sourceCode":"def _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)\nunreduced_psum_p.def_effectful_abstract_eval(_unreduced_psum_abstract_eval)\n\ndef _unreduced_psum_lowering(ctx, arg, *, axes):\n  return _all_reduce_lowering(lax.add_p, lax.reduce_sum, ctx, arg,\n                             axes=axes, axis_index_groups=None)","sourceCodeStart":2755,"sourceCodeEnd":2791,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/parallel.py#L2755-L2791","documentation":"Unlike some collectives that accept integer positional axis indices, the unreduced p-sum/max/min collectives only accept named axes (strings). This ValueError fires when any element of the axis_name argument is an int.","triggerScenarios":"Calling `unreduced_psum(x, 0)` or passing an integer axis (e.g. from `range(x.ndim)`) instead of a named axis string.","commonSituations":"Copy-pasting code that uses `psum(x, axis=0)`-style positional axes; loops that build axis lists from integers; refactoring from lax.reduce-style APIs that use axis indices.","solutions":["Pass the named mesh axis (string) instead of an integer","If you meant a positional array axis, use a different API (e.g. `jax.lax.psum` with a name bound via mesh, or plain `jnp.sum(x, axis=0)`)"],"exampleFix":"// before\njax.lax.unreduced_psum(x, 0)\n// after\njax.lax.unreduced_psum(x, 'data')","handlingStrategy":"type-guard","validationCode":"assert all(isinstance(a, str) for a in jax.tree.leaves(axis_name)), 'axis_name must be strings, not ints'","typeGuard":"def is_named_axes(axis_name) -> bool:\n    return all(isinstance(a, str) for a in (axis_name if isinstance(axis_name, (tuple, list)) else (axis_name,)))","tryCatchPattern":null,"preventionTips":["Never pass integer axes to named collectives","Use mesh axis name constants instead of indices","Lint for integer axis arguments in collective wrappers"],"tags":["jax","axis-name","collectives","api-misuse"],"backgroundTag":"invalid-axis-argument","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}