jax-ml/jax · error · ValueError

argmin and argmax require non-empty reduced dimension. opera

Error message

argmin and argmax require non-empty reduced dimension. operand.shape={operand.shape} {axis=}

What it means

argmin/argmax require the reduced dimension to have size >= 1; reducing an empty dimension has no well-defined index. The check is on the static shape known at trace time.

Source

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

  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:

  def __init__(self, value_comparator: Callable[[Any, Any], Any]):
    self._value_comparator = value_comparator

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Check for empty dimensions before calling argmax and handle the empty case separately (return a sentinel or skip).
2. Fix upstream filtering so the batch is non-empty, or pad with a dummy row before reducing.
3. If using vmap/jit, add a size guard on the concrete shape via jax.lax.cond or host-side branching.

Example fix

# before
best = jnp.argmax(scores, axis=0)  # scores.shape[0] == 0
# after
best = jnp.argmax(scores, axis=0) if scores.shape[0] > 0 else -1

if x.shape[axis] == 0:
    best = -1  # or skip
else:
    best = jnp.argmax(x, axis=axis)

def nonempty_axis(x, axis):
    return x.shape[axis] >= 1

When it happens

Trigger: jnp.argmax(x, axis=0) where x.shape[0] == 0 (e.g., an empty batch after filtering, or a zero-length sequence dimension).

Common situations: Data pipelines where a filter/batch can legitimately produce zero rows; sequence models with length-0 sequences; dynamic shapes that statically fold to 0.

Related errors

• attempt to get argmax of an empty sequence
• zero-size arrays not supported in convolutions, got shapes {
• scan got `length` argument of {} which disagrees with leadin
• conv_general_dilated batch_group_count must divide lhs batch
• conv_general_dilated rhs output feature dimension size must

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