Repository: nv-legate/cunumeric

Issue:
Refactor the axis sanitization
Many operations optionally take an `axis` argument to denote a specific dimension of an ndarray. Although these operations sanitize the `axis` value in almost the same way, they implemented the logic individually and there's no common utility. We need to refactor the code so that all those operations use the same sanitization implementation.


Discussion:
@magnatelee I have one up-front question about philosophy/policy. Is our goal to match numpy's handling of axis args *exactly*, including any limitations? Or is it acceptable for our APIs to be "wider"? I am thinking specifically of the possibility that some numpy method accepts a strict axis type, but we want to accept a more general (convertible) "axis-like" everywhere for simplicity. 

Put another way: is it OK if our functions accept everything numpy accepts, *and more*? (I'm also speculating here, perhaps numpy already accepts a more general "axis-like" everywhere and the issue is moot).
The goal is to match NumPy's behavior. In most cases, `axis` can be either None or an integer, and if the value is negative, it follows the usual Python rule of indexing (i.e., `-k` means the `k`-th value from the back). And an `IndexError` should be raised for out-of-bounds accesses (in this case, non-existent axes). There are some places where the `axis` can be a tuple of integers, in which case the logic applies element-wise. I feel there should be two functions, one that accepts tuples and one that doesn't, and we should choose the right one to call depending on the context. And of course, we first need to identity all places where an `axis` can be passed.
`grep` etc. was not really workable for a census. I wrote a little script to check more thoroughly based on the actual modules:
```python
from inspect import isclass, isfunction, signature
from pkgutil import walk_packages

import cunumeric

PATH = cunumeric.__path__
PREFIX = f"{cunumeric.__name__}."
SKIP = ("cunumeric.config", "cunumeric.runtime")

for loader, name, _ in walk_packages(path=PATH, prefix=PREFIX):
    if name in SKIP: continue

    module = loader.find_module(name).load_module(name)

    for k, v in vars(module).items():
        if k.startswith("_"): continue

        if isfunction(v) and "axis" in signature(v).parameters:
            print(f"{name}.{k}")

        if isclass(v) and v.__module__ == name:
            for kk, vv in vars(v).items():
                if kk.startswith("_"): continue

                try:
                    if isfunction(vv) and "axis" in signature(vv).parameters:
                        print(f"{name}.{k}.{kk}")
                except ValueError:
                    pass
```
And the results appear to be:
```
cunumeric._ufunc.ufunc.binary_ufunc.reduce
cunumeric.array.ndarray.all
cunumeric.array.ndarray.any
cunumeric.array.ndarray.argmax
cunumeric.array.ndarray.argmin
cunumeric.array.ndarray.take
cunumeric.array.ndarray.compress
cunumeric.array.ndarray.max
cunumeric.array.ndarray.mean
cunumeric.array.ndarray.min
cunumeric.array.ndarray.partition
cunumeric.array.ndarray.argpartition
cunumeric.array.ndarray.prod
cunumeric.array.ndarray.sort
cunumeric.array.ndarray.argsort
cunumeric.array.ndarray.squeeze
cunumeric.array.ndarray.sum
cunumeric.array.ndarray.flip
cunumeric.deferred.sort
cunumeric.deferred.DeferredArray.squeeze
cunumeric.deferred.DeferredArray.repeat
cunumeric.eager.EagerArray.squeeze
cunumeric.eager.EagerArray.repeat
cunumeric.eager.EagerArray.sort
cunumeric.eager.EagerArray.partition
cunumeric.fft.fft
cunumeric.fft.ifft
cunumeric.fft.rfft
cunumeric.fft.irfft
cunumeric.fft.hfft
cunumeric.fft.ihfft
cunumeric.fft.fft.fft
cunumeric.fft.fft.ifft
cunumeric.fft.fft.rfft
cunumeric.fft.fft.irfft
cunumeric.fft.fft.hfft
cunumeric.fft.fft.ihfft
cunumeric.linalg.norm
cunumeric.linalg.linalg.norm
cunumeric.module.normalize_axis_tuple
cunumeric.module.linspace
cunumeric.module.squeeze
cunumeric.module.check_shape_dtype
cunumeric.module.append
cunumeric.module.concatenate
cunumeric.module.stack
cunumeric.module.split
cunumeric.module.array_split
cunumeric.module.repeat
cunumeric.module.flip
cunumeric.module.take
cunumeric.module.compress
cunumeric.module.all
cunumeric.module.any
cunumeric.module.prod
cunumeric.module.sum
cunumeric.module.amax
cunumeric.module.max
cunumeric.module.amin
cunumeric.module.min
cunumeric.module.unique
cunumeric.module.argsort
cunumeric.module.sort
cunumeric.module.argpartition
cunumeric.module.partition
cunumeric.module.argmax
cunumeric.module.argmin
cunumeric.module.count_nonzero
cunumeric.module.mean
cunumeric.sort.sort
cunumeric.thunk.NumPyThunk.repeat
cunumeric.thunk.NumPyThunk.squeeze
```
@magnatelee I just wrote 
```
def _sanitize_axes(arg: AxisLike, ndim: int, name: str = "axis") -> tuple[int, ...]:
    if arg is None:
        axes = tuple(range(ndim))
    elif isinstance(arg, int):
        axes = (arg,)

    if any(type(ax) != int for ax in axes):
        raise TypeError(
            f"{name!r} must be an integer or a tuple of integers, but got {arg}"
        )

    computed_axes = tuple(ax + ndim if ax < 0 else ax for ax in axes)

    if any(ax < 0 for ax in computed_axes):
        raise ValueError(f"Invalid {name!r} value {arg}")

    return computed_axes
```
when I discovered that Numpy itself defines an almost identical [`normalize_axis_tuple`](https://github.com/numpy/numpy/blob/a4a757d3b971a6b25608a91e9985c4ff0b10dc82/numpy/core/numeric.py#L1347-L1403), which we already happen to use in a couple of places (e.g. `moveaxis`). The Numpy function does not handle the None case, and in Numpy we often see that case handled explicitly
```
    if axis is None:
        return roll(a.ravel(), shift, 0).reshape(a.shape)

    else:
        axis = normalize_axis_tuple(axis, a.ndim, allow_duplicate=True)
        ....
```
or 
```
    if axis is None:
        axis = tuple(range(x.ndim))
    else:
        axis = normalize_axis_tuple(axis, x.ndim)
```
Rather than duplicate what is already in numpy I would suggest first just apply `normalize_axis_tuple` in places where we are already doing that work by hand (and also handling None case explicitly at site, following Numpy's pattern)

Resolving patch:
diff --git a/cunumeric/array.py b/cunumeric/array.py
--- a/cunumeric/array.py
+++ b/cunumeric/array.py
@@ -17,10 +17,12 @@
 from collections.abc import Iterable
 from functools import reduce, wraps
 from inspect import signature
-from typing import Callable, Optional, Set, Tuple, TypeVar
+from typing import Callable, Optional, Set, TypeVar
 
 import numpy as np
 import pyarrow
+from numpy.core.multiarray import normalize_axis_index
+from numpy.core.numeric import normalize_axis_tuple
 from typing_extensions import ParamSpec
 
 from legate.core import Array
@@ -1782,10 +1784,8 @@ def take(self, indices, axis=None, out=None, mode="raise"):
         if axis is None:
             self = self.ravel()
             axis = 0
-        elif axis < 0:
-            axis = self.ndim + axis
-        if axis < 0 or axis >= self.ndim:
-            raise ValueError("axis argument is out of bounds")
+        else:
+            axis = normalize_axis_index(axis, self.ndim)
 
         # TODO remove "raise" logic when bounds check for advanced
         # indexing is implementd
@@ -1935,7 +1935,7 @@ def choose(self, choices, out=None, mode="raise"):
     def compress(self, condition, axis=None, out=None):
         """a.compress(self, condition, axis=None, out=None)
 
-        Return selected slices of an array along given axis..
+        Return selected slices of an array along given axis.
 
         Refer to :func:`cunumeric.compress` for full documentation.
 
@@ -1959,9 +1959,12 @@ def compress(self, condition, axis=None, out=None):
                 category=RuntimeWarning,
             )
             condition = condition.astype(bool)
+
         if axis is None:
             axis = 0
             a = self.ravel()
+        else:
+            axis = normalize_axis_index(axis, self.ndim)
 
         if a.shape[axis] < condition.shape[0]:
             raise ValueError(
@@ -2530,7 +2533,7 @@ def getfield(self, dtype, offset=0):
             "for ndarray.getfield"
         )
 
-    def _convert_singleton_key(self, args: Tuple):
+    def _convert_singleton_key(self, args: tuple):
         if len(args) == 0 and self.size == 1:
             return (0,) * self.ndim
         if len(args) == 1 and isinstance(args[0], int):
@@ -3070,9 +3073,7 @@ def squeeze(self, axis=None):
                         "all axis to squeeze must be less than ndim"
                     )
                 if self.shape[axis] != 1:
-                    raise ValueError(
-                        "axis to squeeze must have extent " "of one"
-                    )
+                    raise ValueError("axis to squeeze must have extent of one")
             elif isinstance(axis, tuple):
                 for ax in axis:
                     if ax >= self.ndim:
@@ -3593,23 +3594,11 @@ def _perform_unary_reduction(
             raise NotImplementedError(
                 "(arg)max/min not supported for complex-type arrays"
             )
-        # Compute the output shape
-        axes = axis
-        if axes is None:
-            axes = tuple(range(src.ndim))
-        elif not isinstance(axes, tuple):
-            axes = (axes,)
 
-        if any(type(ax) != int for ax in axes):
-            raise TypeError(
-                "'axis' must be an integer or a tuple of integers, "
-                f"but got {axis}"
-            )
-
-        axes = tuple(ax + src.ndim if ax < 0 else ax for ax in axes)
-
-        if any(ax < 0 for ax in axes):
-            raise ValueError(f"Invalid 'axis' value {axis}")
+        if axis is None:
+            axes = tuple(range(src.ndim))
+        else:
+            axes = normalize_axis_tuple(axis, src.ndim)
 
         out_shape = ()
         for dim in range(src.ndim):
diff --git a/cunumeric/linalg/linalg.py b/cunumeric/linalg/linalg.py
--- a/cunumeric/linalg/linalg.py
+++ b/cunumeric/linalg/linalg.py
@@ -20,9 +20,8 @@
 from cunumeric._ufunc.math import add, sqrt as _sqrt
 from cunumeric.array import add_boilerplate, convert_to_cunumeric_ndarray
 from cunumeric.module import dot, empty_like, eye, matmul, ndarray
-from numpy.core.multiarray import (  # type: ignore [attr-defined]
-    normalize_axis_index,
-)
+from numpy.core.multiarray import normalize_axis_index  # type: ignore
+from numpy.core.numeric import normalize_axis_tuple  # type: ignore
 
 if TYPE_CHECKING:
     import numpy.typing as npt
@@ -424,14 +423,11 @@ def norm(
                 ret = ret.reshape(ndim * [1])
             return ret
 
-    # Normalize the `axis` argument to a tuple.
-    nd = x.ndim
     if axis is None:
-        computed_axis = tuple(range(nd))
-    elif not isinstance(axis, tuple):
-        computed_axis = (axis,)
+        computed_axis = tuple(range(x.ndim))
     else:
-        computed_axis = axis
+        computed_axis = normalize_axis_tuple(axis, x.ndim)
+
     for ax in computed_axis:
         if not isinstance(ax, int):
             raise TypeError(
@@ -469,8 +465,8 @@ def norm(
             return ret
     elif len(computed_axis) == 2:
         row_axis, col_axis = computed_axis
-        row_axis = normalize_axis_index(row_axis, nd)
-        col_axis = normalize_axis_index(col_axis, nd)
+        row_axis = normalize_axis_index(row_axis, x.ndim)
+        col_axis = normalize_axis_index(col_axis, x.ndim)
         if row_axis == col_axis:
             raise ValueError("Duplicate axes given")
         if ord == 2:
diff --git a/cunumeric/sort.py b/cunumeric/sort.py
--- a/cunumeric/sort.py
+++ b/cunumeric/sort.py
@@ -15,6 +15,7 @@
 from __future__ import annotations
 
 from cunumeric.config import CuNumericOpCode
+from numpy.core.multiarray import normalize_axis_index
 
 from legate.core import types as ty
 
@@ -32,7 +33,7 @@ def sort_flattened(output, input, argsort, stable):
 
 
 def sort_swapped(output, input, argsort, sort_axis, stable):
-    assert sort_axis < input.ndim - 1 and sort_axis >= 0
+    sort_axis = normalize_axis_index(sort_axis, input.ndim)
 
     # swap axes
     swapped = input.swapaxes(sort_axis, input.ndim - 1)
@@ -97,12 +98,10 @@ def sort(output, input, argsort, axis=-1, stable=False):
     else:
         if axis is None:
             axis = 0
-        elif axis < 0:
-            axis = input.ndim + axis
-
-        if axis is not input.ndim - 1:
-            sort_swapped(output, input, argsort, axis, stable)
-
         else:
-            # run actual sort task
+            axis = normalize_axis_index(axis, input.ndim)
+
+        if axis == input.ndim - 1:
             sort_task(output, input, argsort, stable)
+        else:
+            sort_swapped(output, input, argsort, axis, stable)


Test patch: