Skip to content

Commit 86a8823

Browse files
authored
Leverage on dpctl.tensor implementation in dpnp.diff (#1963)
* Leverage on dpctl.tensor implementation of count_nonzero * Extend dpnp.get_result_array() to accept dpt.usm_ndarray * Leverage on dpctl.tensor implementation in dpnp.diff * Remove w/a since dpctl issue is resolved
1 parent cb56101 commit 86a8823

File tree

2 files changed

+13
-59
lines changed

2 files changed

+13
-59
lines changed

dpnp/dpnp_iface_mathematical.py

Lines changed: 8 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -130,31 +130,6 @@
130130
]
131131

132132

133-
def _append_to_diff_array(a, axis, combined, values):
134-
"""
135-
Append `values` to `combined` list based on data of array `a`.
136-
137-
Scalar value (including case with 0d array) is expanded to an array
138-
with length=1 in the direction of axis and the shape of the input array `a`
139-
along all other axes.
140-
Note, if `values` is a scalar, then it is converted to 0d array allocating
141-
on the same SYCL queue as the input array `a` and with the same USM type.
142-
143-
"""
144-
145-
dpnp.check_supported_arrays_type(values, scalar_type=True, all_scalars=True)
146-
if dpnp.isscalar(values):
147-
values = dpnp.asarray(
148-
values, sycl_queue=a.sycl_queue, usm_type=a.usm_type
149-
)
150-
151-
if values.ndim == 0:
152-
shape = list(a.shape)
153-
shape[axis] = 1
154-
values = dpnp.broadcast_to(values, tuple(shape))
155-
combined.append(values)
156-
157-
158133
def _get_reduction_res_dt(a, dtype, _out):
159134
"""Get a data type used by dpctl for result array in reduction function."""
160135

@@ -1206,39 +1181,14 @@ def diff(a, n=1, axis=-1, prepend=None, append=None):
12061181
12071182
"""
12081183

1209-
dpnp.check_supported_arrays_type(a)
1210-
if n == 0:
1211-
return a
1212-
if n < 0:
1213-
raise ValueError(f"order must be non-negative but got {n}")
1214-
1215-
nd = a.ndim
1216-
if nd == 0:
1217-
raise ValueError("diff requires input that is at least one dimensional")
1218-
axis = normalize_axis_index(axis, nd)
1219-
1220-
combined = []
1221-
if prepend is not None:
1222-
_append_to_diff_array(a, axis, combined, prepend)
1223-
1224-
combined.append(a)
1225-
if append is not None:
1226-
_append_to_diff_array(a, axis, combined, append)
1227-
1228-
if len(combined) > 1:
1229-
a = dpnp.concatenate(combined, axis=axis)
1230-
1231-
slice1 = [slice(None)] * nd
1232-
slice2 = [slice(None)] * nd
1233-
slice1[axis] = slice(1, None)
1234-
slice2[axis] = slice(None, -1)
1235-
slice1 = tuple(slice1)
1236-
slice2 = tuple(slice2)
1237-
1238-
op = dpnp.not_equal if a.dtype == numpy.bool_ else dpnp.subtract
1239-
for _ in range(n):
1240-
a = op(a[slice1], a[slice2])
1241-
return a
1184+
usm_a = dpnp.get_usm_ndarray(a)
1185+
usm_pre = (
1186+
None if prepend is None else dpnp.get_usm_ndarray_or_scalar(prepend)
1187+
)
1188+
usm_app = None if append is None else dpnp.get_usm_ndarray_or_scalar(append)
1189+
1190+
usm_res = dpt.diff(usm_a, axis=axis, n=n, prepend=usm_pre, append=usm_app)
1191+
return dpnp_array._create_from_usm_ndarray(usm_res)
12421192

12431193

12441194
_DIVIDE_DOCSTRING = """

tests/third_party/cupy/math_tests/test_sumprod.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -687,7 +687,11 @@ def test_diff_2dim_with_append(self, xp, dtype):
687687
@testing.numpy_cupy_allclose(type_check=has_support_aspect64())
688688
def test_diff_2dim_with_scalar_append(self, xp, dtype):
689689
a = testing.shaped_arange((4, 5), xp, dtype)
690-
return xp.diff(a, prepend=1, append=0)
690+
res = xp.diff(a, prepend=1, append=0)
691+
if xp is numpy and a.dtype != res.dtype:
692+
# res.dtype must follow a.dtype
693+
res = res.astype(a.dtype)
694+
return res
691695

692696
@testing.with_requires("numpy>=1.16")
693697
def test_diff_invalid_axis(self):

0 commit comments

Comments
 (0)