Merge pull request #2856 from emcastillo/bp-2815-v6-fix-split-unordered

[backport] Fix split and array_split with unordered indices supplied

cupy/core/_routines_manipulation.pyx

@@ -434,7 +434,7 @@ cpdef array_split(ndarray ary, indices_or_sections, Py_ssize_t axis):
         stride = 0
     for index in indices:
         index = min(index, size)
-        shape[axis] = index - prev
+        shape[axis] = max(index - prev, 0)
         v = ary.view()
         v.data = ary.data + prev * stride
         # TODO(niboshi): Confirm update_x_contiguity flags

tests/cupy_tests/manipulation_tests/test_split.py

@@ -38,6 +38,11 @@ def test_array_split_out_of_bound2(self, xp):
         a = testing.shaped_arange((0,), xp)
         return xp.array_split(a, [1])
 
+    @testing.numpy_cupy_array_list_equal()
+    def test_array_split_unordered_sections(self, xp):
+        a = testing.shaped_arange((5,), xp)
+        return xp.array_split(a, [4, 2])
+
     @testing.numpy_cupy_array_list_equal()
     def test_array_split_non_divisible(self, xp):
         a = testing.shaped_arange((5, 3), xp)
@@ -84,6 +89,11 @@ def test_split_out_of_bound2(self, xp):
         a = testing.shaped_arange((0,), xp)
         return xp.split(a, [1])
 
+    @testing.numpy_cupy_array_list_equal()
+    def test_split_unordered_sections(self, xp):
+        a = testing.shaped_arange((5,), xp)
+        return xp.split(a, [4, 2])
+
     @testing.numpy_cupy_array_list_equal()
     def test_vsplit(self, xp):
         a = testing.shaped_arange((12, 3), xp)