jax-ml/jax · error · ValueError

Invalid axis {axis} for operand shape {operand.shape}

Error message

Invalid axis {axis} for operand shape {operand.shape}

What it means

argmin/argmax shape validation found the reduction axis outside the valid range for the operand's shape. The single axis must satisfy 0 <= axis < len(operand.shape).

Source

Thrown at jax/_src/lax/lax.py:8659

def _reduce_min_ur_rule(operand, *, axes, out_sharding):
  out_unreduced, kind = _reduce_op_unreduced_rule(
      operand, axes, out_sharding, UnreducedKind.min, 'reduce_min')
  out_reduced = _reduce_op_reduced_rule(operand, out_sharding, 'reduce_min')
  return out_unreduced, out_reduced, kind

reduce_min_p = standard_primitive(
    _reduce_op_shape_rule, input_dtype, 'reduce_min',
    sharding_rule=_reduce_op_sharding_rule_with_out_sharding,
    vma_rule=partial(core.standard_vma_rule, 'reduce_min'),
    ur_rule=_reduce_min_ur_rule)
ad.defjvp2(reduce_min_p, _reduce_chooser_jvp_rule)
batching.defreducer(reduce_min_p)

def _argminmax_shape_rule(operand, *, axes, index_dtype):
  axis, = axes
  if not (0 <= axis < len(operand.shape)):
    raise ValueError(f"Invalid axis {axis} for operand shape {operand.shape}")
  if operand.shape[axis] < 1:
    raise ValueError("argmin and argmax require non-empty reduced dimension. "
                     f"operand.shape={operand.shape} {axis=}")
  return util.tuple_delete(operand.shape, axis)

def _argminmax_sharding_rule(operand, *, axes, index_dtype):
  axis, = axes
  return operand.sharding.update(spec=
      util.tuple_delete(operand.sharding.spec, axis))

def _argminmax_dtype_rule(operand, *, axes, index_dtype):
  if not dtypes.issubdtype(index_dtype, np.integer):
    raise TypeError("index_dtype must be an integer type, but got {}"
                    .format(dtype_to_string(index_dtype)))
  return index_dtype

class _ArgMinMaxReducer:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Log operand.shape and axis; fix the axis to a valid one or use -1 for the last axis.
2. Guard: axis = axis % x.ndim before calling argmin/argmax.
3. If rank can vary, branch on x.ndim instead of assuming fixed rank.

Example fix

# before
i = jnp.argmax(x, axis=2)  # x.shape == (8, 4)
# after
i = jnp.argmax(x, axis=-1)

axis = axis % x.ndim if isinstance(axis, int) else axis
assert 0 <= axis < x.ndim, (axis, x.shape)
i = jnp.argmax(x, axis=axis)

def valid_axis(x, axis):
    return -x.ndim <= axis < x.ndim

When it happens

Trigger: lax.argmin(x, axis=2) on a 2-D array; jnp.argmax(x, axis=x.ndim) after squeeze removed a dimension. The check runs during abstract evaluation, so it fires at trace time.

Common situations: Hardcoded axis constants broken by a reshape/squeeze upstream; axis derived from a configuration dict; code ported from NumPy where a later check would have caught it differently.

Related errors

• axis_index_groups only supported for sums over just named ax
• {full_name} must be a pytree prefix with bool leaves or a tu
• Axes mentioned in `manual_axis_type` field of ShapedArray sh
• varying and unreduced cannot have common mesh axes. Got vary
• varying and reduced cannot have common mesh axes. Got varyin

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