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

top_k does not support complex dtypes because 'largest' is not a total order over complex numbers. The abstract evaluator rejects complexfloating operands immediately.

Source

Thrown at jax/_src/lax/lax.py:9029

  with ir.InsertionPoint(comparator):
    lower_comparator = mlir.lower_fun(partial(_sort_lt_comparator),
                                      multiple_results=False)
    sub_ctx = ctx.replace(primitive=None,
                          avals_in=util.flatten(zip(scalar_avals, scalar_avals)),
                          avals_out=[core.ShapedArray((), np.bool_)])

    out = lower_comparator(sub_ctx, *comparator.arguments, num_keys=num_keys)
    flat_out, _ = mlir.ir_tree_registry.flatten(out)
    hlo.return_(flat_out)
  return [mlir.lower_with_sharding_in_types(ctx, op, aval)
          for op, aval in zip(sort.results, ctx.avals_out)]

mlir.register_lowering(sort_p, _sort_lower)


def _top_k_abstract_eval(operand, *, k, axis, is_stable):
  if dtypes.issubdtype(operand.dtype, np.complexfloating):
    raise ValueError("top_k is not compatible with complex inputs.")
  if k < 0:
    raise ValueError(f"k argument to top_k must be nonnegative, got {k}")
  if len(operand.shape) == 0:
    raise TypeError("top_k operand must have >= 1 dimension, got {}"
                    .format(operand.shape))
  if not (0 <= axis < len(operand.shape)):
    raise ValueError(f"axis argument out of range: {axis=} for {operand.shape=}")
  shape = list(operand.shape)
  if shape[axis] < k:
    raise ValueError("k argument to top_k must be no larger than size along axis;"
                     f" got {k=} with {shape=} and {axis=}")
  int32_max = dtypes.iinfo('int32').max
  try:
    too_large = (shape[axis] > int32_max + 1)
  except core.InconclusiveDimensionOperation:
    pass
  else:
    if too_large:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Rank by magnitude or real part: jnp.top_k(jnp.abs(x), k) (add angle as a tiebreak if needed).
2. If you need a complex-ordering, define a key (abs, then real/imag) and use lax.sort with a custom comparator.
3. Prevent implicit promotion to complex upstream (check x.dtype before top_k).

Example fix

# before
vals, idx = jnp.top_k(jnp.fft.rfft(signal), k)
# after
spectrum = jnp.fft.rfft(signal)
vals, idx = jnp.top_k(jnp.abs(spectrum), k)

import jax.numpy as jnp, numpy as np
from jax import dtypes
if dtypes.issubdtype(x.dtype, np.complexfloating):
    x = jnp.abs(x)
vals, idx = jnp.top_k(x, k)

def is_complex(x):
    import numpy as np
    from jax import dtypes
    return dtypes.issubdtype(x.dtype, np.complexfloating)

When it happens

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

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

Related errors

• jax.numpy.var does not yet support real dtype parameters whe
• Pallas top_k only supports float32 and bfloat16, got {input_
• primal and tangent arguments to jax.jvp do not match; dtypes
• unexpected JAX type (e.g. shape/dtype) for gradient ref pass
• Accumulator aval mismatch: expected {aval}, got {acc.aval}

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