jax-ml/jax · error · ValueError

Invalid axis for operand shape

Error message

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

What it means

A reduction operation (reduce_min/reduce_max family) was given an axis index that is outside the valid range for the operand's rank. The shape rule canonicalizes and validates each entry of the axes argument against operand.shape and rejects out-of-bounds or non-integer axes.

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…
• A compound factor should contain at least two factors
• accumulate does not allow multiple axes
• All arrays must have matching shapes along non-concatenated…
• All arrays must have the same dtype, got

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

Appendix: 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)