{"record":{"id":"17ee7b4d97312c70","repo":"jax-ml/jax","slug":"top-k-returns-int32-indices-which-will-overflow-f","errorCode":null,"errorMessage":"top_k returns int32 indices, which will overflow for array dimensions larger than the maximum int32 ({int32_max}). Got {operand.shape=}","messagePattern":"top_k returns int32 indices, which will overflow for array dimensions larger than the maximum int32 \\((.+?)\\)\\. Got (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":9048,"sourceCode":"  if k < 0:\n    raise ValueError(f\"k argument to top_k must be nonnegative, got {k}\")\n  if len(operand.shape) == 0:\n    raise TypeError(\"top_k operand must have >= 1 dimension, got {}\"\n                    .format(operand.shape))\n  if not (0 <= axis < len(operand.shape)):\n    raise ValueError(f\"axis argument out of range: {axis=} for {operand.shape=}\")\n  shape = list(operand.shape)\n  if shape[axis] < k:\n    raise ValueError(\"k argument to top_k must be no larger than size along axis;\"\n                     f\" got {k=} with {shape=} and {axis=}\")\n  int32_max = dtypes.iinfo('int32').max\n  try:\n    too_large = (shape[axis] > int32_max + 1)\n  except core.InconclusiveDimensionOperation:\n    pass\n  else:\n    if too_large:\n      raise ValueError(\n          'top_k returns int32 indices, which will overflow for array'\n          f' dimensions larger than the maximum int32 ({int32_max}). Got'\n          f' {operand.shape=}')\n  shape[axis] = k\n  if operand.sharding.spec[axis] is not None:\n    raise core.ShardingTypeError(\n        'The input should be unsharded over the axis along which to compute the'\n        f' top_k values. Got input type={operand} and axis={axis}')\n  return (operand.update(shape=shape),\n          operand.update(shape=shape, dtype=np.dtype(np.int32)))\n\ndef _top_k_jvp(primals, tangents, *, k, axis, is_stable):\n  operand, = primals\n  tangent, = tangents\n  primals_out = top_k(operand, k, axis=axis, is_stable=is_stable)\n  if type(tangent) is ad_util.Zero:\n    tangent_out = ad_util.p2tz(primals_out[0])\n  else:","sourceCodeStart":9030,"sourceCodeEnd":9066,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L9030-L9066","documentation":"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.","triggerScenarios":"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).","commonSituations":"Very large embedding tables or datasets sharded across TPU pods; symbolic shape arithmetic under jit that concludes the dimension is enormous.","solutions":["Split the axis into chunks smaller than 2**31, top_k each chunk, then merge candidates.","Use a partition/argpartition-style approach or gather-based selection that returns int64.","If the dimension estimate is wrong (symbolic shapes), concretize or reshape so the axis bound is accurate."],"exampleFix":"# before\nvals, idx = lax.top_k(huge_1d, k)  # huge_1d.size > 2**31\n# after\nchunks = huge_1d.reshape(-1, 2**30)\nv, i = jax.vmap(lambda c: lax.top_k(c, k))(chunks)\n# then merge chunk results","handlingStrategy":"validation","validationCode":"import jax.numpy as jnp\nint32_max = jnp.iinfo(jnp.int32).max\nsize = x.shape[axis]\nif size > int32_max + 1:\n    x = x.reshape(-1, int32_max)  # chunk before top_k\nvals, idx = jnp.top_k(x, k, axis=-1)","typeGuard":"def axis_within_int32(x, axis):\n    return x.shape[axis] <= jnp.iinfo(jnp.int32).max + 1","tryCatchPattern":null,"preventionTips":["Chunk huge axes below 2**31 and merge results.","Use int64-producing selection (argpartition/gather) for giant dims."],"tags":["jax","top-k","int32-overflow","large-arrays"],"backgroundTag":"integer-overflow","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}