{"record":{"id":"592d23278a57b183","repo":"jax-ml/jax","slug":"k-argument-to-top-k-must-be-no-larger-than-size-al","errorCode":null,"errorMessage":"k argument to top_k must be no larger than size along axis; got {k=} with {shape=} and {axis=}","messagePattern":"k argument to top_k must be no larger than size along axis; got (.+?) with (.+?) and (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":9039,"sourceCode":"  return [mlir.lower_with_sharding_in_types(ctx, op, aval)\n          for op, aval in zip(sort.results, ctx.avals_out)]\n\nmlir.register_lowering(sort_p, _sort_lower)\n\n\ndef _top_k_abstract_eval(operand, *, k, axis, is_stable):\n  if dtypes.issubdtype(operand.dtype, np.complexfloating):\n    raise ValueError(\"top_k is not compatible with complex inputs.\")\n  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),","sourceCodeStart":9021,"sourceCodeEnd":9057,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L9021-L9057","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Clamp k to the axis size: k = min(k, x.shape[axis]).","For variable-length inputs, either pad to a minimum length or gather per-example top-k with a mask.","Check off-by-one: use k = size, not size + 1, when you want everything."],"exampleFix":"# before\nvals, idx = jnp.top_k(scores, k=100)  # scores: (batch, 50)\n# after\nvals, idx = jnp.top_k(scores, k=min(100, scores.shape[-1]))","handlingStrategy":"validation","validationCode":"k = min(int(k), x.shape[axis])\nvals, idx = jnp.top_k(x, k, axis=axis)","typeGuard":"def k_fits(x, k, axis=-1):\n    return 0 <= k <= x.shape[axis]","tryCatchPattern":null,"preventionTips":["Clamp k with min(k, size) for variable-length inputs.","Pad short sequences before batched top-k."],"tags":["jax","top-k","shape-validation","off-by-one"],"backgroundTag":"k-exceeds-array-size","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}