jax-ml/jax · error · ValueError

k argument to top_k must be no larger than size along axis;

Error message

k argument to top_k must be no larger than size along axis; got {k=} with {shape=} and {axis=}

What it means

top_k cannot return more elements than exist along the chosen axis; k must satisfy k <= shape[axis]. The evaluator checks the static shape, so this fires even before any data exists.

Source

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

  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=}')
  shape[axis] = k
  if operand.sharding.spec[axis] is not None:
    raise core.ShardingTypeError(
        'The input should be unsharded over the axis along which to compute the'
        f' top_k values. Got input type={operand} and axis={axis}')
  return (operand.update(shape=shape),

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Clamp k to the axis size: k = min(k, x.shape[axis]).
2. For variable-length inputs, either pad to a minimum length or gather per-example top-k with a mask.
3. Check off-by-one: use k = size, not size + 1, when you want everything.

Example fix

# before
vals, idx = jnp.top_k(scores, k=100)  # scores: (batch, 50)
# after
vals, idx = jnp.top_k(scores, k=min(100, scores.shape[-1]))

k = min(int(k), x.shape[axis])
vals, idx = jnp.top_k(x, k, axis=axis)

def k_fits(x, k, axis=-1):
    return 0 <= k <= x.shape[axis]

When it happens

Trigger: jnp.top_k(x, k=10) where x.shape[axis] == 5; k hardcoded larger than a variable-length axis; k = x.shape[axis] + 1 off-by-one.

Common situations: Fixed k (e.g., top-100) applied to short sequences or small batches; sequence-length-dependent data where some examples are shorter than k; unit tests using tiny arrays.

Related errors

• 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
• conv_general_dilated window and window_strides must have the
• Wrong number of explicit pads for convolution: expected {},

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