Repository: numba/numba

Issue:
Segfault in np.split with negative axis
- [x] I have tried using the latest released version of Numba (most recent is
 visible in the change log (https://github.com/numba/numba/blob/main/CHANGE_LOG).
- [x] I have included a self contained code sample to reproduce the problem.
  i.e. it's possible to run as 'python bug.py'.

```python
import numba
import numpy as np

@numba.njit
def foo():
    return np.split(np.zeros((10, 10)), np.array([5]), axis=-1)

foo()  # segfault
```

According to gdb the segfault occurs [in this function](https://github.com/numba/numba/blob/61ec1fd0f69aeadece218dccf4c39ebc5c7dfbc4/numba/core/runtime/_nrt_python.c#L332)

numba version '0.55.2'




Discussion:
Thanks for the report, the above reproduces. The issue is in here:
https://github.com/numba/numba/blob/88fef5b350f5b9b26b26ec963147c9dbaf3d0998/numba/np/arrayobj.py#L5642-L5653
I think the negative axis should be wrapped around before handing it to `tuple_setitem`, something like:
```diff
diff --git a/numba/np/arrayobj.py b/numba/np/arrayobj.py
--- a/numba/np/arrayobj.py
+++ b/numba/np/arrayobj.py
@@ -5643,6 +5643,8 @@ def np_array_split(ary, indices_or_sections, axis=0):
             slice_tup = build_full_slice_tuple(ary.ndim)
             out = list()
             prev = 0
+            if axis < 0:
+                axis += ary.ndim
             for cur in indices_or_sections:
                 idx = tuple_setitem(slice_tup, axis, slice(prev, cur))
                 out.append(ary[idx])
```
should fix.
That would still segfault for `axis = -3`.
Maybe there should be a helper function that checks axis arguments and raises errors?

```python
def normalize_axis(ndim, axis):
    if axis < 0:
        axis += ndim
    if axis < 0 or axis >= ndim:
        raise IndexError("Invalid axis argument")
    return axis
```

Ah, maybe this one looks right? https://github.com/numba/numba/blob/main/numba/np/arrayobj.py#L4990

I'm not really sure about how this should be used in the split implementation.
> Ah, maybe this one looks right? https://github.com/numba/numba/blob/main/numba/np/arrayobj.py#L4990
> 
> I'm not really sure about how this should be used in the split implementation.

It'll need wrapping in a `numba.extending.intrinsic` and then it can just be called within the `np_array_split` implementation.

Resolving patch:
diff --git a/numba/np/arrayobj.py b/numba/np/arrayobj.py
--- a/numba/np/arrayobj.py
+++ b/numba/np/arrayobj.py
@@ -26,6 +26,7 @@
                                  impl_ret_new_ref, impl_ret_untracked,
                                  RefType)
 from numba.core.typing import signature
+from numba.core.types import StringLiteral
 from numba.core.extending import (register_jitable, overload, overload_method,
                                   intrinsic)
 from numba.misc import quicksort, mergesort
@@ -270,6 +271,34 @@ def update_array_info(aryty, array):
                                           get_itemsize(context, aryty))
 
 
+def normalize_axis(func_name, arg_name, ndim, axis):
+    """Constrain axis values to valid positive values."""
+    raise NotImplementedError()
+
+
+@overload(normalize_axis)
+def normalize_axis_overloads(func_name, arg_name, ndim, axis):
+    if not isinstance(func_name, StringLiteral):
+        raise errors.TypingError("func_name must be a str literal.")
+    if not isinstance(arg_name, StringLiteral):
+        raise errors.TypingError("arg_name must be a str literal.")
+
+    msg = (
+        f"{func_name.literal_value}: Argument {arg_name.literal_value} "
+        "out of bounds for dimensions of the array"
+    )
+
+    def impl(func_name, arg_name, ndim, axis):
+        if axis < 0:
+            axis += ndim
+        if axis < 0 or axis >= ndim:
+            raise ValueError(msg)
+
+        return axis
+
+    return impl
+
+
 @lower_builtin('getiter', types.Buffer)
 def getiter_array(context, builder, sig, args):
     [arrayty] = sig.args
@@ -5939,7 +5968,8 @@ def impl(ary, indices_or_sections, axis=0):
     ):
         def impl(ary, indices_or_sections, axis=0):
             slice_tup = build_full_slice_tuple(ary.ndim)
