fix: float16/float128/complex256 reducers, sort, and argsort no longer raise KeyError (#4259)

[ianna]

Fixes #4259

Problem

Every reducer except ak.count, plus ak.sort and ak.argsort, raise a raw KeyError on float16, float128, and complex256 arrays — the message is just the kernel-table lookup key:

>>> ak.sort(ak.Array(np.array([1.0, 2.0], dtype=np.float16)))
KeyError: ('awkward_sort', <class 'numpy.float16'>, <class 'numpy.float16'>, <class 'numpy.int64'>)
>>> ak.sum(ak.values_astype(ak.Array([[1.0, 2.0]]), np.float16), axis=1)
KeyError: ('awkward_reduce_sum', <class 'numpy.float16'>, ...)

kernel-specification.yml has no float16/float128/complex256 entries, so neither the CPU nor CUDA backend has kernels for them. Unlike datetime64/timedelta64 (viewed as int64) and complex128/complex64 (viewed as float64/float32), these dtypes are neither mapped onto an existing kernel nor refused with a clear error — and float16 in particular reaches this failure without the user choosing the dtype (ak.from_arrow maps Arrow halffloat to it, so Arrow/Parquet input hits it).

The existing _dtype_for_kernel mapping uses .view() (a bit-reinterpret), which can't help here: float16 → float32 is a real numeric conversion, not a reinterpret.

Fix

Cast the unsupported dtype to the nearest supported one at the three NumpyArray kernel-dispatch sites (_reduce_next, _sort_next, _argsort_next), and cast value-preserving results back:

• float16 → float32 (lossless)
• float128/longdouble → float64, complex256/clongdouble → complex128 (these lose precision — there is no native kernel for them; detected by kind/itemsize, so platform-independent).

_reduce_next casts self, delegates, and restores result leaves whose dtype equals the cast target — so sum/prod/min/max return the original dtype while argmin/argmax/count_nonzero/any/all keep their int/bool result. _sort_next casts the sorted values back (order-preserving/monotonic conversion). _argsort_next needs no cast-back (positions are unchanged). The hook is backend-agnostic and uses self._backend.nplike, so it also fixes float16 on the cuda backend (the cast runs on-device; float128/complex256 are CPU-only NumPy dtypes and never reach cupy).

Complex arrays have no total order, so sort/argsort on any complex width (64/128/256) now raise a clear TypeError ("complex numbers have no total order") instead of the confusing KeyError — previously plain complex128 already KeyErrord here.

Result

	before	after
reducers (sum, prod, any, all, min, max, argmin, argmax, count_nonzero) on float16/float128/complex256	KeyError	work; match NumPy; value dtype preserved
derived (mean, std, var, ptp, moment, corr, covar, linear_fit, nan*)	KeyError	work
sort/argsort on float16/float128	KeyError	work; sort preserves dtype
sort/argsort on complex (any width)	KeyError	clear TypeError

Tests

• tests/test_4259_float16_float128_complex256_reducers.py — every reducer + sort/argsort on float16 (axis None/1) vs NumPy, dtype preservation, mean/std/var, the issue's exact repro, complex sort TypeError, and platform-gated float128/complex256 coverage.
• tests-cuda/test_4259_cuda_float16_reducers.py — float16 reducers + sort/argsort GPU vs CPU, device residency, complex sort TypeError.

Notes / follow-ups

• Precision: float128/complex256 reductions accumulate in float64/complex128, and min/max return the float64-rounded value — there is no native kernel for extended precision (per the Complex numbers (in ak.from_iter and elsewhere) #392 roadmap). Sorting is still correct (float64 rounding is monotonic).
• GPU float16 copies: the cast materializes one transient float32 device buffer per float16 reduce. Fusing the widen into the cuda reductions (a TransformIterator float16→float32 with a float32 accumulator, like the sum-of-squares path) would make it copy-free — a focused, GPU-tested follow-up, since most cuda reduce impls (min/max/argmin/…) bind h_init to the input dtype and can't take a widened iterator uniformly.

AI disclosure

Parts of this PR were developed with AI assistance (per CONTRIBUTING.md).

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 83.93%. Comparing base (3ff7ff8) to head (be44b0b).
⚠️ Report is 12 commits behind head on main.
✅ All tests successful. No failed tests found.

Additional details and impacted files

Files with missing lines	Coverage Δ
src/awkward/contents/numpyarray.py	92.37% <100.00%> (+0.45%)	⬆️

[github-actions]

The documentation preview is ready to be viewed at http://preview.awkward-array.org.s3-website.us-east-1.amazonaws.com/PR4270

[TaiSakuma]

Review notes on the approach, with behavior probes run against main (v2.12.0) to confirm the before-state and the cast-back assumptions.

Verified

• complex128 ak.sort/ak.argsort raise KeyError on main, so the new TypeError replaces an error, not working behavior. All reducers, including ak.min/ak.max/ak.argmin, work on complex128, so the complex256 → complex128 reduce cast reaches working kernels.
• The _reduce_next cast-back predicate (node.dtype == cast) fires exactly for the value-carrying results: ak.sum/ak.prod on float32 return float32, while argmin/argmax/count_nonzero return int64 and any/all return bool, so those keep their result dtype.
• Accumulating float16 in float32 matches NumPy, which also uses a wider intermediate for half precision: np.sum of [60000, 60000, -60000] float16 gives 60000, not inf, and so does the cast path.
• The hook is backend-agnostic, so the typetracer path predicts the same restored dtype that real execution produces.

Main concern: float128 ak.sort alters the values

_sort_next sorts the float64-cast data and then casts the sorted float64 values back to float128, rather than permuting the original values, so every element comes back rounded to float64 precision and relabeled float128. A sort should return the same multiset reordered. The tests only use values exactly representable in float64, so they cannot detect this; sorting [float128(1) + float128(2)**-60, float128(1)] is expected to return two elements that both equal 1.0. (float16 is unaffected — the float32 round trip is exact — and complex256 never reaches sort.) One fix: gather the original data through the argsort carry (per-bin local indices plus bin starts) instead of casting values back; at minimum, document the rounding and pin it with a test. A related, milder caveat: float128 values that are distinct but equal after float64 rounding sort by original position rather than by true float128 order (both sort and argsort; inherent to the cast approach, worth a sentence in the code comment).

Suggestion: complex TypeError message

"complex numbers have no total order" conflicts with NumPy (np.sort orders complex lexicographically) and with awkward itself: ak.min/ak.max on complex128 work today and return the lexicographic extremum. A message such as "sorting complex numbers is not supported" refuses without asserting a rationale the library's own reducers contradict. The parenthetical "(this applies to complex64/128/256 alike)" could be dropped. A NumPy-compatible lexicographic complex sort could be a follow-up issue.

Test suggestions

• The complex256 reducer test covers sum/prod/count_nonzero; min/max/argmin/all/any also work through the cast (confirmed on complex128) and could be added with one more parametrize entry.
• There is no NaN/inf case for float16 sort or reducers.
• tests-cuda/ does not run in the pull-request CI — has the new file been run on real hardware?

Follow-ups (not blockers)

• _unique/_is_unique in NumpyArray also dispatch awkward_sort by value dtype, reached via ak.validity_error/ak.is_valid on __array__="categorical" content, so a categorical float16 array still raises KeyError after this PR. The new helper could be reused there.
• Once this merges, has_issue_4259 in tests/properties/operations/known_issues.py should be removed so the property tests exercise the fixed path.

🤖 Generated with Claude Code