regularize axis makes sure now that its type is either int or None

src/awkward/_regularize.py

@@ -51,11 +51,19 @@ def is_non_string_like_sequence(obj) -> bool:
     return not isinstance(obj, (str, bytes)) and isinstance(obj, Sequence)
 
 
-def regularize_axis(axis: Any) -> int | Any:
+def regularize_axis(axis: Any, none_allowed: bool = True) -> int | None:
     """
-    This function's purpose is to convert "0" to 0, "1" to 1, etc., but leave any other value as it is.
+    This function's main purpose is to convert [np,cp,...].array(0) to 0.
     """
-    try:
-        return int(axis)
-    except (TypeError, ValueError):
-        return axis
+    if is_integer_like(axis):
+        regularized_axis = int(axis)
+    else:
+        regularized_axis = axis
+    cond = is_integer(regularized_axis)
+    msg = f"'axis' must be an integer, not {axis!r}"
+    if none_allowed:
+        cond = cond or regularized_axis is None
+        msg = f"'axis' must be an integer or None, not {axis!r}"
+    if not cond:
+        raise TypeError(msg)
+    return regularized_axis

src/awkward/operations/ak_all.py

@@ -13,7 +13,7 @@
     _remove_named_axis,
 )
 from awkward._nplikes.numpy_like import NumpyMetadata
-from awkward._regularize import is_integer, regularize_axis
+from awkward._regularize import regularize_axis
 
 __all__ = ("all",)
 
@@ -76,8 +76,6 @@ def _impl(array, axis, keepdims, mask_identity, highlevel, behavior, attrs):
     with HighLevelContext(behavior=behavior, attrs=attrs) as ctx:
         layout = ctx.unwrap(array, allow_record=False, primitive_policy="error")
 
-    axis = regularize_axis(axis)
-
     # Handle named axis
     named_axis = _get_named_axis(ctx)
     # Step 1: Normalize named axis to positional axis
@@ -93,8 +91,7 @@ def _impl(array, axis, keepdims, mask_identity, highlevel, behavior, attrs):
             total=layout.minmax_depth[1],
         )
 
-    if not is_integer(axis) and axis is not None:
-        raise TypeError(f"'axis' must be an integer or None by now, not {axis!r}")
+    axis = regularize_axis(axis, none_allowed=True)
 
     reducer = ak._reducers.All()
 

src/awkward/operations/ak_any.py

@@ -13,7 +13,7 @@
     _remove_named_axis,
 )
 from awkward._nplikes.numpy_like import NumpyMetadata
-from awkward._regularize import is_integer, regularize_axis
+from awkward._regularize import regularize_axis
 
 __all__ = ("any",)
 
@@ -76,8 +76,6 @@ def _impl(array, axis, keepdims, mask_identity, highlevel, behavior, attrs):
     with HighLevelContext(behavior=behavior, attrs=attrs) as ctx:
         layout = ctx.unwrap(array, allow_record=False, primitive_policy="error")
 
-    axis = regularize_axis(axis)
-
     # Handle named axis
     named_axis = _get_named_axis(ctx)
     # Step 1: Normalize named axis to positional axis
@@ -93,8 +91,7 @@ def _impl(array, axis, keepdims, mask_identity, highlevel, behavior, attrs):
             total=layout.minmax_depth[1],
         )
 
-    if not is_integer(axis) and axis is not None:
-        raise TypeError(f"'axis' must be an integer or None by now, not {axis!r}")
+    axis = regularize_axis(axis, none_allowed=True)
 
     reducer = ak._reducers.Any()
 

src/awkward/operations/ak_argcartesian.py

@@ -7,7 +7,6 @@
 import awkward as ak
 from awkward._dispatch import high_level_function
 from awkward._nplikes.numpy_like import NumpyMetadata
-from awkward._regularize import regularize_axis
 
 __all__ = ("argcartesian",)
 
