{"record":{"id":"24c717833fc714cd","repo":"jax-ml/jax","slug":"top-k-is-not-compatible-with-complex-inputs","errorCode":null,"errorMessage":"top_k is not compatible with complex inputs.","messagePattern":"top_k is not compatible with complex inputs\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":9029,"sourceCode":"  with ir.InsertionPoint(comparator):\n    lower_comparator = mlir.lower_fun(partial(_sort_lt_comparator),\n                                      multiple_results=False)\n    sub_ctx = ctx.replace(primitive=None,\n                          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:","sourceCodeStart":9011,"sourceCodeEnd":9047,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L9011-L9047","documentation":"top_k does not support complex dtypes because 'largest' is not a total order over complex numbers. The abstract evaluator rejects complexfloating operands immediately.","triggerScenarios":"lax.top_k(complex_array, k) or jnp.top_k on a complex-valued array; also when a real array is implicitly promoted to complex by an upstream operation (e.g., FFT output).","commonSituations":"Taking top-k of FFT magnitudes but forgetting the abs(); complex weights in signal-processing models; any pipeline ending in jnp.fft.* followed by ranking.","solutions":["Rank by magnitude or real part: jnp.top_k(jnp.abs(x), k) (add angle as a tiebreak if needed).","If you need a complex-ordering, define a key (abs, then real/imag) and use lax.sort with a custom comparator.","Prevent implicit promotion to complex upstream (check x.dtype before top_k)."],"exampleFix":"# before\nvals, idx = jnp.top_k(jnp.fft.rfft(signal), k)\n# after\nspectrum = jnp.fft.rfft(signal)\nvals, idx = jnp.top_k(jnp.abs(spectrum), k)","handlingStrategy":"type-guard","validationCode":"import jax.numpy as jnp, numpy as np\nfrom jax import dtypes\nif dtypes.issubdtype(x.dtype, np.complexfloating):\n    x = jnp.abs(x)\nvals, idx = jnp.top_k(x, k)","typeGuard":"def is_complex(x):\n    import numpy as np\n    from jax import dtypes\n    return dtypes.issubdtype(x.dtype, np.complexfloating)","tryCatchPattern":null,"preventionTips":["Apply abs() to FFT outputs before ranking.","Define an explicit key (magnitude, angle) for complex ordering."],"tags":["jax","top-k","complex-numbers","dtype"],"backgroundTag":"unsupported-dtype-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}