jax-ml/jax · error · TypeError

top_k operand must have >= 1 dimension, got {}

Error message

top_k operand must have >= 1 dimension, got {}

What it means

top_k requires an operand with at least one dimension; a scalar (0-d array) has no axis to rank along. This is a TypeError from shape evaluation.

Source

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

                          avals_in=util.flatten(zip(scalar_avals, scalar_avals)),
                          avals_out=[core.ShapedArray((), np.bool_)])

    out = lower_comparator(sub_ctx, *comparator.arguments, num_keys=num_keys)
    flat_out, _ = mlir.ir_tree_registry.flatten(out)
    hlo.return_(flat_out)
  return [mlir.lower_with_sharding_in_types(ctx, op, aval)
          for op, aval in zip(sort.results, ctx.avals_out)]

mlir.register_lowering(sort_p, _sort_lower)


def _top_k_abstract_eval(operand, *, k, axis, is_stable):
  if dtypes.issubdtype(operand.dtype, np.complexfloating):
    raise ValueError("top_k is not compatible with complex inputs.")
  if k < 0:
    raise ValueError(f"k argument to top_k must be nonnegative, got {k}")
  if len(operand.shape) == 0:
    raise TypeError("top_k operand must have >= 1 dimension, got {}"
                    .format(operand.shape))
  if not (0 <= axis < len(operand.shape)):
    raise ValueError(f"axis argument out of range: {axis=} for {operand.shape=}")
  shape = list(operand.shape)
  if shape[axis] < k:
    raise ValueError("k argument to top_k must be no larger than size along axis;"
                     f" got {k=} with {shape=} and {axis=}")
  int32_max = dtypes.iinfo('int32').max
  try:
    too_large = (shape[axis] > int32_max + 1)
  except core.InconclusiveDimensionOperation:
    pass
  else:
    if too_large:
      raise ValueError(
          'top_k returns int32 indices, which will overflow for array'
          f' dimensions larger than the maximum int32 ({int32_max}). Got'
          f' {operand.shape=}')

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Keep at least one dimension: use x.sum(axis=1) or keepdims=True before top_k.
2. Reshape scalars: x.reshape(1) then top_k(x, 1).
3. Check for stray squeezes/drop_axis in vmap that reduce rank to 0.

Example fix

# before
vals, idx = jnp.top_k(scores.sum(), k)  # scalar
# after
vals, idx = jnp.top_k(scores.sum(axis=-1), k)

assert x.ndim >= 1, x.shape
if x.ndim == 0:
    x = x.reshape(1)
vals, idx = jnp.top_k(x, k)

def has_rank_at_least(x, n):
    return x.ndim >= n

When it happens

Trigger: jnp.top_k(jnp.asarray(3.0), k=1), or top_k applied after an operation that collapses all dims (e.g., x.sum() or x.mean() producing a scalar).

Common situations: Per-example scores reduced to scalars before ranking instead of after; a vmap'd function where the mapped axis was squeezed away; batch-of-one reshapes to ().

Related errors

• unstack requires arrays with rank > 0, however a scalar arra
• iteration over a 0-d array
• Invalid scalar value {x}
• cond_fun must return a boolean scalar, but got output type(s
• length of padding_config must equal the number of axes of op

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