feat: add axis=None reducer specializations

src/awkward/_do.py

@@ -226,6 +226,10 @@ def reduce(
 ):
     reducer = layout.backend.prepare_reducer(reducer)
 
+    # a flat array can only be reduced with axis=None
+    if isinstance(layout, ak.contents.NumpyArray) and layout.data.ndim <= 1:
+        axis = None
+
     if axis is None:
         parts = remove_structure(
             layout,
@@ -247,6 +251,16 @@ def reduce(
         else:
             (layout,) = parts
 
+        # Check if we're running with concrete data and if the reducer has a axis=None specialization.
+        # If both are true, we use the specialized reducer. This allows us to use optimized implementations
+        # from e.g. NumPy, but also make use of potentially better algorithms, i.e. Kahan summation for sum.
+        if (
+            layout.backend.nplike.known_data
+            and (specialization := reducer.axis_none_reducer()) is not None
+        ):
+            # overwrite reducer if it has an axis=None version
+            reducer = specialization
+
         starts = ak.index.Index64.zeros(1, layout.backend.nplike)
         parents = ak.index.Index64.zeros(layout.length, layout.backend.nplike)
         shifts = None

src/awkward/_nplikes/array_module.py

@@ -665,6 +665,28 @@ def min(
         (x,) = maybe_materialize(x)
         return self._module.min(x, axis=axis, keepdims=keepdims, out=maybe_out)
 
+    def sum(
+        self,
+        x: ArrayLikeT,
+        *,
+        axis: ShapeItem | tuple[ShapeItem, ...] | None = None,
+        keepdims: bool = False,
+        maybe_out: ArrayLikeT | None = None,
+    ) -> ArrayLikeT:
+        (x,) = maybe_materialize(x)
+        return self._module.sum(x, axis=axis, keepdims=keepdims, out=maybe_out)
+
+    def prod(
+        self,
+        x: ArrayLikeT,
+        *,
+        axis: ShapeItem | tuple[ShapeItem, ...] | None = None,
+        keepdims: bool = False,
+        maybe_out: ArrayLikeT | None = None,
+    ) -> ArrayLikeT:
+        (x,) = maybe_materialize(x)
+        return self._module.prod(x, axis=axis, keepdims=keepdims, out=maybe_out)
+
     def max(
         self,
         x: ArrayLikeT,

src/awkward/_nplikes/cupy.py

@@ -139,6 +139,36 @@ def min(
         else:
             return out
 
+    def sum(
+        self,
+        x: ArrayLike,
+        *,
+        axis: ShapeItem | tuple[ShapeItem, ...] | None = None,
+        keepdims: bool = False,
+        maybe_out: ArrayLike | None = None,
+    ) -> ArrayLike:
+        (x,) = maybe_materialize(x)
+        out = self._module.sum(x, axis=axis, out=maybe_out)
+        if axis is None and isinstance(out, self._module.ndarray):
+            return out.item()
+        else:
+            return out
+
+    def prod(
+        self,
+        x: ArrayLike,
+        *,
+        axis: ShapeItem | tuple[ShapeItem, ...] | None = None,
+        keepdims: bool = False,
+        maybe_out: ArrayLike | None = None,
+    ) -> ArrayLike:
+        (x,) = maybe_materialize(x)
+        out = self._module.prod(x, axis=axis, out=maybe_out)
+        if axis is None and isinstance(out, self._module.ndarray):
+            return out.item()
+        else:
+            return out
+
     def max(
         self,
         x: ArrayLike,

src/awkward/_nplikes/numpy_like.py

@@ -467,6 +467,26 @@ def max(
         maybe_out: ArrayLikeT | None = None,
     ) -> ArrayLikeT: ...
 
+    @abstractmethod
+    def sum(
+        self,
+        x: ArrayLikeT,
+        *,
+        axis: int | tuple[int, ...] | None = None,
+        keepdims: bool = False,
+        maybe_out: ArrayLikeT | None = None,
+    ) -> ArrayLikeT: ...
+
+    @abstractmethod
+    def prod(
+        self,
+        x: ArrayLikeT,
+        *,
+        axis: int | tuple[int, ...] | None = None,
+        keepdims: bool = False,
+        maybe_out: ArrayLikeT | None = None,
+    ) -> ArrayLikeT: ...
+
     @abstractmethod
     def count_nonzero(
         self, x: ArrayLikeT, *, axis: int | tuple[int, ...] | None = None

src/awkward/_nplikes/typetracer.py

@@ -1646,6 +1646,26 @@ def max(
     ) -> TypeTracerArray:
         return self.min(x, axis=axis, keepdims=keepdims, maybe_out=maybe_out)
 
+    def sum(
+        self,
+        x: TypeTracerArray,
+        *,
+        axis: int | tuple[int, ...] | None = None,
+        keepdims: bool = False,
+        maybe_out: TypeTracerArray | None = None,
+    ) -> TypeTracerArray:
+        return self.min(x, axis=axis, keepdims=keepdims, maybe_out=maybe_out)
+
+    def prod(
+        self,
+        x: TypeTracerArray,
+        *,
+        axis: int | tuple[int, ...] | None = None,
+        keepdims: bool = False,
+        maybe_out: TypeTracerArray | None = None,
+    ) -> TypeTracerArray:
+        return self.min(x, axis=axis, keepdims=keepdims, maybe_out=maybe_out)
+
     def array_str(
         self,
         x: TypeTracerArray,

src/awkward/_reducers.py

@@ -33,6 +33,11 @@ def preferred_dtype(self) -> DTypeLike: ...
     def highlevel_function(cls):
         return getattr(ak.operations, cls.name)
 
+    @classmethod
+    def axis_none_reducer(cls) -> Reducer | None:
+        """A specialized version for axis=None reductions, or None if there is none."""
+        return None
+
     @abstractmethod
     def apply(
         self,
@@ -336,6 +341,10 @@ class Sum(KernelReducer):
     preferred_dtype: Final = np.float64
     needs_position: Final = False
 
+    @classmethod
+    def axis_none_reducer(cls):
+        return AxisNoneSum()
+
     def apply(
         self,
         array: ak.contents.NumpyArray,
@@ -435,11 +444,34 @@ def apply(
             )
 
 
+class AxisNoneSum(Sum):
+    def apply(
+        self,
+        array: ak.contents.NumpyArray,
+        parents: ak.index.Index,
+        starts: ak.index.Index,
+        shifts: ak.index.Index | None,
+        outlength: ShapeItem,
+    ) -> ak.contents.NumpyArray:
+        del parents, starts, shifts, outlength  # Unused
+        assert isinstance(array, ak.contents.NumpyArray)
+        if array.dtype.kind == "M":
+            raise ValueError(f"cannot compute the sum (ak.sum) of {array.dtype!r}")
+        reduce_fn = getattr(array.backend.nplike, self.name)
+        return ak.contents.NumpyArray(
+            [reduce_fn(array.data, axis=None)], backend=array.backend
+        )
+
+
 class Prod(KernelReducer):
     name: Final = "prod"
     preferred_dtype: Final = np.float64
     needs_position: Final = False
 
+    @classmethod
+    def axis_none_reducer(cls) -> Reducer | None:
+        return AxisNoneProd()
+
     def apply(
         self,
         array: ak.contents.NumpyArray,
@@ -526,6 +558,25 @@ def apply(
             )
 
 
+class AxisNoneProd(Prod):
+    def apply(
+        self,
+        array: ak.contents.NumpyArray,
+        parents: ak.index.Index,
+        starts: ak.index.Index,
+        shifts: ak.index.Index | None,
+        outlength: ShapeItem,
+    ) -> ak.contents.NumpyArray:
+        del parents, starts, shifts, outlength  # Unused
+        assert isinstance(array, ak.contents.NumpyArray)
+        if array.dtype.kind.upper() == "M":
+            raise ValueError(f"cannot compute the product (ak.prod) of {array.dtype!r}")
+        reduce_fn = getattr(array.backend.nplike, self.name)
+        return ak.contents.NumpyArray(
+            [reduce_fn(array.data, axis=None)], backend=array.backend
+        )
+
+
 class Any(KernelReducer):
     name: Final = "any"
     preferred_dtype: Final = np.bool_