{"record":{"id":"3516312512e2adb6","repo":"jax-ml/jax","slug":"jnp-argpartition-for-complex-dtype-is-not-implemen","errorCode":null,"errorMessage":"jnp.argpartition for complex dtype is not implemented.","messagePattern":"jnp\\.argpartition for complex dtype is not implemented\\.","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/numpy/sorting.py","lineNumber":306,"sourceCode":"\n    The result is a sequence of indices that partially sort the input. All indices\n    before ``kth`` are of values smaller than the pivot value, and all indices\n    after ``kth`` are of values larger than the pivot value:\n\n    >>> x_partitioned = x[idx]\n    >>> smallest_values = x_partitioned[:kth]\n    >>> pivot_value = x_partitioned[kth]\n    >>> largest_values = x_partitioned[kth + 1:]\n    >>> print(smallest_values, pivot_value, largest_values)\n    [1 2 3 3] 4 [6 8 9 7 5]\n\n    Notice that among ``smallest_values`` and ``largest_values``, the returned\n    order is arbitrary and implementation-dependent.\n  \"\"\"\n  # TODO(jakevdp): handle NaN values like numpy.\n  arr = util.ensure_arraylike(\"partition\", a)\n  if dtypes.issubdtype(arr.dtype, np.complexfloating):\n    raise NotImplementedError(\"jnp.argpartition for complex dtype is not implemented.\")\n  axis = canonicalize_axis(axis, arr.ndim)\n  kth = canonicalize_axis(kth, arr.shape[axis])\n\n  arr = arr.swapaxes(axis, -1)\n  if dtypes.isdtype(arr.dtype, \"unsigned integer\"):\n    # Here, we apply a trick to handle correctly 0 values for unsigned integers\n    bottom_ind = lax.top_k(-(arr + 1), kth + 1)[1]\n  else:\n    bottom_ind = lax.top_k(-arr, kth + 1)[1]\n\n  # To avoid issues with duplicate values, we compute the top indices via a proxy\n  set_to_zero = lambda a, i: a.at[i].set(0)\n  for _ in range(arr.ndim - 1):\n    set_to_zero = api.vmap(set_to_zero)\n  proxy = set_to_zero(lax.full(arr.shape, 1.0), bottom_ind)\n  top_ind = lax.top_k(proxy, arr.shape[-1] - kth - 1)[1]\n  out = lax.concatenate([bottom_ind, top_ind], dimension=arr.ndim - 1)\n  return out.swapaxes(-1, axis)","sourceCodeStart":288,"sourceCodeEnd":324,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/sorting.py#L288-L324","documentation":"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.","triggerScenarios":"Calling jnp.argpartition(a, kth) where a is complex64 or complex128; e.g. jnp.argpartition(jnp.fft.fft(x), 2) on FFT output.","commonSituations":"Post-FFT processing where top-frequency selection is done with argpartition; porting NumPy spectral analysis code; RF/audio pipelines using complex baseband samples.","solutions":["Partition a real key derived from the array: idx = jnp.argsort(jnp.abs(cplx)) and slice the k smallest","Use jnp.top_k(jnp.abs(cplx), k) to get values and indices directly","If ordering by real part is acceptable, run argpartition on arr.real","Move data to NumPy and use np.argpartition, then convert indices back with jnp.asarray"],"exampleFix":"// before\nidx = jnp.argpartition(fft_out, -5)[-5:]\n// after\nvals, idx = jnp.top_k(jnp.abs(fft_out), 5)","handlingStrategy":"type-guard","validationCode":"if jnp.issubdtype(arr.dtype, jnp.complexfloating):\n    idx = jnp.argsort(jnp.abs(arr))  # real-key fallback path","typeGuard":"def real_or_key(a):\n    return jnp.abs(a) if jnp.issubdtype(np.asarray(a).dtype, np.complexfloating) else a","tryCatchPattern":"try:\n    idx = jnp.argpartition(arr, k)\nexcept NotImplementedError:\n    idx = jnp.argsort(jnp.abs(arr))","preventionTips":["Wrap FFT outputs with real projections before selection","Unit-test selection helpers with complex inputs","Document ordering semantics for complex data"],"tags":["jax","argpartition","complex-dtype","not-implemented","sorting"],"backgroundTag":"complex-dtype-unsupported-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}