diff --git a/dpnp/dpnp_iface_mathematical.py b/dpnp/dpnp_iface_mathematical.py index 98f4c6c8fa8..43f2b0b57dc 100644 --- a/dpnp/dpnp_iface_mathematical.py +++ b/dpnp/dpnp_iface_mathematical.py @@ -308,6 +308,30 @@ def _gradient_num_diff_edges( ) +def _process_ediff1d_args(arg, arg_name, ary_dtype, ary_sycl_queue, usm_type): + """Process the argument for ediff1d.""" + if not dpnp.is_supported_array_type(arg): + arg = dpnp.asarray(arg, usm_type=usm_type, sycl_queue=ary_sycl_queue) + else: + usm_type = dpu.get_coerced_usm_type([usm_type, arg.usm_type]) + # check that arrays have the same allocation queue + if dpu.get_execution_queue([ary_sycl_queue, arg.sycl_queue]) is None: + raise dpu.ExecutionPlacementError( + f"ary and {arg_name} must be allocated on the same SYCL queue" + ) + + if not dpnp.can_cast(arg, ary_dtype, casting="same_kind"): + raise TypeError( + f"dtype of {arg_name} must be compatible " + "with input ary under the `same_kind` rule." + ) + + if arg.ndim > 1: + arg = dpnp.ravel(arg) + + return arg, usm_type + + _ABS_DOCSTRING = """ Calculates the absolute value for each element `x_i` of input array `x`. @@ -1332,52 +1356,30 @@ def ediff1d(ary, to_end=None, to_begin=None): return ary[1:] - ary[:-1] ary_dtype = ary.dtype - ary_usm_type = ary.usm_type ary_sycl_queue = ary.sycl_queue + usm_type = ary.usm_type if to_begin is None: l_begin = 0 else: - if not dpnp.is_supported_array_type(to_begin): - to_begin = dpnp.asarray( - to_begin, usm_type=ary_usm_type, sycl_queue=ary_sycl_queue - ) - if not dpnp.can_cast(to_begin, ary_dtype, casting="same_kind"): - raise TypeError( - "dtype of `to_begin` must be compatible " - "with input `ary` under the `same_kind` rule." - ) - - to_begin_ndim = to_begin.ndim - - if to_begin_ndim > 1: - to_begin = dpnp.ravel(to_begin) - + to_begin, usm_type = _process_ediff1d_args( + to_begin, "to_begin", ary_dtype, ary_sycl_queue, usm_type + ) l_begin = to_begin.size if to_end is None: l_end = 0 else: - if not dpnp.is_supported_array_type(to_end): - to_end = dpnp.asarray( - to_end, usm_type=ary_usm_type, sycl_queue=ary_sycl_queue - ) - if not dpnp.can_cast(to_end, ary_dtype, casting="same_kind"): - raise TypeError( - "dtype of `to_end` must be compatible " - "with input `ary` under the `same_kind` rule." - ) - - to_end_ndim = to_end.ndim - - if to_end_ndim > 1: - to_end = dpnp.ravel(to_end) - + to_end, usm_type = _process_ediff1d_args( + to_end, "to_end", ary_dtype, ary_sycl_queue, usm_type + ) l_end = to_end.size # calculating using in place operation l_diff = max(len(ary) - 1, 0) - result = dpnp.empty_like(ary, shape=l_diff + l_begin + l_end) + result = dpnp.empty_like( + ary, shape=l_diff + l_begin + l_end, usm_type=usm_type + ) if l_begin > 0: result[:l_begin] = to_begin diff --git a/tests/test_mathematical.py b/tests/test_mathematical.py index b78371033ad..b2a90f23477 100644 --- a/tests/test_mathematical.py +++ b/tests/test_mathematical.py @@ -2151,6 +2151,18 @@ def test_ediff1d_errors(self): to_end = dpnp.array([5], dtype="f4") assert_raises(TypeError, dpnp.ediff1d, a_dp, to_end=to_end) + # another `to_begin` sycl queue + to_begin = dpnp.array([-20, -15], sycl_queue=dpctl.SyclQueue()) + assert_raises( + ExecutionPlacementError, dpnp.ediff1d, a_dp, to_begin=to_begin + ) + + # another `to_end` sycl queue + to_end = dpnp.array([15, 20], sycl_queue=dpctl.SyclQueue()) + assert_raises( + ExecutionPlacementError, dpnp.ediff1d, a_dp, to_end=to_end + ) + @pytest.mark.usefixtures("allow_fall_back_on_numpy") class TestTrapz: diff --git a/tests/test_sycl_queue.py b/tests/test_sycl_queue.py index 902e5c9ba8d..9f40f6b4ec2 100644 --- a/tests/test_sycl_queue.py +++ b/tests/test_sycl_queue.py @@ -2407,12 +2407,7 @@ def test_nan_to_num(copy, device): @pytest.mark.parametrize( - "device_x", - valid_devices, - ids=[device.filter_string for device in valid_devices], -) -@pytest.mark.parametrize( - "device_args", + "device", valid_devices, ids=[device.filter_string for device in valid_devices], ) @@ -2424,15 +2419,15 @@ def test_nan_to_num(copy, device): (10, -10), ], ) -def test_ediff1d(device_x, device_args, to_end, to_begin): +def test_ediff1d(device, to_end, to_begin): data = [1, 3, 5, 7] - x = dpnp.array(data, device=device_x) + x = dpnp.array(data, device=device) if to_end: - to_end = dpnp.array(to_end, device=device_args) + to_end = dpnp.array(to_end, device=device) if to_begin: - to_begin = dpnp.array(to_begin, device=device_args) + to_begin = dpnp.array(to_begin, device=device) res = dpnp.ediff1d(x, to_end=to_end, to_begin=to_begin) diff --git a/tests/test_usm_type.py b/tests/test_usm_type.py index ca9dda5d198..dd4cacc253f 100644 --- a/tests/test_usm_type.py +++ b/tests/test_usm_type.py @@ -1431,4 +1431,4 @@ def test_ediff1d(usm_type_x, usm_type_args, to_end, to_begin): res = dp.ediff1d(x, to_end=to_end, to_begin=to_begin) - assert res.usm_type == x.usm_type + assert res.usm_type == du.get_coerced_usm_type([usm_type_x, usm_type_args])