diff --git a/dpnp/dpnp_iface_manipulation.py b/dpnp/dpnp_iface_manipulation.py index 6e75caeb7f92..6c078425c0ca 100644 --- a/dpnp/dpnp_iface_manipulation.py +++ b/dpnp/dpnp_iface_manipulation.py @@ -41,7 +41,9 @@ import math import operator import warnings +from typing import NamedTuple +import dpctl import dpctl.tensor as dpt import numpy from dpctl.tensor._numpy_helper import AxisError, normalize_axis_index @@ -54,6 +56,21 @@ from .dpnp_utils import get_usm_allocations from .dpnp_utils.dpnp_utils_pad import dpnp_pad + +class InsertDeleteParams(NamedTuple): + """Parameters used for ``dpnp.delete`` and ``dpnp.insert``.""" + + a: dpnp_array + a_ndim: int + order: str + axis: int + slobj: list + n: int + a_shape: list + exec_q: dpctl.SyclQueue + usm_type: str + + __all__ = [ "append", "array_split", @@ -79,6 +96,7 @@ "flipud", "hsplit", "hstack", + "insert", "matrix_transpose", "moveaxis", "ndim", @@ -120,10 +138,11 @@ def _check_stack_arrays(arrays): ) -def _delete_with_slice(a, obj, axis): +def _delete_with_slice(params, obj, axis): """Utility function for ``dpnp.delete`` when obj is slice.""" - a, a_ndim, order, axis, slobj, n, a_shape = _calc_parameters(a, axis) + a, a_ndim, order, axis, slobj, n, newshape, exec_q, usm_type = params + start, stop, step = obj.indices(n) xr = range(start, stop, step) num_del = len(xr) @@ -137,13 +156,13 @@ def _delete_with_slice(a, obj, axis): start = xr[-1] stop = xr[0] + 1 - a_shape[axis] -= num_del + newshape[axis] -= num_del new = dpnp.empty( - a_shape, - dtype=a.dtype, + newshape, order=order, - sycl_queue=a.sycl_queue, - usm_type=a.usm_type, + dtype=a.dtype, + sycl_queue=exec_q, + usm_type=usm_type, ) # copy initial chunk if start == 0: @@ -166,8 +185,8 @@ def _delete_with_slice(a, obj, axis): keep = dpnp.ones( stop - start, dtype=dpnp.bool, - sycl_queue=a.sycl_queue, - usm_type=a.usm_type, + sycl_queue=exec_q, + usm_type=usm_type, ) keep[: stop - start : step] = False slobj[axis] = slice(start, stop - num_del) @@ -180,10 +199,11 @@ def _delete_with_slice(a, obj, axis): return new -def _delete_without_slice(a, obj, axis, single_value, exec_q, usm_type): +def _delete_without_slice(params, obj, axis, single_value): """Utility function for ``dpnp.delete`` when obj is int or array of int.""" - a, a_ndim, order, axis, slobj, n, a_shape = _calc_parameters(a, axis) + a, a_ndim, order, axis, slobj, n, newshape, exec_q, usm_type = params + if single_value: # optimization for a single value if obj < -n or obj >= n: @@ -193,11 +213,11 @@ def _delete_without_slice(a, obj, axis, single_value, exec_q, usm_type): ) if obj < 0: obj += n - a_shape[axis] -= 1 + newshape[axis] -= 1 new = dpnp.empty( - a_shape, - dtype=a.dtype, + newshape, order=order, + dtype=a.dtype, sycl_queue=exec_q, usm_type=usm_type, ) @@ -229,7 +249,7 @@ def _delete_without_slice(a, obj, axis, single_value, exec_q, usm_type): return new -def _calc_parameters(a, axis): +def _calc_parameters(a, axis, obj, values=None): """Utility function for ``dpnp.delete`` and ``dpnp.insert``.""" a_ndim = a.ndim @@ -246,7 +266,118 @@ def _calc_parameters(a, axis): n = a.shape[axis] a_shape = list(a.shape) - return a, a_ndim, order, axis, slobj, n, a_shape + usm_type, exec_q = get_usm_allocations([a, obj, values]) + + return InsertDeleteParams( + a, a_ndim, order, axis, slobj, n, a_shape, exec_q, usm_type + ) + + +def _insert_array_indices(parameters, indices, values, obj): + """ + Utility function for ``dpnp.insert`` when indices is an array with + multiple elements. + + """ + + a, a_ndim, order, axis, slobj, n, newshape, exec_q, usm_type = parameters + + is_array = isinstance(obj, (dpnp_array, numpy.ndarray, dpt.usm_ndarray)) + if indices.size == 0 and not is_array: + # Can safely cast the empty list to intp + indices = indices.astype(dpnp.intp) + + indices[indices < 0] += n + + numnew = len(indices) + ind_sort = indices.argsort(kind="stable") + indices[ind_sort] += dpnp.arange( + numnew, dtype=indices.dtype, sycl_queue=exec_q, usm_type=usm_type + ) + + newshape[axis] += numnew + old_mask = dpnp.ones( + newshape[axis], dtype=dpnp.bool, sycl_queue=exec_q, usm_type=usm_type + ) + old_mask[indices] = False + + new = dpnp.empty( + newshape, + order=order, + dtype=a.dtype, + sycl_queue=exec_q, + usm_type=usm_type, + ) + slobj2 = [slice(None)] * a_ndim + slobj[axis] = indices + slobj2[axis] = old_mask + new[tuple(slobj)] = values + new[tuple(slobj2)] = a + + return new + + +def _insert_singleton_index(parameters, indices, values, obj): + """ + Utility function for ``dpnp.insert`` when indices is an array with + one element. + + """ + + a, a_ndim, order, axis, slobj, n, newshape, exec_q, usm_type = parameters + + # In dpnp, `.item()` calls `.wait()`, so it is preferred to avoid it + # When possible (i.e. for numpy arrays, lists, etc), it is preferred + # to use `.item()` on a NumPy array + if dpnp.is_supported_array_type(obj): + index = indices.item() + else: + if isinstance(obj, slice): + obj = numpy.arange(*obj.indices(n), dtype=dpnp.intp) + index = numpy.asarray(obj).item() + + if index < -n or index > n: + raise IndexError( + f"index {index} is out of bounds for axis {axis} with size {n}" + ) + if index < 0: + index += n + + # Need to change the dtype of values to input array dtype and update + # its shape to make ``input_arr[..., index, ...] = values`` legal + values = dpnp.array( + values, + copy=None, + ndmin=a_ndim, + dtype=a.dtype, + sycl_queue=exec_q, + usm_type=usm_type, + ) + if indices.ndim == 0: + # numpy.insert behave differently if obj is an scalar or an array + # with one element, so, this change is needed to align with NumPy + values = dpnp.moveaxis(values, 0, axis) + + numnew = values.shape[axis] + newshape[axis] += numnew + new = dpnp.empty( + newshape, + order=order, + dtype=a.dtype, + sycl_queue=exec_q, + usm_type=usm_type, + ) + + slobj[axis] = slice(None, index) + new[tuple(slobj)] = a[tuple(slobj)] + slobj[axis] = slice(index, index + numnew) + new[tuple(slobj)] = values + slobj[axis] = slice(index + numnew, None) + slobj2 = [slice(None)] * a_ndim + slobj2[axis] = slice(index, None) + new[tuple(slobj)] = a[tuple(slobj2)] + + return new def _unique_1d( @@ -1442,14 +1573,10 @@ def delete(arr, obj, axis=None): """ dpnp.check_supported_arrays_type(arr) + params = _calc_parameters(arr, axis, obj) if isinstance(obj, slice): - return _delete_with_slice(arr, obj, axis) - - if dpnp.is_supported_array_type(obj): - usm_type, exec_q = get_usm_allocations([arr, obj]) - else: - usm_type, exec_q = arr.usm_type, arr.sycl_queue + return _delete_with_slice(params, obj, axis) if isinstance(obj, (int, dpnp.integer)) and not isinstance(obj, bool): single_value = True @@ -1457,7 +1584,9 @@ def delete(arr, obj, axis=None): else: single_value = False is_array = isinstance(obj, (dpnp_array, numpy.ndarray, dpt.usm_ndarray)) - indices = dpnp.asarray(obj, sycl_queue=exec_q, usm_type=usm_type) + indices = dpnp.asarray( + obj, sycl_queue=params.exec_q, usm_type=params.usm_type + ) # if `obj` is originally an empty list, after converting it into # an array, it will have float dtype, so we need to change its dtype # to integer. However, if `obj` is originally an empty array with @@ -1467,15 +1596,16 @@ def delete(arr, obj, axis=None): elif indices.size == 1 and indices.dtype.kind in "ui": # For a size 1 integer array we can use the single-value path # (most dtypes, except boolean, should just fail later). - if isinstance(obj, (dpnp_array, dpt.usm_ndarray)): + single_value = True + # In dpnp, `.item()` calls `.wait()`, so it is preferred to avoid it + # When possible (i.e. for numpy arrays, lists, etc), it is + # preferred to use `.item()` on a NumPy array + if dpnp.is_supported_array_type(obj): indices = indices.item() else: indices = numpy.asarray(obj).item() - single_value = True - return _delete_without_slice( - arr, indices, axis, single_value, exec_q, usm_type - ) + return _delete_without_slice(params, indices, axis, single_value) def dsplit(ary, indices_or_sections): @@ -2030,6 +2160,138 @@ def hstack(tup, *, dtype=None, casting="same_kind"): return dpnp.concatenate(arrs, axis=1, dtype=dtype, casting=casting) +def insert(arr, obj, values, axis=None): + """ + Insert values along the given axis before the given indices. + + For full documentation refer to :obj:`numpy.insert`. + + Parameters + ---------- + arr : array_like + Input array. + obj : {slice, int, array-like of ints} + Object that defines the index or indices before which `values` is + inserted. It supports multiple insertions when `obj` is a single + scalar or a sequence with one element (similar to calling insert + multiple times). + values : array_like + Values to insert into `arr`. If the type of `values` is different + from that of `arr`, `values` is converted to the type of `arr`. + `values` should be shaped so that ``arr[..., obj, ...] = values`` + is legal. + axis : {None, int}, optional + Axis along which to insert `values`. If `axis` is ``None`` then `arr` + is flattened first. + Default: ``None``. + + Returns + ------- + out : dpnp.ndarray + A copy of `arr` with `values` inserted. Note that :obj:`dpnp.insert` + does not occur in-place: a new array is returned. If + `axis` is ``None``, `out` is a flattened array. + + See Also + -------- + :obj:`dpnp.append` : Append elements at the end of an array. + :obj:`dpnp.concatenate` : Join a sequence of arrays along an existing axis. + :obj:`dpnp.delete` : Delete elements from an array. + + Notes + ----- + Note that for higher dimensional inserts ``obj=0`` behaves very different + from ``obj=[0]`` just like ``arr[:, 0, :] = values`` is different from + ``arr[:, [0], :] = values``. + + Examples + -------- + >>> import dpnp as np + >>> a = np.array([[1, 1], [2, 2], [3, 3]]) + >>> a + array([[1, 1], + [2, 2], + [3, 3]]) + >>> np.insert(a, 1, 5) + array([1, 5, 1, 2, 2, 3, 3]) + >>> np.insert(a, 1, 5, axis=1) + array([[1, 5, 1], + [2, 5, 2], + [3, 5, 3]]) + + Difference between sequence and scalars: + + >>> np.insert(a, [1], [[1],[2],[3]], axis=1) + array([[1, 1, 1], + [2, 2, 2], + [3, 3, 3]]) + >>> np.array_equal(np.insert(a, 1, [1, 2, 3], axis=1), + ... np.insert(a, [1], [[1],[2],[3]], axis=1)) + array(True) + + >>> b = a.flatten() + >>> b + array([1, 1, 2, 2, 3, 3]) + >>> np.insert(b, [2, 2], [5, 6]) + array([1, 1, 5, 6, 2, 2, 3, 3]) + + >>> np.insert(b, slice(2, 4), [5, 6]) + array([1, 1, 5, 2, 6, 2, 3, 3]) + + >>> np.insert(b, [2, 2], [7.13, False]) # dtype casting + array([1, 1, 7, 0, 2, 2, 3, 3]) + + >>> x = np.arange(8).reshape(2, 4) + >>> idx = (1, 3) + >>> np.insert(x, idx, 999, axis=1) + array([[ 0, 999, 1, 2, 999, 3], + [ 4, 999, 5, 6, 999, 7]]) + + """ + + dpnp.check_supported_arrays_type(arr) + params = _calc_parameters(arr, axis, obj, values) + + if isinstance(obj, slice): + # turn it into a range object + indices = dpnp.arange( + *obj.indices(params.n), + dtype=dpnp.intp, + sycl_queue=params.exec_q, + usm_type=params.usm_type, + ) + else: + # need to copy obj, because indices will be changed in-place + indices = dpnp.copy( + obj, sycl_queue=params.exec_q, usm_type=params.usm_type + ) + if indices.dtype == dpnp.bool: + warnings.warn( + "In the future insert will treat boolean arrays and array-likes" + " as a boolean index instead of casting it to integers", + FutureWarning, + stacklevel=2, + ) + indices = indices.astype(dpnp.intp) + # TODO: Code after warning period: + # if indices.ndim != 1: + # raise ValueError( + # "boolean array argument `obj` to insert must be " + # "one-dimensional" + # ) + # indices = dpnp.nonzero(indices)[0] + elif indices.ndim > 1: + raise ValueError( + "index array argument `obj` to insert must be one-dimensional " + "or scalar" + ) + + if indices.size == 1: + return _insert_singleton_index(params, indices, values, obj) + + return _insert_array_indices(params, indices, values, obj) + + def matrix_transpose(x, /): """ Transposes a matrix (or a stack of matrices) `x`. diff --git a/tests/test_manipulation.py b/tests/test_manipulation.py index 2bab4d754f9e..7d53d7efcbeb 100644 --- a/tests/test_manipulation.py +++ b/tests/test_manipulation.py @@ -524,6 +524,190 @@ def test_3D_array(self): _compare_results(result, expected) +class TestInsert: + @pytest.mark.usefixtures("suppress_complex_warning") + @pytest.mark.parametrize( + "obj", [slice(0, 4, 2), 3, [2, 3]], ids=["slice", "scalar", "list"] + ) + @pytest.mark.parametrize("dt1", get_all_dtypes(no_none=True)) + @pytest.mark.parametrize("dt2", get_all_dtypes(no_none=True)) + def test_dtype(self, obj, dt1, dt2): + a = numpy.array([0, 1, 2, 3, 4, 5], dtype=dt1) + a_dp = dpnp.array(a) + + values = numpy.array(3, dtype=dt2) + + expected = numpy.insert(a, obj, values) + result = dpnp.insert(a_dp, obj, values) + assert result.dtype == dt1 + assert_array_equal(result, expected) + + @pytest.mark.parametrize("order", ["C", "F"]) + def test_order(self, order): + a = numpy.arange(10).reshape(2, 5, order=order) + a_dp = dpnp.array(a) + + expected = numpy.insert(a, slice(3, None), -1, axis=1) + result = dpnp.insert(a_dp, slice(3, None), -1, axis=1) + + assert_equal(result.flags.c_contiguous, expected.flags.c_contiguous) + assert_equal(result.flags.f_contiguous, expected.flags.f_contiguous) + + @pytest.mark.parametrize( + "obj", + [numpy.array([2]), dpnp.array([0, 2]), dpnp.asarray([1])], + ) + @pytest.mark.parametrize( + "values", + [numpy.array([-1]), dpnp.array([-2, -3])], + ) + def test_ndarray_obj_values(self, obj, values): + a = numpy.array([1, 2, 3]) + ia = dpnp.array(a) + + values_np = ( + values if isinstance(values, numpy.ndarray) else values.asnumpy() + ) + obj_np = obj if isinstance(obj, numpy.ndarray) else obj.asnumpy() + expected = numpy.insert(a, obj_np, values_np) + result = dpnp.insert(ia, obj, values) + assert_equal(result, expected) + + @pytest.mark.filterwarnings("ignore::FutureWarning") + @pytest.mark.parametrize( + "obj", + [True, [False], numpy.array([True] * 4), [True, False, True, False]], + ) + def test_boolean_obj(self, obj): + a = numpy.array([1, 2, 3]) + ia = dpnp.array(a) + assert_equal(dpnp.insert(ia, obj, 9), numpy.insert(a, obj, 9)) + + def test_1D_array(self): + a = numpy.array([1, 2, 3]) + ia = dpnp.array(a) + + assert_equal(dpnp.insert(ia, 0, 1), numpy.insert(a, 0, 1)) + assert_equal(dpnp.insert(ia, 3, 1), numpy.insert(a, 3, 1)) + assert_equal( + dpnp.insert(ia, -1, [1, 2, 3]), numpy.insert(a, -1, [1, 2, 3]) + ) + assert_equal( + dpnp.insert(ia, [1, -1, 3], 9), numpy.insert(a, [1, -1, 3], 9) + ) + + expected = numpy.insert(a, [1, 1, 1], [1, 2, 3]) + result = dpnp.insert(ia, [1, 1, 1], [1, 2, 3]) + assert_equal(result, expected) + + expected = numpy.insert(a, slice(-1, None, -1), 9) + result = dpnp.insert(ia, slice(-1, None, -1), 9) + assert_equal(result, expected) + + expected = numpy.insert(a, slice(None, 1, None), 9) + result = dpnp.insert(ia, slice(None, 1, None), 9) + assert_equal(result, expected) + + expected = numpy.insert(a, [-1, 1, 3], [7, 8, 9]) + result = dpnp.insert(ia, [-1, 1, 3], [7, 8, 9]) + assert_equal(result, expected) + + b = numpy.array([0, 1], dtype=numpy.float32) + ib = dpnp.array(b) + assert_equal(dpnp.insert(ib, 0, ib[0]), numpy.insert(b, 0, b[0])) + assert_equal(dpnp.insert(ib, [], []), numpy.insert(b, [], [])) + + def test_ND_array(self): + a = numpy.array([[1, 1, 1]]) + ia = dpnp.array(a) + + assert_equal(dpnp.insert(ia, 0, [1]), numpy.insert(a, 0, [1])) + assert_equal( + dpnp.insert(ia, 0, 2, axis=0), numpy.insert(a, 0, 2, axis=0) + ) + assert_equal( + dpnp.insert(ia, 2, 2, axis=1), numpy.insert(a, 2, 2, axis=1) + ) + expected = numpy.insert(a, 0, [2, 2, 2], axis=0) + result = dpnp.insert(ia, 0, [2, 2, 2], axis=0) + assert_equal(result, expected) + + a = numpy.array([[1, 1], [2, 2], [3, 3]]) + ia = dpnp.array(a) + expected = numpy.insert(a, [1], [[1], [2], [3]], axis=1) + result = dpnp.insert(ia, [1], [[1], [2], [3]], axis=1) + assert_equal(result, expected) + + expected = numpy.insert(a, [1], [1, 2, 3], axis=1) + result = dpnp.insert(ia, [1], [1, 2, 3], axis=1) + assert_equal(result, expected) + # scalars behave differently + expected = numpy.insert(a, 1, [1, 2, 3], axis=1) + result = dpnp.insert(ia, 1, [1, 2, 3], axis=1) + assert_equal(result, expected) + + expected = numpy.insert(a, 1, [[1], [2], [3]], axis=1) + result = dpnp.insert(ia, 1, [[1], [2], [3]], axis=1) + assert_equal(result, expected) + + a = numpy.arange(4).reshape(2, 2) + ia = dpnp.array(a) + expected = numpy.insert(a[:, :1], 1, a[:, 1], axis=1) + result = dpnp.insert(ia[:, :1], 1, ia[:, 1], axis=1) + assert_equal(result, expected) + + expected = numpy.insert(a[:1, :], 1, a[1, :], axis=0) + result = dpnp.insert(ia[:1, :], 1, ia[1, :], axis=0) + assert_equal(result, expected) + + # negative axis value + a = numpy.arange(24).reshape((2, 3, 4)) + ia = dpnp.array(a) + expected = numpy.insert(a, 1, a[:, :, 3], axis=-1) + result = dpnp.insert(ia, 1, ia[:, :, 3], axis=-1) + assert_equal(result, expected) + + expected = numpy.insert(a, 1, a[:, 2, :], axis=-2) + result = dpnp.insert(ia, 1, ia[:, 2, :], axis=-2) + assert_equal(result, expected) + + # invalid axis value + assert_raises(AxisError, dpnp.insert, ia, 1, ia[:, 2, :], axis=3) + assert_raises(AxisError, dpnp.insert, ia, 1, ia[:, 2, :], axis=-4) + + def test_index_array_copied(self): + a = dpnp.array([0, 1, 2]) + x = dpnp.array([1, 1, 1]) + dpnp.insert(a, x, [3, 4, 5]) + assert_equal(x, dpnp.array([1, 1, 1])) + + def test_error(self): + a = dpnp.array([0, 1, 2]) + + # index float + with pytest.raises(IndexError): + dpnp.insert(a, dpnp.array([1.0, 2.0]), [10, 20]) + with pytest.raises(IndexError): + dpnp.insert(a, dpnp.array([], dtype=dpnp.float32), []) + + # index 2d + with pytest.raises(ValueError): + dpnp.insert(a, dpnp.array([[1.0], [2.0]]), [10, 20]) + + # incorrect axis + a = dpnp.array(1) + with pytest.raises(AxisError): + dpnp.insert(a, [], 2, axis=0) + with pytest.raises(TypeError): + dpnp.insert(a, [], 2, axis="nonsense") + + @pytest.mark.parametrize("idx", [4, -4]) + def test_index_out_of_bounds(self, idx): + a = dpnp.array([0, 1, 2]) + with pytest.raises(IndexError, match="out of bounds"): + dpnp.insert(a, [idx], [3, 4]) + + # array_split has more comprehensive test of splitting. # only do simple test on hsplit, vsplit, and dsplit class TestHsplit: diff --git a/tests/test_sycl_queue.py b/tests/test_sycl_queue.py index 277c1b6ee10e..a4e969f68430 100644 --- a/tests/test_sycl_queue.py +++ b/tests/test_sycl_queue.py @@ -2078,6 +2078,70 @@ def test_obj_ndarray(self, device): assert_sycl_queue_equal(result.sycl_queue, y.sycl_queue) +class TestInsert: + @pytest.mark.parametrize( + "obj", + [slice(None, None, 2), 3, [2, 3]], + ids=["slice", "scalar", "list"], + ) + @pytest.mark.parametrize( + "device", + valid_devices, + ids=[device.filter_string for device in valid_devices], + ) + def test_basic(self, obj, device): + x = dpnp.arange(5, device=device) + result = dpnp.insert(x, obj, 3) + assert_sycl_queue_equal(result.sycl_queue, x.sycl_queue) + + @pytest.mark.parametrize( + "obj", + [slice(None, None, 3), 3, [2, 3]], + ids=["slice", "scalar", "list"], + ) + @pytest.mark.parametrize( + "device", + valid_devices, + ids=[device.filter_string for device in valid_devices], + ) + def test_values_ndarray(self, obj, device): + x = dpnp.arange(5, device=device) + y = dpnp.array([1, 4], device=device) + result = dpnp.insert(x, obj, y) + + assert_sycl_queue_equal(result.sycl_queue, x.sycl_queue) + assert_sycl_queue_equal(result.sycl_queue, y.sycl_queue) + + @pytest.mark.parametrize("values", [-2, [-1, -2]], ids=["scalar", "list"]) + @pytest.mark.parametrize( + "device", + valid_devices, + ids=[device.filter_string for device in valid_devices], + ) + def test_obj_ndarray(self, values, device): + x = dpnp.arange(5, device=device) + y = dpnp.array([1, 4], device=device) + result = dpnp.insert(x, y, values) + + assert_sycl_queue_equal(result.sycl_queue, x.sycl_queue) + assert_sycl_queue_equal(result.sycl_queue, y.sycl_queue) + + @pytest.mark.parametrize( + "device", + valid_devices, + ids=[device.filter_string for device in valid_devices], + ) + def test_obj_values_ndarray(self, device): + x = dpnp.arange(5, device=device) + y = dpnp.array([1, 4], device=device) + z = dpnp.array([-1, -3], device=device) + result = dpnp.insert(x, y, z) + + assert_sycl_queue_equal(result.sycl_queue, x.sycl_queue) + assert_sycl_queue_equal(result.sycl_queue, y.sycl_queue) + assert_sycl_queue_equal(result.sycl_queue, z.sycl_queue) + + @pytest.mark.parametrize( "func,data1", [ diff --git a/tests/test_usm_type.py b/tests/test_usm_type.py index 2f1b2ca7bd65..aa76443dad40 100644 --- a/tests/test_usm_type.py +++ b/tests/test_usm_type.py @@ -904,6 +904,81 @@ def test_einsum(usm_type): assert result.usm_type == usm_type +class TestInsert: + @pytest.mark.parametrize( + "usm_type", list_of_usm_types, ids=list_of_usm_types + ) + @pytest.mark.parametrize( + "obj", + [slice(None, None, 2), 3, [2, 3]], + ids=["slice", "scalar", "list"], + ) + def test_bacis(self, usm_type, obj): + x = dp.arange(5, usm_type=usm_type) + result = dp.insert(x, obj, 3) + + assert x.usm_type == usm_type + assert result.usm_type == usm_type + + @pytest.mark.parametrize( + "obj", + [slice(None, None, 3), 3, [2, 3]], + ids=["slice", "scalar", "list"], + ) + @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_values_ndarray(self, obj, usm_type_x, usm_type_y): + x = dp.arange(5, usm_type=usm_type_x) + y = dp.array([1, 4], usm_type=usm_type_y) + z = dp.insert(x, obj, y) + + assert x.usm_type == usm_type_x + assert y.usm_type == usm_type_y + assert z.usm_type == du.get_coerced_usm_type([usm_type_x, usm_type_y]) + + @pytest.mark.parametrize("values", [-2, [-1, -2]], ids=["scalar", "list"]) + @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_obj_ndarray(self, values, usm_type_x, usm_type_y): + x = dp.arange(5, usm_type=usm_type_x) + y = dp.array([1, 4], usm_type=usm_type_y) + z = dp.insert(x, y, values) + + assert x.usm_type == usm_type_x + assert y.usm_type == usm_type_y + assert z.usm_type == du.get_coerced_usm_type([usm_type_x, usm_type_y]) + + @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 + ) + @pytest.mark.parametrize( + "usm_type_z", list_of_usm_types, ids=list_of_usm_types + ) + def test_obj_values_ndarray(self, usm_type_x, usm_type_y, usm_type_z): + x = dp.arange(5, usm_type=usm_type_x) + y = dp.array([1, 4], usm_type=usm_type_y) + z = dp.array([-1, -3], usm_type=usm_type_z) + res = dp.insert(x, y, z) + + assert x.usm_type == usm_type_x + assert y.usm_type == usm_type_y + assert z.usm_type == usm_type_z + assert res.usm_type == du.get_coerced_usm_type( + [usm_type_x, usm_type_y, usm_type_z] + ) + + @pytest.mark.parametrize("func", ["take", "take_along_axis"]) @pytest.mark.parametrize("usm_type_x", list_of_usm_types, ids=list_of_usm_types) @pytest.mark.parametrize(