diff --git a/dpnp/dpnp_iface_manipulation.py b/dpnp/dpnp_iface_manipulation.py index ee6812f44c79..be9765c3420b 100644 --- a/dpnp/dpnp_iface_manipulation.py +++ b/dpnp/dpnp_iface_manipulation.py @@ -50,6 +50,7 @@ __all__ = [ "append", + "array_split", "asarray_chkfinite", "asfarray", "atleast_1d", @@ -62,11 +63,13 @@ "concat", "concatenate", "copyto", + "dsplit", "dstack", "expand_dims", "flip", "fliplr", "flipud", + "hsplit", "hstack", "moveaxis", "ndim", @@ -80,6 +83,7 @@ "row_stack", "shape", "size", + "split", "squeeze", "stack", "swapaxes", @@ -87,6 +91,7 @@ "transpose", "trim_zeros", "unique", + "vsplit", "vstack", ] @@ -409,7 +414,7 @@ def asarray_chkfinite( >>> np.asarray_chkfinite(a, dtype=np.float32) array([1., 2.]) - Raises ``ValueError`` if array_like contains Nans or Infs. + Raises ``ValueError`` if array_like contains NaNs or Infs. >>> a = [1, 2, np.inf] >>> try: @@ -447,6 +452,89 @@ def asarray_chkfinite( return a +def array_split(ary, indices_or_sections, axis=0): + """ + Split an array into multiple sub-arrays. + + Please refer to the :obj:`dpnp.split` documentation. The only difference + between these functions is that ``dpnp.array_split`` allows + `indices_or_sections` to be an integer that does *not* equally divide the + axis. For an array of length l that should be split into n sections, it + returns ``l % n`` sub-arrays of size ``l//n + 1`` and the rest of size + ``l//n``. + + For full documentation refer to :obj:`numpy.array_split`. + + Parameters + ---------- + ary : {dpnp.ndarray, usm_ndarray} + Array to be divided into sub-arrays. + indices_or_sections : {int, sequence of ints} + If `indices_or_sections` is an integer, N, and array length is l, it + returns ``l % n`` sub-arrays of size ``l//n + 1`` and the rest of size + ``l//n``. + + If `indices_or_sections` is a sequence of sorted integers, the entries + indicate where along `axis` the array is split. + axis : int, optional + The axis along which to split. + Default: ``0``. + + Returns + ------- + sub-arrays : list of dpnp.ndarray + A list of sub arrays. Each array is a view of the corresponding input + array. + + See Also + -------- + :obj:`dpnp.split` : Split array into multiple sub-arrays of equal size. + + Examples + -------- + >>> import dpnp as np + >>> x = np.arange(8.0) + >>> np.array_split(x, 3) + [array([0., 1., 2.]), array([3., 4., 5.]), array([6., 7.])] + + >>> x = np.arange(9) + >>> np.array_split(x, 4) + [array([0, 1, 2]), array([3, 4]), array([5, 6]), array([7, 8])] + + """ + + dpnp.check_supported_arrays_type(ary) + n_tot = ary.shape[axis] + try: + # handle array case. + n_sec = len(indices_or_sections) + 1 + div_points = [0] + list(indices_or_sections) + [n_tot] + except TypeError: + # indices_or_sections is a scalar, not an array. + n_sec = int(indices_or_sections) + if n_sec <= 0: + raise ValueError("number sections must be larger than 0.") from None + n_each_sec, extras = numpy.divmod(n_tot, n_sec) + section_sizes = ( + [0] + extras * [n_each_sec + 1] + (n_sec - extras) * [n_each_sec] + ) + div_points = dpnp.array( + section_sizes, + dtype=dpnp.intp, + usm_type=ary.usm_type, + sycl_queue=ary.sycl_queue, + ).cumsum() + + sub_arys = [] + sary = dpnp.swapaxes(ary, axis, 0) + for i in range(n_sec): + st = div_points[i] + end = div_points[i + 1] + sub_arys.append(dpnp.swapaxes(sary[st:end], axis, 0)) + + return sub_arys + + def asfarray(a, dtype=None, *, device=None, usm_type=None, sycl_queue=None): """ Return an array converted to a float type. @@ -1121,6 +1209,75 @@ def copyto(dst, src, casting="same_kind", where=True): dst_usm[mask_usm] = src_usm[mask_usm] +def dsplit(ary, indices_or_sections): + """ + Split array into multiple sub-arrays along the 3rd axis (depth). + + Please refer to the :obj:`dpnp.split` documentation. ``dsplit`` + is equivalent to ``split`` with ``axis=2``, the array is always + split along the third axis provided the array dimension is greater than + or equal to 3. + + For full documentation refer to :obj:`numpy.dsplit`. + + Parameters + ---------- + ary : {dpnp.ndarray, usm_ndarray} + Array to be divided into sub-arrays. + indices_or_sections : {int, sequence of ints} + If `indices_or_sections` is an integer, N, the array will be divided + into N equal arrays along the third axis. If such a split is not + possible, an error is raised. + If `indices_or_sections` is a sequence of sorted integers, the entries + indicate where along the third axis the array is split. + + Returns + ------- + sub-arrays : list of dpnp.ndarray + A list of sub arrays. Each array is a view of the corresponding input + array. + + See Also + -------- + :obj:`dpnp.split` : Split array into multiple sub-arrays of equal size. + + Examples + -------- + >>> import dpnp as np + >>> x = np.arange(16.0).reshape(2, 2, 4) + >>> x + array([[[ 0., 1., 2., 3.], + [ 4., 5., 6., 7.]], + [[ 8., 9., 10., 11.], + [12., 13., 14., 15.]]]) + >>> np.dsplit(x, 2) + [array([[[ 0., 1.], + [ 4., 5.]], + [[ 8., 9.], + [12., 13.]]]), + array([[[ 2., 3.], + [ 6., 7.]], + [[10., 11.], + [14., 15.]]])] + >>> np.dsplit(x, np.array([3, 6])) + [array([[[ 0., 1., 2.], + [ 4., 5., 6.]], + [[ 8., 9., 10.], + [12., 13., 14.]]]), + array([[[ 3.], + [ 7.]], + [[11.], + [15.]]]), + array([])] + + """ + + dpnp.check_supported_arrays_type(ary) + if ary.ndim < 3: + raise ValueError("dsplit only works on arrays of 3 or more dimensions") + return split(ary, indices_or_sections, 2) + + def dstack(tup): """ Stack arrays in sequence depth wise (along third axis). @@ -1445,6 +1602,99 @@ def flipud(m): return m[::-1, ...] +def hsplit(ary, indices_or_sections): + """ + Split an array into multiple sub-arrays horizontally (column-wise). + + Please refer to the :obj:`dpnp.split` documentation. ``hsplit`` + is equivalent to ``dpnp.split`` with ``axis=1``, the array is always + split along the second axis except for 1-D arrays, where it is split at + ``axis=0``. + + For full documentation refer to :obj:`numpy.hsplit`. + + Parameters + ---------- + ary : {dpnp.ndarray, usm_ndarray} + Array to be divided into sub-arrays. + indices_or_sections : {int, sequence of ints} + If `indices_or_sections` is an integer, N, the array will be divided + into N equal arrays along the second axis except for 1-D arrays, where + it is split at the first axis. If such a split is not possible, + an error is raised. + If `indices_or_sections` is a sequence of sorted integers, the entries + indicate where along the second axis the array is split. For 1-D arrays, + the entries indicate where along the first axis the array is split. + + Returns + ------- + sub-arrays : list of dpnp.ndarray + A list of sub arrays. Each array is a view of the corresponding input + array. + + See Also + -------- + :obj:`dpnp.split` : Split array into multiple sub-arrays of equal size. + + Examples + -------- + >>> import dpnp as np + >>> x = np.arange(16.0).reshape(4, 4) + >>> x + array([[ 0., 1., 2., 3.], + [ 4., 5., 6., 7.], + [ 8., 9., 10., 11.], + [12., 13., 14., 15.]]) + >>> np.hsplit(x, 2) + [array([[ 0., 1.], + [ 4., 5.], + [ 8., 9.], + [12., 13.]]), + array([[ 2., 3.], + [ 6., 7.], + [10., 11.], + [14., 15.]])] + >>> np.hsplit(x, np.array([3, 6])) + [array([[ 0., 1., 2.], + [ 4., 5., 6.], + [ 8., 9., 10.], + [12., 13., 14.]]), + array([[ 3.], + [ 7.], + [11.], + [15.]]), + array([])] + + With a higher dimensional array the split is still along the second axis. + + >>> x = np.arange(8.0).reshape(2, 2, 2) + >>> x + array([[[0., 1.], + [2., 3.]], + [[4., 5.], + [6., 7.]]]) + >>> np.hsplit(x, 2) + [array([[[0., 1.]], + [[4., 5.]]]), + array([[[2., 3.]], + [[6., 7.]]])] + + With a 1-D array, the split is along axis 0. + + >>> x = np.array([0, 1, 2, 3, 4, 5]) + >>> np.hsplit(x, 2) + [array([0, 1, 2]), array([3, 4, 5])] + + """ + + dpnp.check_supported_arrays_type(ary) + if ary.ndim == 0: + raise ValueError("hsplit only works on arrays of 1 or more dimensions") + if ary.ndim > 1: + return split(ary, indices_or_sections, 1) + return split(ary, indices_or_sections, 0) + + def hstack(tup, *, dtype=None, casting="same_kind"): """ Stack arrays in sequence horizontally (column wise). @@ -2045,6 +2295,96 @@ def size(a, axis=None): return numpy.size(a, axis) +def split(ary, indices_or_sections, axis=0): + """ + Split an array into multiple sub-arrays as views into `ary`. + + For full documentation refer to :obj:`numpy.split`. + + Parameters + ---------- + ary : {dpnp.ndarray, usm_ndarray} + Array to be divided into sub-arrays. + indices_or_sections : {int, sequence of ints} + If `indices_or_sections` is an integer, N, the array will be divided + into N equal arrays along `axis`. If such a split is not possible, + an error is raised. + + If `indices_or_sections` is a sequence of sorted integers, the entries + indicate where along `axis` the array is split. For example, + ``[2, 3]`` would, for ``axis=0``, result in + + - ary[:2] + - ary[2:3] + - ary[3:] + + If an index exceeds the dimension of the array along `axis`, + an empty sub-array is returned correspondingly. + axis : int, optional + The axis along which to split. + Default: ``0``. + + Returns + ------- + sub-arrays : list of dpnp.ndarray + A list of sub arrays. Each array is a view of the corresponding input + array. + + Raises + ------ + ValueError + If `indices_or_sections` is given as an integer, but + a split does not result in equal division. + + See Also + -------- + :obj:`dpnp.array_split` : Split an array into multiple sub-arrays of equal + or near-equal size. Does not raise an exception if an + equal division cannot be made. + :obj:`dpnp.hsplit` : Split array into multiple sub-arrays horizontally + (column-wise). + :obj:`dpnp.vsplit` : Split array into multiple sub-arrays vertically + (row wise). + :obj:`dpnp.dsplit` : Split array into multiple sub-arrays along the 3rd + axis (depth). + :obj:`dpnp.concatenate` : Join a sequence of arrays along an existing axis. + :obj:`dpnp.stack` : Join a sequence of arrays along a new axis. + :obj:`dpnp.hstack` : Stack arrays in sequence horizontally (column wise). + :obj:`dpnp.vstack` : Stack arrays in sequence vertically (row wise). + :obj:`dpnp.dstack` : Stack arrays in sequence depth wise + (along third dimension). + + Examples + -------- + >>> import dpnp as np + >>> x = np.arange(9.0) + >>> np.split(x, 3) + [array([0., 1., 2.]), array([3., 4., 5.]), array([6., 7., 8.])] + + >>> x = np.arange(8.0) + >>> np.split(x, [3, 5, 6, 10]) + [array([0., 1., 2.]), array([3., 4.]), array([5.]), array([6., 7.]), \ + array([])] + + """ + + dpnp.check_supported_arrays_type(ary) + if ary.ndim <= axis: + raise IndexError("Axis exceeds ndim") + + try: + len(indices_or_sections) + except TypeError: + if ary.shape[axis] % indices_or_sections != 0: + raise ValueError( + "indices_or_sections must divide the size along the axes.\n" + "If you want to split the array into non-equally-sized " + "arrays, use array_split instead." + ) from None + + return array_split(ary, indices_or_sections, axis) + + def squeeze(a, /, axis=None): """ Removes singleton dimensions (axes) from array `a`. @@ -2631,6 +2971,80 @@ def unique( return _unpack_tuple(result) +def vsplit(ary, indices_or_sections): + """ + Split an array into multiple sub-arrays vertically (row-wise). + + Please refer to the :obj:`dpnp.split` documentation. ``vsplit`` + is equivalent to ``split`` with ``axis=0``(default), the array + is always split along the first axis regardless of the array dimension. + + For full documentation refer to :obj:`numpy.vsplit`. + + Parameters + ---------- + ary : {dpnp.ndarray, usm_ndarray} + Array to be divided into sub-arrays. + indices_or_sections : {int, sequence of ints} + If `indices_or_sections` is an integer, N, the array will be divided + into N equal arrays along the first axis. If such a split is not + possible, an error is raised. + If `indices_or_sections` is a sequence of sorted integers, the entries + indicate where along the first axis the array is split. + + Returns + ------- + sub-arrays : list of dpnp.ndarray + A list of sub arrays. Each array is a view of the corresponding input + array. + + See Also + -------- + :obj:`dpnp.split` : Split array into multiple sub-arrays of equal size. + + Examples + -------- + >>> import dpnp as np + >>> x = np.arange(16.0).reshape(4, 4) + >>> x + array([[ 0., 1., 2., 3.], + [ 4., 5., 6., 7.], + [ 8., 9., 10., 11.], + [12., 13., 14., 15.]]) + >>> np.vsplit(x, 2) + [array([[0., 1., 2., 3.], + [4., 5., 6., 7.]]), + array([[ 8., 9., 10., 11.], + [12., 13., 14., 15.]])] + >>> np.vsplit(x, np.array([3, 6])) + [array([[ 0., 1., 2., 3.], + [ 4., 5., 6., 7.], + [ 8., 9., 10., 11.]]), + array([[12., 13., 14., 15.]]), + array([], shape=(0, 4), dtype=float64)] + + With a higher dimensional array the split is still along the first axis. + + >>> x = np.arange(8.0).reshape(2, 2, 2) + >>> x + array([[[0., 1.], + [2., 3.]], + [[4., 5.], + [6., 7.]]]) + >>> np.vsplit(x, 2) + [array([[[0., 1.], + [2., 3.]]]), + array([[[4., 5.], + [6., 7.]]])] + + """ + + dpnp.check_supported_arrays_type(ary) + if ary.ndim < 2: + raise ValueError("vsplit only works on arrays of 2 or more dimensions") + return split(ary, indices_or_sections, 0) + + def vstack(tup, *, dtype=None, casting="same_kind"): """ Stack arrays in sequence vertically (row wise). diff --git a/tests/test_manipulation.py b/tests/test_manipulation.py index a8e64717b9df..b9afeef7f290 100644 --- a/tests/test_manipulation.py +++ b/tests/test_manipulation.py @@ -5,6 +5,7 @@ from numpy.testing import assert_array_equal, assert_raises import dpnp +from tests.third_party.cupy import testing from .helper import ( assert_dtype_allclose, @@ -29,6 +30,15 @@ testdata += [([1j, -1j, 1 - 2j], dtype) for dtype in get_complex_dtypes()] +def _compare_results(result, expected): + """Compare lists of arrays.""" + if len(result) != len(expected): + raise ValueError("Iterables have different lengths") + + for x, y in zip(result, expected): + assert_array_equal(x, y) + + @pytest.mark.parametrize("in_obj, out_dtype", testdata) def test_copyto_dtype(in_obj, out_dtype): ndarr = numpy.array(in_obj) @@ -186,6 +196,223 @@ def test_axis(self, dtype1, dtype2): assert_array_equal(result, expected) +class TestArraySplit: + @pytest.mark.parametrize("xp", [numpy, dpnp]) + def test_error(self, xp): + # 0 split is not allowed + a = xp.arange(10) + assert_raises(ValueError, xp.array_split, a, 0) + + # invalid indices_or_sections + a = xp.arange(10) + assert_raises(TypeError, xp.array_split, a, "wrong") + + # non-integer sequence + a = xp.arange(10) + assert_raises(TypeError, xp.array_split, a, [3, 5.0]) + + # not 1D array + a = xp.arange(10) + indices = dpnp.array([[1, 5], [7, 9]]) + assert_raises(ValueError, xp.array_split, a, indices) + + @pytest.mark.parametrize( + "indices", + [ + 1, + 2, + 3.0, + dpnp.int64(5), + dpnp.int32(5), + dpnp.array(6), + numpy.array(7), + numpy.int32(5), + 9, + 10, + 11, + ], + ) + def test_integer_split(self, indices): + a = numpy.arange(10) + a_dp = dpnp.array(a) + + expected = numpy.array_split(a, indices) + result = dpnp.array_split(a_dp, indices) + _compare_results(result, expected) + + def test_integer_split_2D_rows(self): + a = numpy.array([numpy.arange(10), numpy.arange(10)]) + a_dp = dpnp.array(a) + expected = numpy.array_split(a, 3, axis=0) + result = dpnp.array_split(a_dp, 3, axis=0) + _compare_results(result, expected) + assert a.dtype.type is result[-1].dtype.type + + # Same thing for manual splits: + expected = numpy.array_split(a, [0, 1], axis=0) + result = dpnp.array_split(a_dp, [0, 1], axis=0) + _compare_results(result, expected) + assert a.dtype.type is result[-1].dtype.type + + def test_integer_split_2D_cols(self): + a = numpy.array([numpy.arange(10), numpy.arange(10)]) + a_dp = dpnp.array(a) + expected = numpy.array_split(a, 3, axis=-1) + result = dpnp.array_split(a_dp, 3, axis=-1) + _compare_results(result, expected) + + @testing.slow + def test_integer_split_2D_rows_greater_max_int32(self): + a = numpy.broadcast_to([0], (1 << 32, 2)) + a_dp = dpnp.broadcast_to(dpnp.array([0]), (1 << 32, 2)) + expected = numpy.array_split(a, 4) + result = dpnp.array_split(a_dp, 4) + _compare_results(result, expected) + + @pytest.mark.parametrize( + "indices", + [[1, 5, 7], (1, 5, 7), dpnp.array([1, 5, 7]), numpy.array([1, 5, 7])], + ) + def test_index_split_simple(self, indices): + a = numpy.arange(10) + a_dp = dpnp.array(a) + expected = numpy.array_split(a, indices, axis=-1) + result = dpnp.array_split(a_dp, indices, axis=-1) + _compare_results(result, expected) + + def test_index_split_low_bound(self): + a = numpy.arange(10) + a_dp = dpnp.array(a) + indices = [0, 5, 7] + expected = numpy.array_split(a, indices, axis=-1) + result = dpnp.array_split(a_dp, indices, axis=-1) + _compare_results(result, expected) + + def test_index_split_high_bound(self): + a = numpy.arange(10) + a_dp = dpnp.array(a) + indices = [0, 5, 7, 10, 12] + expected = numpy.array_split(a, indices, axis=-1) + result = dpnp.array_split(a_dp, indices, axis=-1) + _compare_results(result, expected) + + +class TestSplit: + # The split function is essentially the same as array_split, + # except that it test if splitting will result in an + # equal split. Only test for this case. + def test_equal_split(self): + a = numpy.arange(10) + a_dp = dpnp.array(a) + + expected = numpy.split(a, 2) + result = dpnp.split(a_dp, 2) + _compare_results(result, expected) + + @pytest.mark.parametrize("xp", [numpy, dpnp]) + def test_unequal_split(self, xp): + a = xp.arange(10) + assert_raises(ValueError, xp.split, a, 3) + + @pytest.mark.parametrize("xp", [numpy, dpnp]) + def test_error(self, xp): + # axis out of range + a = xp.arange(9) + assert_raises(IndexError, xp.split, a, 3, axis=1) + + @pytest.mark.parametrize( + "indices", + [ + 2, + 3.0, + dpnp.int64(5), + dpnp.int32(5), + dpnp.array(6), + numpy.array(7), + numpy.int32(5), + ], + ) + def test_integer_split(self, indices): + a = numpy.arange(10) + a_dp = dpnp.array(a) + + expected = numpy.array_split(a, indices) + result = dpnp.array_split(a_dp, indices) + _compare_results(result, expected) + + +# array_split has more comprehensive test of splitting. +# only do simple test on hsplit, vsplit, and dsplit +class TestHsplit: + @pytest.mark.parametrize("xp", [numpy, dpnp]) + def test_error(self, xp): + # 0D array + a = xp.array(1) + assert_raises(ValueError, xp.hsplit, a, 2) + + def test_1D_array(self): + a = numpy.array([1, 2, 3, 4]) + a_dp = dpnp.array(a) + + expected = numpy.hsplit(a, 2) + result = dpnp.hsplit(a_dp, 2) + _compare_results(result, expected) + + def test_2D_array(self): + a = numpy.array([[1, 2, 3, 4], [1, 2, 3, 4]]) + a_dp = dpnp.array(a) + + expected = numpy.hsplit(a, 2) + result = dpnp.hsplit(a_dp, 2) + _compare_results(result, expected) + + +class TestVsplit: + @pytest.mark.parametrize("xp", [numpy, dpnp]) + def test_error(self, xp): + # 0D array + a = xp.array(1) + assert_raises(ValueError, xp.vsplit, a, 2) + + # 1D array + a = xp.array([1, 2, 3, 4]) + assert_raises(ValueError, xp.vsplit, a, 2) + + def test_2D_array(self): + a = numpy.array([[1, 2, 3, 4], [1, 2, 3, 4]]) + a_dp = dpnp.array(a) + + expected = numpy.vsplit(a, 2) + result = dpnp.vsplit(a_dp, 2) + _compare_results(result, expected) + + +class TestDsplit: + @pytest.mark.parametrize("xp", [numpy, dpnp]) + def test_error(self, xp): + # 0D array + a = xp.array(1) + assert_raises(ValueError, xp.dsplit, a, 2) + + # 1D array + a = xp.array([1, 2, 3, 4]) + assert_raises(ValueError, xp.dsplit, a, 2) + + # 2D array + a = xp.array([[1, 2, 3, 4], [1, 2, 3, 4]]) + assert_raises(ValueError, xp.dsplit, a, 2) + + def test_3D_array(self): + a = numpy.array( + [[[1, 2, 3, 4], [1, 2, 3, 4]], [[1, 2, 3, 4], [1, 2, 3, 4]]] + ) + a_dp = dpnp.array(a) + + expected = numpy.dsplit(a, 2) + result = dpnp.dsplit(a_dp, 2) + _compare_results(result, expected) + + class TestAsarrayCheckFinite: @pytest.mark.parametrize("dtype", get_all_dtypes()) def test_basic(self, dtype): diff --git a/tests/test_sycl_queue.py b/tests/test_sycl_queue.py index e88a09fba95d..eb54065e1df7 100644 --- a/tests/test_sycl_queue.py +++ b/tests/test_sycl_queue.py @@ -1951,7 +1951,6 @@ def test_concat_stack(func, data1, data2, device): result = getattr(dpnp, func)((x1, x2)) assert_allclose(result, expected) - assert_sycl_queue_equal(result.sycl_queue, x1.sycl_queue) assert_sycl_queue_equal(result.sycl_queue, x2.sycl_queue) @@ -1975,6 +1974,36 @@ def test_append(device): assert_sycl_queue_equal(result.sycl_queue, x2.sycl_queue) +@pytest.mark.parametrize( + "func,data1", + [ + pytest.param("array_split", [1, 2, 3, 4]), + pytest.param("split", [1, 2, 3, 4]), + pytest.param("hsplit", [1, 2, 3, 4]), + pytest.param( + "dsplit", + [[[1, 2, 3, 4], [1, 2, 3, 4]], [[1, 2, 3, 4], [1, 2, 3, 4]]], + ), + pytest.param("vsplit", [[1, 2, 3, 4], [1, 2, 3, 4]]), + ], +) +@pytest.mark.parametrize( + "device", + valid_devices, + ids=[device.filter_string for device in valid_devices], +) +def test_split(func, data1, device): + x1_orig = numpy.array(data1) + x1 = dpnp.array(data1, device=device) + expected = getattr(numpy, func)(x1_orig, 2) + result = getattr(dpnp, func)(x1, 2) + + assert_allclose(result[0], expected[0]) + assert_allclose(result[1], expected[1]) + assert_sycl_queue_equal(result[0].sycl_queue, x1.sycl_queue) + assert_sycl_queue_equal(result[1].sycl_queue, x1.sycl_queue) + + @pytest.mark.parametrize( "device_x", valid_devices, diff --git a/tests/test_usm_type.py b/tests/test_usm_type.py index 9d8eee83b0d8..f83b3df5f02f 100644 --- a/tests/test_usm_type.py +++ b/tests/test_usm_type.py @@ -768,6 +768,29 @@ def test_concat_stack(func, data1, data2, usm_type_x, usm_type_y): assert z.usm_type == du.get_coerced_usm_type([usm_type_x, usm_type_y]) +@pytest.mark.parametrize( + "func,data1", + [ + pytest.param("array_split", [1, 2, 3, 4]), + pytest.param("split", [1, 2, 3, 4]), + pytest.param("hsplit", [1, 2, 3, 4]), + pytest.param( + "dsplit", + [[[1, 2, 3, 4], [1, 2, 3, 4]], [[1, 2, 3, 4], [1, 2, 3, 4]]], + ), + pytest.param("vsplit", [[1, 2, 3, 4], [1, 2, 3, 4]]), + ], +) +@pytest.mark.parametrize("usm_type", list_of_usm_types, ids=list_of_usm_types) +def test_split(func, data1, usm_type): + x = dp.array(data1, usm_type=usm_type) + y = getattr(dp, func)(x, 2) + + assert x.usm_type == usm_type + assert y[0].usm_type == usm_type + assert y[1].usm_type == usm_type + + @pytest.mark.parametrize("usm_type_x", list_of_usm_types, ids=list_of_usm_types) @pytest.mark.parametrize("usm_type_y", list_of_usm_types, ids=list_of_usm_types) def test_append(usm_type_x, usm_type_y): diff --git a/tests/third_party/cupy/manipulation_tests/test_split.py b/tests/third_party/cupy/manipulation_tests/test_split.py new file mode 100644 index 000000000000..c7e7a59ae9b1 --- /dev/null +++ b/tests/third_party/cupy/manipulation_tests/test_split.py @@ -0,0 +1,96 @@ +import unittest + +import dpnp as cupy +from tests.third_party.cupy import testing + + +class TestSplit(unittest.TestCase): + @testing.numpy_cupy_array_equal() + def test_array_split1(self, xp): + a = testing.shaped_arange((3, 11), xp) + return xp.array_split(a, 4, 1) + + @testing.numpy_cupy_array_equal() + def test_array_split2(self, xp): + a = testing.shaped_arange((3, 11), xp) + return xp.array_split(a, 4, -1) + + @testing.numpy_cupy_array_equal() + def test_array_split_empty_array(self, xp): + a = testing.shaped_arange((5, 0), xp) + return xp.array_split(a, [2, 4], 0) + + @testing.numpy_cupy_array_equal() + def test_array_split_empty_sections(self, xp): + a = testing.shaped_arange((3, 11), xp) + return xp.array_split(a, []) + + @testing.numpy_cupy_array_equal() + def test_array_split_out_of_bound1(self, xp): + a = testing.shaped_arange((2, 3), xp) + return xp.array_split(a, [3]) + + @testing.numpy_cupy_array_equal() + 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_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_equal() + def test_array_split_non_divisible(self, xp): + a = testing.shaped_arange((5, 3), xp) + return xp.array_split(a, 4) + + @testing.numpy_cupy_array_equal() + def test_dsplit(self, xp): + a = testing.shaped_arange((3, 3, 12), xp) + return xp.dsplit(a, 4) + + @testing.numpy_cupy_array_equal() + def test_hsplit_vectors(self, xp): + a = testing.shaped_arange((12,), xp) + return xp.hsplit(a, 4) + + @testing.numpy_cupy_array_equal() + def test_hsplit(self, xp): + a = testing.shaped_arange((3, 12), xp) + return xp.hsplit(a, 4) + + @testing.numpy_cupy_array_equal() + def test_split_by_sections1(self, xp): + a = testing.shaped_arange((3, 11), xp) + return xp.split(a, (2, 4, 9), 1) + + @testing.numpy_cupy_array_equal() + def test_split_by_sections2(self, xp): + a = testing.shaped_arange((3, 11), xp) + return xp.split(a, (2, 4, 9), -1) + + @testing.numpy_cupy_array_equal() + def test_split_by_sections3(self, xp): + a = testing.shaped_arange((3, 11), xp) + return xp.split(a, (-9, 4, -2), 1) + + @testing.numpy_cupy_array_equal() + def test_split_out_of_bound1(self, xp): + a = testing.shaped_arange((2, 3), xp) + return xp.split(a, [3]) + + @testing.numpy_cupy_array_equal() + def test_split_out_of_bound2(self, xp): + a = testing.shaped_arange((0,), xp) + return xp.split(a, [1]) + + @testing.numpy_cupy_array_equal() + def test_split_unordered_sections(self, xp): + a = testing.shaped_arange((5,), xp) + return xp.split(a, [4, 2]) + + @testing.numpy_cupy_array_equal() + def test_vsplit(self, xp): + a = testing.shaped_arange((12, 3), xp) + return xp.vsplit(a, 4)