Skip to content

feat: add a dtype= accumulator to ak.sum (overflow-safe, copy-free integer sums)#4232

Open
ianna wants to merge 5 commits into
mainfrom
ianna/copy_free_integer_sums
Open

feat: add a dtype= accumulator to ak.sum (overflow-safe, copy-free integer sums)#4232
ianna wants to merge 5 commits into
mainfrom
ianna/copy_free_integer_sums

Conversation

@ianna

@ianna ianna commented Jul 23, 2026

Copy link
Copy Markdown
Member

Motivation

Integer reductions in the descriptive statistics overflow because the sum accumulates at the input dtype in the segmented kernel. The Python-level fix (#4223) casts integer input to float64 up front, which allocates a full-size temporary — the copy raised against #3527. NumPy avoids this by casting inside the reduction (np.add.reduce(x, dtype=np.float64)); awkward couldn't, because it had no float64-accumulator / integer-input reduce kernel.

This PR adds that primitive: ak.sum(x, dtype=np.float64) sums integer/bool input directly into float64, matching NumPy and avoiding a promoted copy.

Changes

  • Kernels: 8 new awkward_reduce_sum_float64_{int,uint}{8,16,32,64}_64 specializations in kernel-specification.yml and registered in awkward-cpp/src/cpu-kernels/awkward_reduce_sum.cpp. The kernel body is unchanged — awkward_reduce_sum already casts per element via static_cast<OUT>, so a float64 accumulator over integer input is a pure specialization. Generated kernels.h and the signature tables are regenerated from the spec.
  • Sum reducer (_reducers.py): Sum(dtype=None). When set, the accumulator/output dtype flows into the segmented kernel selection and the axis=None path; the bool path casts its small result.
  • nplike: dtype= added to sum on all four backends (array_module, cupy, typetracer, and the protocol in numpy_like), forwarding to NumPy/CuPy sum(dtype=).
  • ak.sum: new dtype= keyword + docstring.
  • Tests: tests/test_4223_sum_dtype_accumulator.py.

Scope / follow-ups

  • ak.mean/ak.var are intentionally not rewired here. Doing so would hard-couple the pure-Python package to the new kernel — integer mean/var over an axis would break against any released awkward-cpp predating it. They should adopt dtype= (and drop the fix: overflow-safe stats without needless copies #4223 promotion) in a follow-up, once these kernels ship in a released awkward-cpp.
  • This is Phase 1 of the plan for copy-free overflow-safe stats; var/std additionally need a fused sum-of-squares reducer (Phase 2).

Release ordering

Kernels live in awkward-cpp; the Python wiring is in awkward. The axis=None path works against any awkward-cpp (it maps onto NumPy/CuPy sum(dtype=)), but the segmented axis != None path requires a rebuilt awkward-cpp — so awkward-cpp must be released with these kernels before awkward relies on them. The two segmented tests skip when the kernels are absent.

Non-goals

ak.sum/ak.prod default dtype semantics are unchanged (NumPy preserves integer dtype); the float64 accumulator is opt-in. Complex/datetime unaffected.

AI disclosure

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

@github-actions github-actions Bot added the type/feat PR title type: feat (set automatically) label Jul 23, 2026
@ianna
ianna requested a review from TaiSakuma July 23, 2026 12:00
@codecov

codecov Bot commented Jul 23, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 84.05%. Comparing base (802b733) to head (cc9b934).
✅ All tests successful. No failed tests found.

Additional details and impacted files
Files with missing lines Coverage Δ
src/awkward/_nplikes/array_module.py 93.71% <100.00%> (ø)
src/awkward/_nplikes/cupy.py 91.45% <100.00%> (ø)
src/awkward/_nplikes/numpy_like.py 100.00% <ø> (ø)
src/awkward/_nplikes/typetracer.py 77.38% <100.00%> (+0.14%) ⬆️
src/awkward/_reducers.py 99.03% <100.00%> (+0.71%) ⬆️
src/awkward/operations/ak_sum.py 94.59% <100.00%> (ø)

@github-actions

Copy link
Copy Markdown

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

@ianna
ianna force-pushed the ianna/copy_free_integer_sums branch from ba0b2d5 to 1a19ee4 Compare July 23, 2026 13:30
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

type/feat PR title type: feat (set automatically)

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant