jax-ml/jax · error · ValueError

top_k returns int32 indices, which will overflow for array d

Error message

top_k returns int32 indices, which will overflow for array dimensions larger than the maximum int32 ({int32_max}). Got {operand.shape=}

What it means

top_k returns int32 indices; if the reduced axis is larger than int32 max + 1 the indices could overflow. When the static shape provably exceeds this bound, JAX raises rather than silently produce garbage indices.

Source

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

  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),
          operand.update(shape=shape, dtype=np.dtype(np.int32)))

def _top_k_jvp(primals, tangents, *, k, axis, is_stable):
  operand, = primals
  tangent, = tangents
  primals_out = top_k(operand, k, axis=axis, is_stable=is_stable)
  if type(tangent) is ad_util.Zero:
    tangent_out = ad_util.p2tz(primals_out[0])
  else:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Split the axis into chunks smaller than 2**31, top_k each chunk, then merge candidates.
2. Use a partition/argpartition-style approach or gather-based selection that returns int64.
3. If the dimension estimate is wrong (symbolic shapes), concretize or reshape so the axis bound is accurate.

Example fix

# before
vals, idx = lax.top_k(huge_1d, k)  # huge_1d.size > 2**31
# after
chunks = huge_1d.reshape(-1, 2**30)
v, i = jax.vmap(lambda c: lax.top_k(c, k))(chunks)
# then merge chunk results

import jax.numpy as jnp
int32_max = jnp.iinfo(jnp.int32).max
size = x.shape[axis]
if size > int32_max + 1:
    x = x.reshape(-1, int32_max)  # chunk before top_k
vals, idx = jnp.top_k(x, k, axis=-1)

def axis_within_int32(x, axis):
    return x.shape[axis] <= jnp.iinfo(jnp.int32).max + 1

When it happens

Trigger: lax.top_k on an array whose axis dimension is statically larger than 2**31 (only feasible on accelerators with huge memory or with symbolic dimensions that resolve large).

Common situations: Very large embedding tables or datasets sharded across TPU pods; symbolic shape arithmetic under jit that concludes the dimension is enormous.

Related errors

• top_k is not compatible with complex inputs.
• top_k operand must have >= 1 dimension, got {}
• axis argument out of range: {axis=} for {operand.shape=}
• k argument to top_k must be no larger than size along axis;
• top_k is not compatible with complex inputs.

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