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

awkward-cpp/src/cpu-kernels/awkward_reduce_sum.cpp

@@ -47,3 +47,11 @@ REDUCE_SUM(awkward_reduce_sum_int32_int16_64, int32_t, int16_t)
 REDUCE_SUM(awkward_reduce_sum_uint32_uint16_64, uint32_t, uint16_t)
 REDUCE_SUM(awkward_reduce_sum_int32_int32_64, int32_t, int32_t)
 REDUCE_SUM(awkward_reduce_sum_uint32_uint32_64, uint32_t, uint32_t)
+REDUCE_SUM(awkward_reduce_sum_float64_int8_64, double, int8_t)
+REDUCE_SUM(awkward_reduce_sum_float64_uint8_64, double, uint8_t)
+REDUCE_SUM(awkward_reduce_sum_float64_int16_64, double, int16_t)
+REDUCE_SUM(awkward_reduce_sum_float64_uint16_64, double, uint16_t)
+REDUCE_SUM(awkward_reduce_sum_float64_int32_64, double, int32_t)
+REDUCE_SUM(awkward_reduce_sum_float64_uint32_64, double, uint32_t)
+REDUCE_SUM(awkward_reduce_sum_float64_int64_64, double, int64_t)
+REDUCE_SUM(awkward_reduce_sum_float64_uint64_64, double, uint64_t)

kernel-specification.yml

@@ -5175,6 +5175,54 @@ kernels:
           - {name: fromptr, type: "Const[List[uint32_t]]", dir: in, role: reducer-fromptr}
           - {name: offsets, type: "Const[List[int64_t]]", dir: in, role: reducer-offsets}
           - {name: outlength, type: "int64_t", dir: in, role: reducer-outlength}
+      - name: awkward_reduce_sum_float64_int8_64
+        args:
+          - {name: toptr, type: "List[double]", dir: out}
+          - {name: fromptr, type: "Const[List[int8_t]]", dir: in, role: reducer-fromptr}
+          - {name: offsets, type: "Const[List[int64_t]]", dir: in, role: reducer-offsets}
+          - {name: outlength, type: "int64_t", dir: in, role: reducer-outlength}
+      - name: awkward_reduce_sum_float64_uint8_64
+        args:
+          - {name: toptr, type: "List[double]", dir: out}
+          - {name: fromptr, type: "Const[List[uint8_t]]", dir: in, role: reducer-fromptr}
+          - {name: offsets, type: "Const[List[int64_t]]", dir: in, role: reducer-offsets}
+          - {name: outlength, type: "int64_t", dir: in, role: reducer-outlength}
+      - name: awkward_reduce_sum_float64_int16_64
+        args:
+          - {name: toptr, type: "List[double]", dir: out}
+          - {name: fromptr, type: "Const[List[int16_t]]", dir: in, role: reducer-fromptr}
+          - {name: offsets, type: "Const[List[int64_t]]", dir: in, role: reducer-offsets}
+          - {name: outlength, type: "int64_t", dir: in, role: reducer-outlength}
+      - name: awkward_reduce_sum_float64_uint16_64
+        args:
+          - {name: toptr, type: "List[double]", dir: out}
+          - {name: fromptr, type: "Const[List[uint16_t]]", dir: in, role: reducer-fromptr}
+          - {name: offsets, type: "Const[List[int64_t]]", dir: in, role: reducer-offsets}
+          - {name: outlength, type: "int64_t", dir: in, role: reducer-outlength}
+      - name: awkward_reduce_sum_float64_int32_64
+        args:
+          - {name: toptr, type: "List[double]", dir: out}
+          - {name: fromptr, type: "Const[List[int32_t]]", dir: in, role: reducer-fromptr}
+          - {name: offsets, type: "Const[List[int64_t]]", dir: in, role: reducer-offsets}
+          - {name: outlength, type: "int64_t", dir: in, role: reducer-outlength}
+      - name: awkward_reduce_sum_float64_uint32_64
+        args:
+          - {name: toptr, type: "List[double]", dir: out}
+          - {name: fromptr, type: "Const[List[uint32_t]]", dir: in, role: reducer-fromptr}
+          - {name: offsets, type: "Const[List[int64_t]]", dir: in, role: reducer-offsets}
+          - {name: outlength, type: "int64_t", dir: in, role: reducer-outlength}
+      - name: awkward_reduce_sum_float64_int64_64
+        args:
+          - {name: toptr, type: "List[double]", dir: out}
+          - {name: fromptr, type: "Const[List[int64_t]]", dir: in, role: reducer-fromptr}
+          - {name: offsets, type: "Const[List[int64_t]]", dir: in, role: reducer-offsets}
+          - {name: outlength, type: "int64_t", dir: in, role: reducer-outlength}
+      - name: awkward_reduce_sum_float64_uint64_64
+        args:
+          - {name: toptr, type: "List[double]", dir: out}
+          - {name: fromptr, type: "Const[List[uint64_t]]", dir: in, role: reducer-fromptr}
+          - {name: offsets, type: "Const[List[int64_t]]", dir: in, role: reducer-offsets}
+          - {name: outlength, type: "int64_t", dir: in, role: reducer-outlength}
     description: null
     definition: |
       def awkward_reduce_sum(toptr, fromptr, offsets, outlength):

