{"record":{"id":"c7140a2cf6979d4d","repo":"jax-ml/jax","slug":"top-k-operand-must-have-1-dimension-got","errorCode":null,"errorMessage":"top_k operand must have >= 1 dimension, got {}","messagePattern":"top_k operand must have >= 1 dimension, got (.+?)","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":9033,"sourceCode":"                          avals_in=util.flatten(zip(scalar_avals, scalar_avals)),\n                          avals_out=[core.ShapedArray((), np.bool_)])\n\n    out = lower_comparator(sub_ctx, *comparator.arguments, num_keys=num_keys)\n    flat_out, _ = mlir.ir_tree_registry.flatten(out)\n    hlo.return_(flat_out)\n  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=}')","sourceCodeStart":9015,"sourceCodeEnd":9051,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L9015-L9051","documentation":"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.","triggerScenarios":"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).","commonSituations":"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 ().","solutions":["Keep at least one dimension: use x.sum(axis=1) or keepdims=True before top_k.","Reshape scalars: x.reshape(1) then top_k(x, 1).","Check for stray squeezes/drop_axis in vmap that reduce rank to 0."],"exampleFix":"# before\nvals, idx = jnp.top_k(scores.sum(), k)  # scalar\n# after\nvals, idx = jnp.top_k(scores.sum(axis=-1), k)","handlingStrategy":"validation","validationCode":"assert x.ndim >= 1, x.shape\nif x.ndim == 0:\n    x = x.reshape(1)\nvals, idx = jnp.top_k(x, k)","typeGuard":"def has_rank_at_least(x, n):\n    return x.ndim >= n","tryCatchPattern":null,"preventionTips":["Use keepdims=True or axis-restricted reductions to preserve rank.","Watch vmap out_axes that squeeze mapped dims to scalars."],"tags":["jax","top-k","scalar","rank-error"],"backgroundTag":"zero-dimensional-array-misuse","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}