@@ -107,8 +106,6 @@ def argcartesian(
 
 
 def _impl(arrays, axis, nested, parameters, with_name, highlevel, behavior, attrs):
-    axis = regularize_axis(axis)
-
     if isinstance(arrays, Mapping):
         index_arrays = {n: ak.local_index(x, axis) for n, x in arrays.items()}
     else:

src/awkward/operations/ak_argcombinations.py

@@ -5,6 +5,7 @@
 import awkward as ak
 from awkward._dispatch import high_level_function
 from awkward._layout import HighLevelContext
+from awkward._namedaxis import _get_named_axis, _named_axis_to_positional_axis
 from awkward._nplikes.numpy_like import NumpyMetadata
 from awkward._regularize import regularize_axis
 
@@ -93,15 +94,20 @@ def _impl(
     behavior,
     attrs,
 ):
-    axis = regularize_axis(axis)
-
     if parameters is None:
         parameters = {}
     else:
         parameters = dict(parameters)
     if with_name is not None:
         parameters["__record__"] = with_name
 
+    # Handle named axis
+    named_axis = _get_named_axis(array)
+    # Step 1: Normalize named axis to positional axis
+    axis = _named_axis_to_positional_axis(named_axis, axis)
+
+    axis = regularize_axis(axis, none_allowed=False)
+
     if axis < 0:
         raise ValueError("the 'axis' for argcombinations must be non-negative")
     else:

src/awkward/operations/ak_argmax.py

@@ -13,7 +13,7 @@
     _remove_named_axis,
 )
 from awkward._nplikes.numpy_like import NumpyMetadata
-from awkward._regularize import is_integer, regularize_axis
+from awkward._regularize import regularize_axis
 
 __all__ = ("argmax", "nanargmax")
 
@@ -141,8 +141,6 @@ def _impl(array, axis, keepdims, mask_identity, highlevel, behavior, attrs):
     with HighLevelContext(behavior=behavior, attrs=attrs) as ctx:
         layout = ctx.unwrap(array, allow_record=False, primitive_policy="error")
 
-    axis = regularize_axis(axis)
-
     # Handle named axis
     named_axis = _get_named_axis(ctx)
     # Step 1: Normalize named axis to positional axis
@@ -158,8 +156,7 @@ def _impl(array, axis, keepdims, mask_identity, highlevel, behavior, attrs):
             total=layout.minmax_depth[1],
         )
 
-    if not is_integer(axis) and axis is not None:
-        raise TypeError(f"'axis' must be an integer or None by now, not {axis!r}")
+    axis = regularize_axis(axis, none_allowed=True)
 
     reducer = ak._reducers.ArgMax()
 

src/awkward/operations/ak_argmin.py

@@ -13,7 +13,7 @@
     _remove_named_axis,
 )
 from awkward._nplikes.numpy_like import NumpyMetadata
-from awkward._regularize import is_integer, regularize_axis
+from awkward._regularize import regularize_axis
 
 __all__ = ("argmin", "nanargmin")
 
@@ -138,8 +138,6 @@ def _impl(array, axis, keepdims, mask_identity, highlevel, behavior, attrs):
     with HighLevelContext(behavior=behavior, attrs=attrs) as ctx:
         layout = ctx.unwrap(array, allow_record=False, primitive_policy="error")
 
-    axis = regularize_axis(axis)
-
     # Handle named axis
     named_axis = _get_named_axis(ctx)
     # Step 1: Normalize named axis to positional axis
@@ -155,8 +153,7 @@ def _impl(array, axis, keepdims, mask_identity, highlevel, behavior, attrs):
             total=layout.minmax_depth[1],
         )
 
-    if not is_integer(axis) and axis is not None:
-        raise TypeError(f"'axis' must be an integer or None by now, not {axis!r}")
+    axis = regularize_axis(axis, none_allowed=True)
 
     reducer = ak._reducers.ArgMin()
 

src/awkward/operations/ak_argsort.py

@@ -11,7 +11,7 @@
     _named_axis_to_positional_axis,
 )
 from awkward._nplikes.numpy_like import NumpyMetadata
-from awkward._regularize import is_integer, regularize_axis
+from awkward._regularize import regularize_axis
 
 __all__ = ("argsort",)
 
@@ -77,15 +77,12 @@ def _impl(array, axis, ascending, stable, highlevel, behavior, attrs):
     with HighLevelContext(behavior=behavior, attrs=attrs) as ctx:
         layout = ctx.unwrap(array, allow_record=False, primitive_policy="error")
 
