jax-ml/jax · error · NotImplementedError

jnp.argpartition for complex dtype is not implemented.

Error message

jnp.argpartition for complex dtype is not implemented.

What it means

jnp.argpartition raises NotImplementedError for complex dtypes for the same reason as jnp.partition: XLA has no ordering over complex numbers, so selecting the k-th smallest indices is undefined. The check happens immediately after input validation, before any axis handling.

Source

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

    The result is a sequence of indices that partially sort the input. All indices
    before ``kth`` are of values smaller than the pivot value, and all indices
    after ``kth`` are of values larger than the pivot value:

    >>> x_partitioned = x[idx]
    >>> smallest_values = x_partitioned[:kth]
    >>> pivot_value = x_partitioned[kth]
    >>> largest_values = x_partitioned[kth + 1:]
    >>> print(smallest_values, pivot_value, largest_values)
    [1 2 3 3] 4 [6 8 9 7 5]

    Notice that among ``smallest_values`` and ``largest_values``, the returned
    order is arbitrary and implementation-dependent.
  """
  # TODO(jakevdp): handle NaN values like numpy.
  arr = util.ensure_arraylike("partition", a)
  if dtypes.issubdtype(arr.dtype, np.complexfloating):
    raise NotImplementedError("jnp.argpartition for complex dtype is not implemented.")
  axis = canonicalize_axis(axis, arr.ndim)
  kth = canonicalize_axis(kth, arr.shape[axis])

  arr = arr.swapaxes(axis, -1)
  if dtypes.isdtype(arr.dtype, "unsigned integer"):
    # Here, we apply a trick to handle correctly 0 values for unsigned integers
    bottom_ind = lax.top_k(-(arr + 1), kth + 1)[1]
  else:
    bottom_ind = lax.top_k(-arr, kth + 1)[1]

  # To avoid issues with duplicate values, we compute the top indices via a proxy
  set_to_zero = lambda a, i: a.at[i].set(0)
  for _ in range(arr.ndim - 1):
    set_to_zero = api.vmap(set_to_zero)
  proxy = set_to_zero(lax.full(arr.shape, 1.0), bottom_ind)
  top_ind = lax.top_k(proxy, arr.shape[-1] - kth - 1)[1]
  out = lax.concatenate([bottom_ind, top_ind], dimension=arr.ndim - 1)
  return out.swapaxes(-1, axis)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Partition a real key derived from the array: idx = jnp.argsort(jnp.abs(cplx)) and slice the k smallest
2. Use jnp.top_k(jnp.abs(cplx), k) to get values and indices directly
3. If ordering by real part is acceptable, run argpartition on arr.real
4. Move data to NumPy and use np.argpartition, then convert indices back with jnp.asarray

Example fix

// before
idx = jnp.argpartition(fft_out, -5)[-5:]
// after
vals, idx = jnp.top_k(jnp.abs(fft_out), 5)

if jnp.issubdtype(arr.dtype, jnp.complexfloating):
    idx = jnp.argsort(jnp.abs(arr))  # real-key fallback path

def real_or_key(a):
    return jnp.abs(a) if jnp.issubdtype(np.asarray(a).dtype, np.complexfloating) else a

try:
    idx = jnp.argpartition(arr, k)
except NotImplementedError:
    idx = jnp.argsort(jnp.abs(arr))

When it happens

Trigger: Calling jnp.argpartition(a, kth) where a is complex64 or complex128; e.g. jnp.argpartition(jnp.fft.fft(x), 2) on FFT output.

Common situations: Post-FFT processing where top-frequency selection is done with argpartition; porting NumPy spectral analysis code; RF/audio pipelines using complex baseband samples.

Related errors

• jnp.partition for complex dtype is not implemented.
• top_k is not compatible with complex inputs.
• Unsupported dtype {x_aval.dtype}
• gaussian_kde does not support complex coordinates
• array ref with memory space only works inside of a `jit`.

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