jax-ml/jax · error · ValueError

top_k is not compatible with complex inputs.

Error message

top_k is not compatible with complex inputs.

What it means

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.

Source

Thrown at jax/_src/numpy/sorting.py:507

    Array([[5, 4],
           [5, 4]], dtype=int32)
    >>> indices
    Array([[4, 3],
           [0, 1]], dtype=int32)

    Find the two smallest elements along the first axis:

    >>> values, indices = jnp.top_k(a, 2, axis=0, mode='smallest')
    >>> values
    Array([[1, 2, 3, 2, 1],
           [5, 4, 3, 4, 5]], dtype=int32)
    >>> indices
    Array([[0, 0, 0, 1, 1],
           [1, 1, 1, 0, 0]], dtype=int32)
  """
  arr = util.ensure_arraylike("top_k", a)
  if dtypes.issubdtype(arr.dtype, np.complexfloating):
    raise ValueError("top_k is not compatible with complex inputs.")
  if mode not in ("largest", "smallest"):
    raise ValueError(f"mode must be 'largest' or 'smallest', got {mode!r}")
  axis = canonicalize_axis(axis, arr.ndim)
  if mode == "largest":
    return lax.top_k(arr, k, axis=axis)
  elif dtypes.isdtype(arr.dtype, "bool"):
    inv = lax.bitwise_not(arr)
    vals, indices = lax.top_k(inv, k, axis=axis)
    return lax.bitwise_not(vals), indices
  elif dtypes.isdtype(arr.dtype, "unsigned integer"):
    inv = -(arr + 1)
    vals, indices = lax.top_k(inv, k, axis=axis)
    return -(vals + 1), indices
  else:
    inv = -arr
    vals, indices = lax.top_k(inv, k, axis=axis)
    return -vals, indices

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Apply top_k to a real projection: jnp.top_k(jnp.abs(arr), k) for magnitude ranking
2. Use arr.real or arr.imag if that matches the intended ordering
3. For complex-aware ordering, sort via jnp.argsort on a composite real key
4. Do complex top-k on host with NumPy if device execution is not required

Example fix

# before
vals, idx = jnp.top_k(cplx_signal, 10)
# after
mags, idx = jnp.top_k(jnp.abs(cplx_signal), 10)
vals = cplx_signal[idx]

if jnp.issubdtype(arr.dtype, jnp.complexfloating):
    key = jnp.abs(arr)
else:
    key = arr
vals_or_idx = jnp.top_k(key, k)

def top_k_safe(a, k):
    a = jnp.abs(a) if jnp.issubdtype(np.asarray(a).dtype, np.complexfloating) else a
    return jnp.top_k(a, k)

try:
    v, i = jnp.top_k(arr, k)
except ValueError:
    v, i = jnp.top_k(jnp.abs(arr), k)

When it happens

Trigger: 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).

Common situations: 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.

Related errors

• jnp.partition for complex dtype is not implemented.
• jnp.argpartition for complex dtype is not implemented.
• mode must be 'largest' or 'smallest', got {mode!r}
• repeated axis in lax.expand_dims: {dimensions}
• Invalid dimension range passed to collapse: {operand.shape}[

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/4b1d2754b0a60c17. Report an issue: GitHub.