feat: add axis=None reducer specializations (#3653)

* feat: add axis=None reducer specializations

* only add an axis=None specialization for sum

* fix check for rectangular 1D arrays to determine if axis=None can be used

* fix axis=0,1 to axis=None specialization path

* style: pre-commit fixes

* simplify if condition

---------

Co-authored-by: Ianna Osborne <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 3 people

src/awkward/_do.py

@@ -224,9 +224,13 @@ def reduce(
     keepdims: bool = False,
     behavior: dict | None = None,
 ):
+    # store the original reducer for potential reuse later
+    original_reducer = reducer
     reducer = layout.backend.prepare_reducer(reducer)
 
     if axis is None:
+        del original_reducer  # not used below this point
+
         parts = remove_structure(
             layout,
             flatten_records=False,
@@ -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
@@ -288,6 +302,19 @@ def reduce(
                     f"(which is {depth})"
                 )
 
+            # a flat array can be fully reduced with axis=None or axis=0 or axis=-1,
+            # so we treat them as equivalent and recurse to the axis=None specialization
+            if depth == negaxis == 1:
+                return reduce(
+                    layout=layout,
+                    reducer=original_reducer,
+                    axis=None,
+                    mask=mask,
+                    keepdims=keepdims,
+                    behavior=behavior,
+                )
+            del original_reducer  # not used below this point
+
         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,17 @@ 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 max(
         self,
         x: ArrayLikeT,

src/awkward/_nplikes/cupy.py

@@ -139,6 +139,21 @@ 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 max(
         self,
         x: ArrayLike,

src/awkward/_nplikes/numpy_like.py

@@ -467,6 +467,16 @@ 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 count_nonzero(
         self, x: ArrayLikeT, *, axis: int | tuple[int, ...] | None = None

src/awkward/_nplikes/typetracer.py

@@ -1646,6 +1646,16 @@ 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 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,6 +444,25 @@ 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