-            out = list()
+            axis = normalize_axis("np.split", "axis", ary.ndim, axis)
+            out = []
             prev = 0
             for cur in indices_or_sections:
                 idx = tuple_setitem(slice_tup, axis, slice(prev, cur))
@@ -5956,7 +5986,8 @@ def impl(ary, indices_or_sections, axis=0):
     ):
         def impl(ary, indices_or_sections, axis=0):
             slice_tup = build_full_slice_tuple(ary.ndim)
-            out = list()
+            axis = normalize_axis("np.split", "axis", ary.ndim, axis)
+            out = []
             prev = 0
             for cur in literal_unroll(indices_or_sections):
                 idx = tuple_setitem(slice_tup, axis, slice(prev, cur))
@@ -6222,12 +6253,8 @@ def numpy_swapaxes(arr, axis1, axis2):
     axes_list = tuple(range(ndim))
 
     def impl(arr, axis1, axis2):
-        if axis1 >= ndim or abs(axis1) > ndim:
-            raise ValueError('The second argument "axis1" is out of bounds '
-                             'for array of given dimension')
-        if axis2 >= ndim or abs(axis2) > ndim:
-            raise ValueError('The third argument "axis2" is out of bounds '
-                             'for array of given dimension')
+        axis1 = normalize_axis("np.swapaxes", "axis1", ndim, axis1)
+        axis2 = normalize_axis("np.swapaxes", "axis2", ndim, axis2)
 
         # to ensure tuple_setitem support of negative values
         if axis1 < 0:
@@ -6253,8 +6280,7 @@ def _take_along_axis_impl(
 
     # Wrap axis, it's used in tuple_setitem so must be (axis >= 0) to ensure
     # the GEP is in bounds.
-    if axis < 0:
-        axis = arr.ndim + axis
+    axis = normalize_axis("np.take_along_axis", "axis", arr.ndim, axis)
 
     # Broadcast the two arrays to matching shapes:
     arr_shape = list(arr.shape)


Test patch:
diff --git a/numba/tests/test_np_functions.py b/numba/tests/test_np_functions.py
--- a/numba/tests/test_np_functions.py
+++ b/numba/tests/test_np_functions.py
@@ -2708,6 +2708,10 @@ def args_variations():
             yield a, [1, 3]
             yield a, [1, 3], 1
             yield a, [1, 3], 2
+            yield a, [1], -1
+            yield a, [1], -2
+            yield a, [1], -3
+            yield a, np.array([], dtype=np.int64), 0
 
             a = np.arange(100).reshape(2, -1)
             yield a, 1
@@ -2768,6 +2772,11 @@ def test_split_basic(self):
             str(raises.exception)
         )
 
+        with self.assertRaises(ValueError) as raises:
+            njit(split)(np.ones(5), [3], axis=-3)
+        self.assertIn("np.split: Argument axis out of bounds",
+                      str(raises.exception))
+
     def test_roll_basic(self):
         pyfunc = roll
         cfunc = jit(nopython=True)(pyfunc)
@@ -4651,14 +4660,14 @@ def test_swapaxes_exception(self):
         with self.assertRaises(ValueError) as raises:
             cfunc(np.arange(4), 1, 0)
 
-        self.assertIn('The second argument "axis1" is out of bounds for array'
-                      ' of given dimension', str(raises.exception))
+        self.assertIn('np.swapaxes: Argument axis1 out of bounds',
+                      str(raises.exception))
 
         with self.assertRaises(ValueError) as raises:
             cfunc(np.arange(8).reshape(2, 4), 0, -3)
 
-        self.assertIn('The third argument "axis2" is out of bounds for array'
-                      ' of given dimension', str(raises.exception))
+        self.assertIn('np.swapaxes: Argument axis2 out of bounds',
+                      str(raises.exception))
 
     def test_take_along_axis(self):
         a = np.arange(24).reshape((3, 1, 4, 2))