{"record":{"id":"a324a1ea65fbc280","repo":"jax-ml/jax","slug":"jnp-partition-for-complex-dtype-is-not-implemented","errorCode":null,"errorMessage":"jnp.partition for complex dtype is not implemented.","messagePattern":"jnp\\.partition for complex dtype is not implemented\\.","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/numpy/sorting.py","lineNumber":235,"sourceCode":"    Array([1, 2, 3, 3, 4, 9, 8, 7, 6, 5], dtype=int32)\n\n    The result is a partially-sorted copy of the input. All values before ``kth``\n    are of smaller than the pivot value, and all values after ``kth`` are larger\n    than the pivot value:\n\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 [9 8 7 6 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.partition 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 = -lax.top_k(-(arr + 1), kth + 1)[0] - 1\n  else:\n    bottom = -lax.top_k(-arr, kth + 1)[0]\n  top = lax.top_k(arr, arr.shape[-1] - kth - 1)[0]\n  out = lax.concatenate([bottom, top], dimension=arr.ndim - 1)\n  return out.swapaxes(-1, axis)\n\n\n@export\n@api.jit(static_argnames=['kth', 'axis'])\ndef argpartition(a: ArrayLike, kth: int, axis: int = -1) -> Array:\n  \"\"\"Returns indices that partially sort an array.","sourceCodeStart":217,"sourceCodeEnd":253,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/sorting.py#L217-L253","documentation":"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.","triggerScenarios":"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).","commonSituations":"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.","solutions":["Compute a real-valued key (e.g. jnp.abs(arr)) and partition that instead","Use jnp.top_k on a real projection such as magnitude or real part","Sort by real part: partition on arr.real if that matches your intent","Do the partition in NumPy on host if complex semantics are required"],"exampleFix":"// before\npart = jnp.partition(cplx_arr, kth=3)\n// after\nmags = jnp.abs(cplx_arr)\norder = jnp.argsort(mags)  # or partition mags and index cplx_arr","handlingStrategy":"type-guard","validationCode":"if jnp.issubdtype(arr.dtype, jnp.complexfloating): raise ValueError('partition complex inputs manually via jnp.abs key')","typeGuard":"def is_real_for_partition(a) -> bool:\n    return not jnp.issubdtype(np.asarray(a).dtype, np.complexfloating)","tryCatchPattern":"try:\n    out = jnp.partition(arr, k)\nexcept NotImplementedError:\n    out = arr[jnp.argsort(jnp.abs(arr))]  # fallback ordering","preventionTips":["Check dtype before selection ops in complex pipelines","Centralize top-k/partition in a helper that handles complex","Prefer jnp.top_k(jnp.abs(x), k) for magnitude ranking"],"tags":["jax","partition","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"}