jax-ml/jax · error · ValueError

attempt to get argmax of an empty sequence

Error message

attempt to get argmax of an empty sequence

What it means

Raised by jnp.argmax when the (possibly raveled) array has zero elements along the reduction axis; argmax of an empty sequence is undefined and, unlike NumPy's error-prone behavior, JAX raises deterministically.

Source

Thrown at jax/_src/numpy/lax_numpy.py:8286

    Array([[1],
           [0]], dtype=int32)
  """
  arr = util.ensure_arraylike("argmax", a)
  if out is not None:
    raise NotImplementedError("The 'out' argument to jnp.argmax is not supported.")
  return _argmax(arr, None if axis is None else operator.index(axis),
                 keepdims=bool(keepdims))

@api.jit(static_argnames=('axis', 'keepdims'), inline=True)
def _argmax(a: Array, axis: int | None = None, keepdims: bool = False) -> Array:
  if axis is None:
    dims = list(range(np.ndim(a)))
    a = ravel(a)
    axis = 0
  else:
    dims = [axis]
  if a.shape[axis] == 0:
    raise ValueError("attempt to get argmax of an empty sequence")
  # TODO(phawkins): use an int64 index if the dimension is large enough.
  result = lax.argmax(a, _canonicalize_axis(axis, a.ndim), int)
  return expand_dims(result, dims) if keepdims else result


@export
def argmin(a: ArrayLike, axis: int | None = None, out: None = None,
           keepdims: bool | None = None) -> Array:
  """Return the index of the minimum value of an array.

  JAX implementation of :func:`numpy.argmin`.

  Args:
    a: input array
    axis: optional integer specifying the axis along which to find the minimum
      value. If ``axis`` is not specified, ``a`` will be flattened.
    out: unused by JAX
    keepdims: if True, then return an array with the same number of dimensions

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Guard with a size check: if a.size == 0 (or a.shape[axis] == 0) handle the empty case explicitly
2. Fix upstream filtering so the array cannot be empty
3. Use a sentinel: run argmax on a padded array (e.g. append -inf) and ignore the sentinel index

Example fix

// before
idx = jnp.argmax(masked)  # may be empty
// after
idx = jnp.argmax(masked) if masked.size else -1

if a.size == 0: return -1  # explicit empty handling
return int(jnp.argmax(a))

def nonempty_for_argmax(a, axis=None):
    a = jnp.asarray(a)
    return a.size > 0 if axis is None else a.shape[axis] > 0

try:
    idx = jnp.argmax(a)
except ValueError as e:
    if 'empty sequence' in str(e):
        idx = -1
    else:
        raise

When it happens

Trigger: jnp.argmax(jnp.array([])); jnp.argmax(a, axis=k) where a.shape[k] == 0; a filtered/computed array that became empty at runtime under jit.

Common situations: Data-dependent filtering (a[a > thresh]) leaving zero elements; empty batches in a training loop; edge-case shapes in tests.

Related errors

• argmin and argmax require non-empty reduced dimension. opera
• Invalid axis {axis} for operand shape {operand.shape}
• index is out of bounds for axis {axis} with size 0
• Cannot do a non-empty jnp.take() from an empty axis.
• The 'out' argument to jnp.argmax is not supported.

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