src/awkward/_nplikes/array_module.py

@@ -747,9 +747,12 @@ def sum(
         axis: ShapeItem | tuple[ShapeItem, ...] | None = None,
         keepdims: bool = False,
         maybe_out: ArrayLikeT | None = None,
+        dtype: DTypeLike | None = None,
     ) -> ArrayLikeT:
         (x,) = maybe_materialize(x)
-        return self._module.sum(x, axis=axis, keepdims=keepdims, out=maybe_out)
+        return self._module.sum(
+            x, axis=axis, keepdims=keepdims, out=maybe_out, dtype=dtype
+        )
 
     def max(
         self,

src/awkward/_nplikes/cupy.py

@@ -148,9 +148,12 @@ def sum(
         axis: ShapeItem | tuple[ShapeItem, ...] | None = None,
         keepdims: bool = False,
         maybe_out: ArrayLike | None = None,
+        dtype: DTypeLike | None = None,
     ) -> ArrayLike:
         (x,) = maybe_materialize(x)
-        out = self._module.sum(x, axis=axis, keepdims=keepdims, out=maybe_out)
+        out = self._module.sum(
+            x, axis=axis, keepdims=keepdims, out=maybe_out, dtype=dtype
+        )
         if axis is None and not keepdims and isinstance(out, self._module.ndarray):
             return out.item()
         else:

src/awkward/_nplikes/numpy_like.py

@@ -485,6 +485,7 @@ def sum(
         axis: int | tuple[int, ...] | None = None,
         keepdims: bool = False,
         maybe_out: ArrayLikeT | None = None,
+        dtype: DTypeLike | None = None,
     ) -> ArrayLikeT: ...
 
     @abstractmethod

src/awkward/_nplikes/typetracer.py

@@ -1663,6 +1663,7 @@ def sum(
         axis: int | tuple[int, ...] | None = None,
         keepdims: bool = False,
         maybe_out: TypeTracerArray | None = None,
+        dtype: DTypeLike | None = None,
     ) -> TypeTracerArray:
         assert isinstance(x, TypeTracerArray)
         try_touch_data(x)
@@ -1675,8 +1676,11 @@ def sum(
         # NumPy's `sum` promotes integer and boolean inputs narrower than the
         # platform integer up to the platform integer (e.g. bool/int32 -> intp).
         # Mirror that behaviour so the typetracer reports the same dtype as the
-        # concrete backends.
-        out_dtype = numpy.empty(0, dtype=x.dtype).sum().dtype
+        # concrete backends. An explicit `dtype` overrides the promotion.
+        if dtype is not None:
+            out_dtype = numpy.dtype(dtype)
+        else:
+            out_dtype = numpy.empty(0, dtype=x.dtype).sum().dtype
 
         if axis is None:
             x = cast(TypeTracerArray, self.reshape(x, (-1,)))

src/awkward/_reducers.py

@@ -334,8 +334,15 @@ class Sum(KernelReducer):
     preferred_dtype: Final = np.float64
     needs_position: Final = False
 
+    def __init__(self, dtype: DTypeLike | None = None):
+        # Optional forced accumulator/output dtype. When set (e.g. float64),
+        # integer/bool inputs are summed *directly* into that dtype, matching
+        # NumPy's `np.add.reduce(x, dtype=...)` and avoiding integer overflow
+        # without materialising a promoted copy of the input.
+        self._dtype = None if dtype is None else np.dtype(dtype)
+
     def axis_none_reducer(self) -> AxisNoneSum:
-        return AxisNoneSum()
+        return AxisNoneSum(self._dtype)
 
     def apply(
         self,
@@ -386,12 +393,24 @@ def apply(
                 )
             else:
                 raise NotImplementedError
+            # bool sums cannot overflow (max == segment length); if a float
+            # accumulator was requested, cast the small (outlength) result.
+            if self._dtype is not None:
+                result = result.astype(self._dtype)
             return ak.contents.NumpyArray(result, backend=array.backend)
         else:
+            is_complex = array.dtype.type in (np.complex128, np.complex64)
             kernel_array_data = array.data.view(self._dtype_for_kernel(array.dtype))
+            # A forced float accumulator (self._dtype) sums integers directly
+            # into that dtype via the awkward_reduce_sum_<out>_<in>_64 kernel,
+            # so no promoted input copy is allocated. Not applied to complex.
+            if self._dtype is not None and not is_complex:
+                result_dtype = self._dtype
+            else:
+                result_dtype = self._promote_integer_rank(kernel_array_data.dtype)
             result = array.backend.nplike.empty(
                 self._length_for_kernel(array.dtype.type, outlength),
-                dtype=self._promote_integer_rank(kernel_array_data.dtype),
+                dtype=result_dtype,
             )
             if array.dtype.type in (np.complex128, np.complex64):
                 assert offsets.nplike is array.backend.nplike
@@ -424,8 +443,12 @@ def apply(
                     )
                 )
 
+            if self._dtype is not None and not is_complex:
+                final_dtype = self._dtype
+            else:
+                final_dtype = self._promote_integer_rank(array.dtype)
             return ak.contents.NumpyArray(
-                result.view(self._promote_integer_rank(array.dtype)),
+                result.view(final_dtype),
                 backend=array.backend,
             )
 
@@ -445,7 +468,12 @@ def apply(
 
         nplike = array.backend.nplike
         reduce_fn = getattr(nplike, self.name)
-        result_scalar = reduce_fn(array.data, axis=None)
+        # A forced accumulator dtype maps straight onto NumPy/CuPy's `dtype=`
+        # argument, casting element-by-element inside the reduction (no copy).
+        if self._dtype is not None and array.dtype.kind != "c":
+            result_scalar = reduce_fn(array.data, axis=None, dtype=self._dtype)
+        else:
+            result_scalar = reduce_fn(array.data, axis=None)
         result_array = nplike.reshape(nplike.asarray(result_scalar), (1,))
 
         return ak.contents.NumpyArray(result_array, backend=array.backend)

src/awkward/operations/ak_sum.py

@@ -27,6 +27,7 @@ def sum(
     *,
     keepdims=False,
     mask_identity=False,
+    dtype=None,
     highlevel=True,
     behavior=None,
     attrs=None,
@@ -60,6 +61,11 @@ def sum(
         mask_identity (bool): If True, reducing over empty lists results in
             None (an option type); otherwise, reducing over empty lists
             results in the operation's identity.
+        dtype (None or dtype): If not None, the reduction accumulates directly
+            into this dtype (e.g. `np.float64`), matching NumPy's
+            `np.sum(..., dtype=...)`. Casting integer/bool input to `float64`
+            this way avoids integer overflow without allocating a promoted copy
+            of the input. If None, the default integer-promotion rules apply.
         highlevel (bool): If True, return an #ak.Array; otherwise, return
             a low-level #ak.contents.Content subclass.
         behavior (None or dict): Custom #ak.behavior for the output array, if
@@ -222,7 +228,9 @@ def sum(
     yield (array,)
 
     # Implementation
-    return _impl(array, axis, keepdims, mask_identity, highlevel, behavior, attrs)
+    return _impl(
+        array, axis, keepdims, mask_identity, highlevel, behavior, attrs, dtype
+    )
 
 
 @high_level_function()
@@ -289,7 +297,7 @@ def nansum(
     )
 
 
-def _impl(array, axis, keepdims, mask_identity, highlevel, behavior, attrs):
+def _impl(array, axis, keepdims, mask_identity, highlevel, behavior, attrs, dtype=None):
     with HighLevelContext(behavior=behavior, attrs=attrs) as ctx:
         layout = ctx.unwrap(array, allow_record=False, primitive_policy="error")
 
@@ -310,7 +318,7 @@ def _impl(array, axis, keepdims, mask_identity, highlevel, behavior, attrs):
 
     axis = regularize_axis(axis, none_allowed=True)
 
-    reducer = ak._reducers.Sum()
+    reducer = ak._reducers.Sum(dtype=dtype)
 
     out = ak._do.reduce(
         layout,

tests/test_4100_typetracer_correctness.py

@@ -28,6 +28,31 @@ def test_sum_dtype_promotion_bool(tt):
     assert out.dtype != np.dtype(np.bool_)
 
 
+def test_sum_forced_dtype_axis_none(tt):
+    # Explicit dtype= is reported verbatim, not the default promotion.
+    x = TypeTracerArray._new(np.dtype(np.int32), (unknown_length,))
+    out = tt.sum(x, axis=None, dtype=np.float64)
+    assert out.shape == ()
+    assert out.dtype == np.dtype(np.float64)
+
+
+def test_sum_forced_dtype_overrides_promotion(tt):
+    # For bool input the default would promote to intp; an explicit float64
+    # dtype must win (exercises the `dtype is not None` branch).
+    x = TypeTracerArray._new(np.dtype(np.bool_), (unknown_length,))
+    out = tt.sum(x, axis=None, dtype=np.float64)
+    assert out.dtype == np.dtype(np.float64)
+    assert out.dtype != np.empty(0, dtype=np.bool_).sum().dtype
+
+
+def test_sum_forced_dtype_axis_keepdims(tt):
+    # Forced dtype also applies on the segmented (axis given) path.
+    x = TypeTracerArray._new(np.dtype(np.int64), (3, unknown_length))
+    out = tt.sum(x, axis=1, keepdims=True, dtype=np.float64)
+    assert out.dtype == np.dtype(np.float64)
+    assert out.shape == (3, 1)
+
+
 def test_stack_symbolic_shape(tt):
     a = TypeTracerArray._new(np.dtype(np.int32), (unknown_length,))
     b = TypeTracerArray._new(np.dtype(np.float64), (unknown_length,))

tests/test_4232_sum_dtype_accumulator.py

@@ -0,0 +1,91 @@
+# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+import awkward as ak
+from awkward._reducers import KernelReducer
+
+
+def _has_float64_integer_kernel():
+    # The segmented (axis != None) float64/integer accumulator needs the
+    # awkward_reduce_sum_float64_<int>_64 kernels; skip if awkward-cpp predates
+    # them.
+    try:
+        ak.sum(ak.Array([[1, 2], [3]]), axis=-1, dtype=np.float64)
+    except KeyError:
+        return False
+    return True
+
+
+needs_kernel = pytest.mark.skipif(
+    not _has_float64_integer_kernel(),
+    reason="requires awkward-cpp built with float64/integer reduce_sum kernels",
+)
+
+
+def test_sum_dtype_none_unchanged():
+    a = ak.Array([[1, 2, 3], [], [4, 5]])
+    assert ak.sum(a, axis=-1).to_list() == [6, 0, 9]
+    assert str(ak.sum(a, axis=-1).type) == "3 * int64"
+    assert ak.sum(a) == 15
+
+
+def test_sum_axis_none_float64_accumulator():
+    # int32 total (6e9) exceeds int32 but the float64 accumulator is exact here.
+    data = np.array([2_000_000_000, 2_000_000_000, 2_000_000_000], dtype=np.int32)
+    out = ak.sum(ak.Array(data), axis=None, dtype=np.float64)
+    assert out == pytest.approx(np.sum(data, dtype=np.float64))
+    assert np.asarray(out).dtype == np.dtype(np.float64)
+
+
+def test_sum_axis_none_bool_float64():
+    out = ak.sum(ak.Array(np.array([True, False, True])), axis=None, dtype=np.float64)
+    assert out == 2.0
+
+
+def test_sum_axis_none_float_input_unchanged():
+    data = np.array([1.5, 2.5, 3.5], dtype=np.float64)
+    assert ak.sum(ak.Array(data), axis=None) == pytest.approx(np.sum(data))
+
+
+@needs_kernel
+def test_sum_segmented_float64_accumulator():
+    a = ak.Array([[1, 2, 3], [], [4, 5]])
+    out = ak.sum(a, axis=-1, dtype=np.float64)
+    assert out.to_list() == [6.0, 0.0, 9.0]
+    assert str(out.type) == "3 * float64"
+
+
+@needs_kernel
+def test_sum_segmented_no_overflow():
+    # int32 squares would overflow, but here we sum large int32 values into
+    # float64 directly (no promoted input copy).
+    a = ak.Array([[2_000_000_000, 2_000_000_000], [2_000_000_000]])
+    out = ak.sum(a, axis=-1, dtype=np.float64)
+    assert out.to_list() == [4_000_000_000.0, 2_000_000_000.0]
+
+
+def test_sum_segmented_bool_float64_cast():
+    # Segmented bool reduction with a forced float64 accumulator: the small
+    # (per-segment) int result is cast to float64 (Sum.apply bool branch).
+    a = ak.Array([[True, True, False], [], [True, True]])
+    out = ak.sum(a, axis=-1, dtype=np.float64)
+    assert out.to_list() == [2.0, 0.0, 2.0]
+    assert str(out.type) == "3 * float64"
+
+
+def test_sum_bool_unsupported_accumulator_raises(monkeypatch):
+    # Defensive branch: if the bool reduction's promoted result dtype is neither
+    # a 32- nor 64-bit integer, Sum.apply raises NotImplementedError. Force it
+    # by making _promote_integer_rank return an unsupported dtype.
+    monkeypatch.setattr(
+        KernelReducer,
+        "_promote_integer_rank",
+        staticmethod(lambda dtype: np.dtype(np.float32)),
+    )
+    a = ak.Array([[True, False], [True]])
+    with pytest.raises(NotImplementedError):
+        ak.sum(a, axis=-1)