{"record":{"id":"689985d17fd76446","repo":"jax-ml/jax","slug":"pargmin-only-accepts-a-single-axis-got-axis-name","errorCode":null,"errorMessage":"pargmin only accepts a single axis, got {axis_name}","messagePattern":"pargmin only accepts a single axis, got (.+?)","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/parallel.py","lineNumber":319,"sourceCode":"  if any(isinstance(axis, int) for axis in axis_name) and axis_index_groups is not None:\n    raise ValueError(\"axis_index_groups only supported for sums over just named axes\")\n  _validate_reduce_axis_index_groups(axis_index_groups)\n  axis_index_groups = _canonicalize_axis_index_groups(axis_index_groups)\n  def bind(leaf):\n    from_ = _get_from(core.typeof(leaf), axis_name, 'jax.lax.pmin')\n    if from_ == 'unreduced':\n      if axis_index_groups is not None:\n        raise NotImplementedError\n      return unreduced_pmin(leaf, axis_name)\n    else:\n      leaf = insert_collective_pvary(axis_name, leaf)\n      return pmin_p.bind(leaf, axes=axis_name, axis_index_groups=axis_index_groups)\n  return tree_util.tree_map(bind, x)\n\n# TODO(mattjj): add a pargmin_p, or add named axis support to lax.argmin_p\ndef pargmin(x, axis_name):\n  if isinstance(axis_name, (tuple, list)):\n    raise TypeError(f\"pargmin only accepts a single axis, got {axis_name}\")\n  return _axis_index_of_val(x, pmin(x, axis_name), axis_name)\n\n# TODO(mattjj): add a pargmax_p, or add named axis support to lax.argmax_p\ndef pargmax(x, axis_name):\n  if isinstance(axis_name, (tuple, list)):\n    raise TypeError(f\"pargmin only accepts a single axis, got {axis_name}\")\n  return _axis_index_of_val(x, pmax(x, axis_name), axis_name)\n\ndef _axis_index_of_val(x, val, axis_name):\n  idx = axis_index(axis_name)\n  mask = (val == x)\n  validx = lax.select(mask,\n                      lax.full(mask.shape, idx),\n                      lax.full(mask.shape, dtypes.iinfo(idx.dtype).max, idx.dtype))\n  return pmin(validx, axis_name)\n\ndef _validate_reduce_axis_index_groups(axis_index_groups):\n  if axis_index_groups is None:","sourceCodeStart":301,"sourceCodeEnd":337,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/parallel.py#L301-L337","documentation":"jax.lax.pargmin returns the index (along the mapped axis) of the minimum element; unlike psum/pmax/pmin it cannot reduce over multiple axes at once, so passing a tuple or list of axis names raises TypeError. (Note the message text says 'pargmin' even for pargmax — a copy-paste bug in JAX.)","triggerScenarios":"pargmin(x, ('i', 'j')) or pargmin(x, ['i','j']) — any sequence argument for axis_name.","commonSituations":"Generalizing multi-axis psum calls to arg-variants; reusing axis tuples built for other collectives.","solutions":["Pass a single axis name: pargmin(x, 'i')","For multiple axes, nest calls or combine them manually (compute per-axis pargmin sequentially)","Track upstream JAX issue for the wrong-function-name message if reporting it"],"exampleFix":"// before\nidx = jax.lax.pargmin(x, ('rows', 'cols'))\n\n// after\nidx_rows = jax.lax.pargmin(x, 'rows')","handlingStrategy":"type-guard","validationCode":"assert not isinstance(axis_name, (tuple, list)), 'pargmin takes one axis'","typeGuard":"def single_axis(axis_name):\n    assert not isinstance(axis_name, (tuple, list)), f'single axis required, got {axis_name}'\n    return axis_name","tryCatchPattern":null,"preventionTips":["Remember arg-variants of collectives take exactly one axis"],"tags":["jax","pargmin","axis","typeerror"],"backgroundTag":"invalid-axis-argument","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}