-    axis = regularize_axis(axis)
-
     # Handle named axis
     named_axis = _get_named_axis(ctx)
     # Step 1: Normalize named axis to positional axis
     axis = _named_axis_to_positional_axis(named_axis, axis)
 
-    if not is_integer(axis):
-        raise TypeError(f"'axis' must be an integer by now, not {axis!r}")
+    axis = regularize_axis(axis, none_allowed=False)
 
     out = ak._do.argsort(layout, axis, ascending, stable)
 

src/awkward/operations/ak_cartesian.py

@@ -251,11 +251,10 @@ def _impl(arrays, axis, nested, parameters, with_name, highlevel, behavior, attr
             #   use strategy "unify" (see: awkward._namedaxis)
             out_named_axis = reduce(_unify_named_axis, map(_get_named_axis, arrays))
 
-    axis = regularize_axis(axis)
-
     # Handle named axis
     # Step 1: Normalize named axis to positional axis
     axis = _named_axis_to_positional_axis(out_named_axis, axis)
+    axis = regularize_axis(axis, none_allowed=False)
     max_ndim = max(layout.minmax_depth[1] for layout in layouts)
 
     if with_name is not None:

src/awkward/operations/ak_combinations.py

@@ -221,13 +221,13 @@ def _impl(
     with HighLevelContext(behavior=behavior, attrs=attrs) as ctx:
         layout = ctx.unwrap(array, allow_record=False, primitive_policy="error")
 
-    axis = regularize_axis(axis)
-
     # Handle named axis
     named_axis = _get_named_axis(ctx)
     # Step 1: Normalize named axis to positional axis
     axis = _named_axis_to_positional_axis(named_axis, axis)
 
+    axis = regularize_axis(axis, none_allowed=False)
+
     if with_name is None:
         pass
     elif parameters is None:

src/awkward/operations/ak_concatenate.py

@@ -128,12 +128,11 @@ def _impl(arrays, axis, mergebool, highlevel, behavior, attrs):
             )
         )
 
-    axis = regularize_axis(axis)
-
     # Handle named axis
     merged_named_axis = reduce(_unify_named_axis, map(_get_named_axis, arrays))
     # Step 1: normalize named axis to positional axis
     axis = _named_axis_to_positional_axis(merged_named_axis, axis)
+    axis = regularize_axis(axis, none_allowed=False)
     # Step 2: propagate named axis from input to output,
     #   use strategy "unify" (see: awkward._namedaxis)
     out_named_axis = merged_named_axis

src/awkward/operations/ak_corr.py

@@ -10,7 +10,7 @@
     ensure_same_backend,
     maybe_highlevel_to_lowlevel,
 )
-from awkward._namedaxis import _get_named_axis
+from awkward._namedaxis import _get_named_axis, _is_valid_named_axis
 from awkward._nplikes import ufuncs
 from awkward._nplikes.numpy_like import NumpyMetadata
 from awkward._regularize import regularize_axis
@@ -88,7 +88,10 @@ def corr(
 
 
 def _impl(x, y, weight, axis, keepdims, mask_identity, highlevel, behavior, attrs):
-    axis = regularize_axis(axis)
+    if _is_valid_named_axis(axis):
+        raise NotImplementedError("named axis not yet supported for ak.corr")
+
+    axis = regularize_axis(axis, none_allowed=True)
 
     with HighLevelContext(behavior=behavior, attrs=attrs) as ctx:
         x_layout, y_layout, weight_layout = ensure_same_backend(

src/awkward/operations/ak_count.py

@@ -12,7 +12,7 @@
     _remove_named_axis,
 )
 from awkward._nplikes.numpy_like import NumpyMetadata
-from awkward._regularize import is_integer, regularize_axis
+from awkward._regularize import regularize_axis
 
 __all__ = ("count",)
 
@@ -118,8 +118,6 @@ def _impl(array, axis, keepdims, mask_identity, highlevel, behavior, attrs):
     with HighLevelContext(behavior=behavior, attrs=attrs) as ctx:
         layout = ctx.unwrap(array, allow_record=False, primitive_policy="error")
 
-    axis = regularize_axis(axis)
-
     # Handle named axis
     named_axis = _get_named_axis(ctx)
     # Step 1: Normalize named axis to positional axis
@@ -135,8 +133,7 @@ def _impl(array, axis, keepdims, mask_identity, highlevel, behavior, attrs):
             total=layout.minmax_depth[1],
         )
 
