BUG: fix array_split incorrect behavior with array size bigger MAX_INT32 (#11813)

Fixes #11809.

* BUG: fix array_split incorrect behavior with array size bigger MAX_INT32

* TST: added test for array_split with array size greater MAX_INT32

* addressed review comments

1 files changed, 14 insertions, 0 deletions

diff --git a/numpy/lib/tests/test_shape_base.py b/numpy/lib/tests/test_shape_base.py
index 6d24dd624..6e4cd225d 100644
--- a/numpy/lib/tests/test_shape_base.py
+++ b/numpy/lib/tests/test_shape_base.py
@@ -3,6 +3,8 @@ from __future__ import division, absolute_import, print_function
 import numpy as np
 import warnings
 import functools
+import sys
+import pytest
 
 from numpy.lib.shape_base import (
     apply_along_axis, apply_over_axes, array_split, split, hsplit, dsplit,
@@ -14,6 +16,9 @@ from numpy.testing import (
     )
 
 
+IS_64BIT = sys.maxsize > 2**32
+
+
 def _add_keepdims(func):
     """ hack in keepdims behavior into a function taking an axis """
     @functools.wraps(func)
@@ -403,6 +408,15 @@ class TestArraySplit(object):
         assert_(a.dtype.type is res[-1].dtype.type)
         # perhaps should check higher dimensions
 
+    @pytest.mark.skipif(not IS_64BIT, reason="Needs 64bit platform")
+    def test_integer_split_2D_rows_greater_max_int32(self):
+        a = np.broadcast_to([0], (1 << 32, 2))
+        res = array_split(a, 4)
+        chunk = np.broadcast_to([0], (1 << 30, 2))
+        tgt = [chunk] * 4
+        for i in range(len(tgt)):
+            assert_equal(res[i].shape, tgt[i].shape)
+
     def test_index_split_simple(self):
         a = np.arange(10)
         indices = [1, 5, 7]