jax-ml/jax · error · TypeError

pargmin only accepts a single axis, got {axis_name}

Error message

pargmin only accepts a single axis, got {axis_name}

What it means

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.)

Source

Thrown at jax/_src/lax/parallel.py:319

  if any(isinstance(axis, int) for axis in axis_name) and axis_index_groups is not None:
    raise ValueError("axis_index_groups only supported for sums over just named axes")
  _validate_reduce_axis_index_groups(axis_index_groups)
  axis_index_groups = _canonicalize_axis_index_groups(axis_index_groups)
  def bind(leaf):
    from_ = _get_from(core.typeof(leaf), axis_name, 'jax.lax.pmin')
    if from_ == 'unreduced':
      if axis_index_groups is not None:
        raise NotImplementedError
      return unreduced_pmin(leaf, axis_name)
    else:
      leaf = insert_collective_pvary(axis_name, leaf)
      return pmin_p.bind(leaf, axes=axis_name, axis_index_groups=axis_index_groups)
  return tree_util.tree_map(bind, x)

# TODO(mattjj): add a pargmin_p, or add named axis support to lax.argmin_p
def pargmin(x, axis_name):
  if isinstance(axis_name, (tuple, list)):
    raise TypeError(f"pargmin only accepts a single axis, got {axis_name}")
  return _axis_index_of_val(x, pmin(x, axis_name), axis_name)

# TODO(mattjj): add a pargmax_p, or add named axis support to lax.argmax_p
def pargmax(x, axis_name):
  if isinstance(axis_name, (tuple, list)):
    raise TypeError(f"pargmin only accepts a single axis, got {axis_name}")
  return _axis_index_of_val(x, pmax(x, axis_name), axis_name)

def _axis_index_of_val(x, val, axis_name):
  idx = axis_index(axis_name)
  mask = (val == x)
  validx = lax.select(mask,
                      lax.full(mask.shape, idx),
                      lax.full(mask.shape, dtypes.iinfo(idx.dtype).max, idx.dtype))
  return pmin(validx, axis_name)

def _validate_reduce_axis_index_groups(axis_index_groups):
  if axis_index_groups is None:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass a single axis name: pargmin(x, 'i')
2. For multiple axes, nest calls or combine them manually (compute per-axis pargmin sequentially)
3. Track upstream JAX issue for the wrong-function-name message if reporting it

Example fix

// before
idx = jax.lax.pargmin(x, ('rows', 'cols'))

// after
idx_rows = jax.lax.pargmin(x, 'rows')

assert not isinstance(axis_name, (tuple, list)), 'pargmin takes one axis'

def single_axis(axis_name):
    assert not isinstance(axis_name, (tuple, list)), f'single axis required, got {axis_name}'
    return axis_name

When it happens

Trigger: pargmin(x, ('i', 'j')) or pargmin(x, ['i','j']) — any sequence argument for axis_name.

Common situations: Generalizing multi-axis psum calls to arg-variants; reusing axis tuples built for other collectives.

Related errors

• iteration over a 0-d array
• Formatting arguments to checkify.check need to be PyTrees of
• check_error takes an Error as argument, got type {type(error
• bool() not supported for instances of type '{0}' (did you me
• Default value must be of type bool, got {default} of type {g

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/689985d17fd76446. Report an issue: GitHub.