-    if not is_integer(axis) and axis is not None:
-        raise TypeError(f"'axis' must be an integer or None by now, not {axis!r}")
+    axis = regularize_axis(axis, none_allowed=True)
 
     reducer = ak._reducers.Count()
 

src/awkward/operations/ak_count_nonzero.py

@@ -12,7 +12,7 @@
     _remove_named_axis,
 )
 from awkward._nplikes.numpy_like import NumpyMetadata
-from awkward._regularize import is_integer, regularize_axis
+from awkward._regularize import regularize_axis
 
 __all__ = ("count_nonzero",)
 
@@ -77,8 +77,6 @@ def _impl(array, axis, keepdims, mask_identity, highlevel, behavior, attrs):
     with HighLevelContext(behavior=behavior, attrs=attrs) as ctx:
         layout = ctx.unwrap(array, allow_record=False, primitive_policy="error")
 
-    axis = regularize_axis(axis)
-
     # Handle named axis
     named_axis = _get_named_axis(ctx)
     # Step 1: Normalize named axis to positional axis
@@ -94,8 +92,7 @@ def _impl(array, axis, keepdims, mask_identity, highlevel, behavior, attrs):
             total=layout.minmax_depth[1],
         )
 
-    if not is_integer(axis) and axis is not None:
-        raise TypeError(f"'axis' must be an integer or None by now, not {axis!r}")
+    axis = regularize_axis(axis, none_allowed=True)
 
     with HighLevelContext(behavior=behavior, attrs=attrs) as ctx:
         layout = ctx.unwrap(array, allow_record=False, primitive_policy="error")

src/awkward/operations/ak_covar.py

@@ -10,7 +10,7 @@
     ensure_same_backend,
     maybe_highlevel_to_lowlevel,
 )
-from awkward._namedaxis import _get_named_axis
+from awkward._namedaxis import _get_named_axis, _is_valid_named_axis
 from awkward._nplikes.numpy_like import NumpyMetadata
 from awkward._regularize import regularize_axis
 
@@ -85,7 +85,9 @@ def covar(
 
 
 def _impl(x, y, weight, axis, keepdims, mask_identity, highlevel, behavior, attrs):
-    axis = regularize_axis(axis)
+    if _is_valid_named_axis(axis):
+        raise NotImplementedError("named axis not yet supported for ak.covar")
+    axis = regularize_axis(axis, none_allowed=True)
 
     with HighLevelContext(behavior=behavior, attrs=attrs) as ctx:
         x_layout, y_layout, weight_layout = ensure_same_backend(

src/awkward/operations/ak_drop_none.py

@@ -10,7 +10,7 @@
     _named_axis_to_positional_axis,
 )
 from awkward._nplikes.numpy_like import NumpyMetadata
-from awkward._regularize import is_integer, regularize_axis
+from awkward._regularize import regularize_axis
 from awkward.errors import AxisError
 
 __all__ = ("drop_none",)
@@ -72,15 +72,12 @@ def _impl(array, axis, highlevel, behavior, attrs):
     with HighLevelContext(behavior=behavior, attrs=attrs) as ctx:
         layout = ctx.unwrap(array, allow_record=False, primitive_policy="error")
 
-    axis = regularize_axis(axis)
-
     # Handle named axis
     named_axis = _get_named_axis(ctx)
     # Step 1: Normalize named axis to positional axis
     axis = _named_axis_to_positional_axis(named_axis, axis)
 
-    if not is_integer(axis) and axis is not None:
-        raise TypeError(f"'axis' must be an integer or None by now, not {axis!r}")
+    axis = regularize_axis(axis, none_allowed=True)
 
     if axis is None:
         # if the outer layout is_option, drop_nones without affecting offsets