jax-ml/jax · error · NotImplementedError

jnp.partition for complex dtype is not implemented.

Error message

jnp.partition for complex dtype is not implemented.

What it means

jnp.partition is not implemented for complex-valued inputs because XLA's partitioning/TopK primitives have no complex ordering defined. JAX raises NotImplementedError (there is a TODO in source to also handle NaN like NumPy) rather than silently producing wrong results. Complex numbers lack a total order, so partial-selection sort semantics are undefined.

Source

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

    Array([1, 2, 3, 3, 4, 9, 8, 7, 6, 5], dtype=int32)

    The result is a partially-sorted copy of the input. All values before ``kth``
    are of smaller than the pivot value, and all values after ``kth`` are larger
    than the pivot value:

    >>> 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 [9 8 7 6 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.partition 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 = -lax.top_k(-(arr + 1), kth + 1)[0] - 1
  else:
    bottom = -lax.top_k(-arr, kth + 1)[0]
  top = lax.top_k(arr, arr.shape[-1] - kth - 1)[0]
  out = lax.concatenate([bottom, top], dimension=arr.ndim - 1)
  return out.swapaxes(-1, axis)


@export
@api.jit(static_argnames=['kth', 'axis'])
def argpartition(a: ArrayLike, kth: int, axis: int = -1) -> Array:
  """Returns indices that partially sort an array.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Compute a real-valued key (e.g. jnp.abs(arr)) and partition that instead
2. Use jnp.top_k on a real projection such as magnitude or real part
3. Sort by real part: partition on arr.real if that matches your intent
4. Do the partition in NumPy on host if complex semantics are required

Example fix

// before
part = jnp.partition(cplx_arr, kth=3)
// after
mags = jnp.abs(cplx_arr)
order = jnp.argsort(mags)  # or partition mags and index cplx_arr

if jnp.issubdtype(arr.dtype, jnp.complexfloating): raise ValueError('partition complex inputs manually via jnp.abs key')

def is_real_for_partition(a) -> bool:
    return not jnp.issubdtype(np.asarray(a).dtype, np.complexfloating)

try:
    out = jnp.partition(arr, k)
except NotImplementedError:
    out = arr[jnp.argsort(jnp.abs(arr))]  # fallback ordering

When it happens

Trigger: Calling jnp.partition(complex_array, kth) where the array's dtype is complex64/complex128; e.g. jnp.partition(jnp.array([1+2j, 3-1j]), 1).

Common situations: Signal processing or quantum-computing pipelines with complex tensors that call NumPy-style partition for top-k selection; porting np.partition code from a DSP workflow.

Related errors

• jnp.argpartition 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/a324a1ea65fbc280. Report an issue: GitHub.