{"record":{"id":"4b1d2754b0a60c17","repo":"jax-ml/jax","slug":"top-k-is-not-compatible-with-complex-inputs-4b1d27","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/numpy/sorting.py","lineNumber":507,"sourceCode":"    Array([[5, 4],\n           [5, 4]], dtype=int32)\n    >>> indices\n    Array([[4, 3],\n           [0, 1]], dtype=int32)\n\n    Find the two smallest elements along the first axis:\n\n    >>> values, indices = jnp.top_k(a, 2, axis=0, mode='smallest')\n    >>> values\n    Array([[1, 2, 3, 2, 1],\n           [5, 4, 3, 4, 5]], dtype=int32)\n    >>> indices\n    Array([[0, 0, 0, 1, 1],\n           [1, 1, 1, 0, 0]], dtype=int32)\n  \"\"\"\n  arr = util.ensure_arraylike(\"top_k\", a)\n  if dtypes.issubdtype(arr.dtype, np.complexfloating):\n    raise ValueError(\"top_k is not compatible with complex inputs.\")\n  if mode not in (\"largest\", \"smallest\"):\n    raise ValueError(f\"mode must be 'largest' or 'smallest', got {mode!r}\")\n  axis = canonicalize_axis(axis, arr.ndim)\n  if mode == \"largest\":\n    return lax.top_k(arr, k, axis=axis)\n  elif dtypes.isdtype(arr.dtype, \"bool\"):\n    inv = lax.bitwise_not(arr)\n    vals, indices = lax.top_k(inv, k, axis=axis)\n    return lax.bitwise_not(vals), indices\n  elif dtypes.isdtype(arr.dtype, \"unsigned integer\"):\n    inv = -(arr + 1)\n    vals, indices = lax.top_k(inv, k, axis=axis)\n    return -(vals + 1), indices\n  else:\n    inv = -arr\n    vals, indices = lax.top_k(inv, k, axis=axis)\n    return -vals, indices\n","sourceCodeStart":489,"sourceCodeEnd":525,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/sorting.py#L489-L525","documentation":"jnp.top_k only supports real inputs; complex numbers have no total order, so selecting the k largest elements is undefined. The check runs immediately after array conversion, before mode/axis validation, and raises ValueError.","triggerScenarios":"Calling jnp.top_k(cplx, k) or jnp.partition/argpartition (which route through top_k) on a complex64/complex128 array; e.g. jnp.top_k(jnp.abs(x) * jnp.exp(1j*x), 3).","commonSituations":"Top-k retrieval on FFT output, embeddings represented as complex numbers, or quantum-state amplitudes; also hit indirectly because jnp.partition and argpartition call top_k internally for their selection logic.","solutions":["Apply top_k to a real projection: jnp.top_k(jnp.abs(arr), k) for magnitude ranking","Use arr.real or arr.imag if that matches the intended ordering","For complex-aware ordering, sort via jnp.argsort on a composite real key","Do complex top-k on host with NumPy if device execution is not required"],"exampleFix":"# before\nvals, idx = jnp.top_k(cplx_signal, 10)\n# after\nmags, idx = jnp.top_k(jnp.abs(cplx_signal), 10)\nvals = cplx_signal[idx]","handlingStrategy":"type-guard","validationCode":"if jnp.issubdtype(arr.dtype, jnp.complexfloating):\n    key = jnp.abs(arr)\nelse:\n    key = arr\nvals_or_idx = jnp.top_k(key, k)","typeGuard":"def top_k_safe(a, k):\n    a = jnp.abs(a) if jnp.issubdtype(np.asarray(a).dtype, np.complexfloating) else a\n    return jnp.top_k(a, k)","tryCatchPattern":"try:\n    v, i = jnp.top_k(arr, k)\nexcept ValueError:\n    v, i = jnp.top_k(jnp.abs(arr), k)","preventionTips":["Route all top-k calls through one dtype-aware helper","Decide ordering semantics (abs/real) for complex data early","Add complex-input tests for selection utilities"],"tags":["jax","top-k","complex-dtype","value-error","sorting"],"backgroundTag":"complex-dtype-unsupported-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}