From 9ddecb09a2375fe6e60062edea9442e85d416d6b Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Mon, 29 Jul 2024 14:08:14 +0000
Subject: [PATCH 01/25] Implements `dpctl.tensor.diff`

---
 dpctl/tensor/__init__.py           |   3 +-
 dpctl/tensor/_utility_functions.py | 283 ++++++++++++++++++++++++++++-
 2 files changed, 284 insertions(+), 2 deletions(-)

diff --git a/dpctl/tensor/__init__.py b/dpctl/tensor/__init__.py
index 477a52e7be..28e3cdde23 100644
--- a/dpctl/tensor/__init__.py
+++ b/dpctl/tensor/__init__.py
@@ -94,7 +94,7 @@
 from dpctl.tensor._search_functions import where
 from dpctl.tensor._statistical_functions import mean, std, var
 from dpctl.tensor._usmarray import usm_ndarray
-from dpctl.tensor._utility_functions import all, any
+from dpctl.tensor._utility_functions import all, any, diff
 
 from ._accumulation import cumulative_logsumexp, cumulative_prod, cumulative_sum
 from ._array_api import __array_api_version__, __array_namespace_info__
@@ -373,4 +373,5 @@
     "cumulative_prod",
     "cumulative_sum",
     "nextafter",
+    "diff",
 ]
diff --git a/dpctl/tensor/_utility_functions.py b/dpctl/tensor/_utility_functions.py
index be11f965cc..2ee06ba93e 100644
--- a/dpctl/tensor/_utility_functions.py
+++ b/dpctl/tensor/_utility_functions.py
@@ -14,12 +14,24 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
+import operator
+
 import dpctl.tensor as dpt
 import dpctl.tensor._tensor_impl as ti
 import dpctl.tensor._tensor_reductions_impl as tri
 import dpctl.utils as du
+from dpctl.tensor._clip import (
+    _resolve_one_strong_one_weak_types,
+    _resolve_one_strong_two_weak_types,
+)
+from dpctl.tensor._elementwise_common import (
+    _get_dtype,
+    _get_queue_usm_type,
+    _get_shape,
+    _validate_dtype,
+)
 
-from ._numpy_helper import normalize_axis_tuple
+from ._numpy_helper import normalize_axis_index, normalize_axis_tuple
 
 
 def _boolean_reduction(x, axis, keepdims, func):
@@ -144,3 +156,272 @@ def any(x, /, *, axis=None, keepdims=False):
             containing the results of the logical OR reduction.
     """
     return _boolean_reduction(x, axis, keepdims, tri._any)
+
+
+def _validate_diff_shape(sh1, sh2, axis):
+    if not sh2:
+        # scalars will always be accepted
+        return True
+    else:
+        sh1_ndim = len(sh1)
+        if sh1_ndim == len(sh2) and all(
+            sh1[i] == sh2[i] for i in range(sh1_ndim) if i != axis
+        ):
+            return True
+        else:
+            return False
+
+
+def _concat_diff_input(arr, axis, prepend, append):
+    if prepend is not None and append is not None:
+        q1, x_usm_type = arr.sycl_queue, arr.usm_type
+        q2, prepend_usm_type = _get_queue_usm_type(prepend)
+        q3, append_usm_type = _get_queue_usm_type(append)
+        if q2 is None and q3 is None:
+            exec_q = q1
+            coerced_usm_type = x_usm_type
+        elif q3 is None:
+            exec_q = du.get_execution_queue((q1, q2))
+            if exec_q is None:
+                raise du.ExecutionPlacementError(
+                    "Execution placement can not be unambiguously inferred "
+                    "from input arguments."
+                )
+            coerced_usm_type = du.get_coerced_usm_type(
+                (
+                    x_usm_type,
+                    prepend_usm_type,
+                )
+            )
+        elif q2 is None:
+            exec_q = du.get_execution_queue((q1, q3))
+            if exec_q is None:
+                raise du.ExecutionPlacementError(
+                    "Execution placement can not be unambiguously inferred "
+                    "from input arguments."
+                )
+            coerced_usm_type = du.get_coerced_usm_type(
+                (
+                    x_usm_type,
+                    append_usm_type,
+                )
+            )
+        else:
+            exec_q = du.get_execution_queue((q1, q2, q3))
+            if exec_q is None:
+                raise du.ExecutionPlacementError(
+                    "Execution placement can not be unambiguously inferred "
+                    "from input arguments."
+                )
+            coerced_usm_type = du.get_coerced_usm_type(
+                (
+                    x_usm_type,
+                    prepend_usm_type,
+                    append_usm_type,
+                )
+            )
+        du.validate_usm_type(coerced_usm_type, allow_none=False)
+        arr_shape = arr.shape
+        prepend_shape = _get_shape(prepend)
+        append_shape = _get_shape(append)
+        if not all(
+            isinstance(s, (tuple, list))
+            for s in (
+                prepend_shape,
+                append_shape,
+            )
+        ):
+            raise TypeError(
+                "Shape of arguments can not be inferred. "
+                "Arguments are expected to be "
+                "lists, tuples, or both"
+            )
+        valid_prepend_shape = _validate_diff_shape(
+            arr_shape, prepend_shape, axis
+        )
+        if not valid_prepend_shape:
+            raise ValueError(
+                f"`diff` argument `prepend` with shape {prepend_shape} is "
+                f"invalid for first input with shape {arr_shape}"
+            )
+        valid_append_shape = _validate_diff_shape(arr_shape, append_shape, axis)
+        if not valid_append_shape:
+            raise ValueError(
+                f"`diff` argument `append` with shape {append_shape} is invalid"
+                f" for first input with shape {arr_shape}"
+            )
+        sycl_dev = exec_q.sycl_device
+        arr_dtype = arr.dtype
+        prepend_dtype = _get_dtype(prepend, sycl_dev)
+        append_dtype = _get_dtype(append, sycl_dev)
+        if not all(_validate_dtype(o) for o in (prepend_dtype, append_dtype)):
+            raise ValueError("Operands have unsupported data types")
+        prepend_dtype, append_dtype = _resolve_one_strong_two_weak_types(
+            arr_dtype, prepend_dtype, append_dtype, sycl_dev
+        )
+        if isinstance(prepend, dpt.usm_ndarray):
+            a_prepend = prepend
+        else:
+            a_prepend = dpt.asarray(
+                prepend,
+                dtype=prepend_dtype,
+                usm_type=coerced_usm_type,
+                sycl_queue=exec_q,
+            )
+        if isinstance(append, dpt.usm_ndarray):
+            a_append = append
+        else:
+            a_append = dpt.asarray(
+                prepend,
+                dtype=append_dtype,
+                usm_type=coerced_usm_type,
+                sycl_queue=exec_q,
+            )
+        if not prepend_shape:
+            prepend_shape = arr_shape[:axis] + (1,) + arr_shape[axis + 1 :]
+            a_prepend = dpt.broadcast_to(a_prepend, arr_shape)
+        if not append_shape:
+            append_shape = arr_shape[:axis] + (1,) + arr_shape[axis + 1 :]
+            a_append = dpt.broadcast_to(a_append, arr_shape)
+        return dpt.concat((a_prepend, arr, a_append), axis=axis)
+    elif prepend is not None:
+        q1, x_usm_type = arr.sycl_queue, arr.usm_type
+        q2, prepend_usm_type = _get_queue_usm_type(prepend)
+        if q2 is None:
+            exec_q = q1
+            coerced_usm_type = x_usm_type
+        else:
+            exec_q = du.get_execution_queue((q1, q2))
+            if exec_q is None:
+                raise du.ExecutionPlacementError(
+                    "Execution placement can not be unambiguously inferred "
+                    "from input arguments."
+                )
+            coerced_usm_type = du.get_coerced_usm_type(
+                (
+                    x_usm_type,
+                    prepend_usm_type,
+                )
+            )
+        du.validate_usm_type(coerced_usm_type, allow_none=False)
+        arr_shape = arr.shape
+        prepend_shape = _get_shape(prepend)
+        if not isinstance(prepend_shape, (tuple, list)):
+            raise TypeError(
+                "Shape of argument can not be inferred. "
+                "Argument is expected to be a "
+                "list or tuple"
+            )
+        valid_prepend_shape = _validate_diff_shape(
+            arr_shape, prepend_shape, axis
+        )
+        if not valid_prepend_shape:
+            raise ValueError(
+                f"`diff` argument `prepend` with shape {prepend_shape} is "
+                f"invalid for first input with shape {arr_shape}"
+            )
+        sycl_dev = exec_q.sycl_device
+        arr_dtype = arr.dtype
+        prepend_dtype = _get_dtype(prepend, sycl_dev)
+        if not _validate_dtype(prepend_dtype):
+            raise ValueError("Operand has unsupported data type")
+        prepend_dtype = _resolve_one_strong_one_weak_types(
+            arr_dtype, prepend_dtype, sycl_dev
+        )
+        if isinstance(prepend, dpt.usm_ndarray):
+            a_prepend = prepend
+        else:
+            a_prepend = dpt.asarray(
+                prepend,
+                dtype=prepend_dtype,
+                usm_type=coerced_usm_type,
+                sycl_queue=exec_q,
+            )
+        if not prepend_shape:
+            prepend_shape = arr_shape[:axis] + (1,) + arr_shape[axis + 1 :]
+            a_prepend = dpt.broadcast_to(a_prepend, arr_shape)
+        return dpt.concat((a_prepend, arr), axis=axis)
+    elif append is not None:
+        q1, x_usm_type = arr.sycl_queue, arr.usm_type
+        q2, append_usm_type = _get_queue_usm_type(append)
+        if q2 is None:
+            exec_q = q1
+            coerced_usm_type = x_usm_type
+        else:
+            exec_q = du.get_execution_queue((q1, q2))
+            if exec_q is None:
+                raise du.ExecutionPlacementError(
+                    "Execution placement can not be unambiguously inferred "
+                    "from input arguments."
+                )
+            coerced_usm_type = du.get_coerced_usm_type(
+                (
+                    x_usm_type,
+                    append_usm_type,
+                )
+            )
+        du.validate_usm_type(coerced_usm_type, allow_none=False)
+        arr_shape = arr.shape
+        append_shape = _get_shape(append)
+        if not isinstance(append_shape, (tuple, list)):
+            raise TypeError(
+                "Shape of argument can not be inferred. "
+                "Argument is expected to be a "
+                "list or tuple"
+            )
+        valid_append_shape = _validate_diff_shape(arr_shape, append_shape, axis)
+        if not valid_append_shape:
+            raise ValueError(
+                f"`diff` argument `append` with shape {append_shape} is invalid"
+                f" for first input with shape {arr_shape}"
+            )
+        sycl_dev = exec_q.sycl_device
+        arr_dtype = arr.dtype
+        append_dtype = _get_dtype(append, sycl_dev)
+        if not _validate_dtype(append_dtype):
+            raise ValueError("Operand has unsupported data type")
+        append_dtype = _resolve_one_strong_one_weak_types(
+            arr_dtype, append_dtype, sycl_dev
+        )
+        if isinstance(append, dpt.usm_ndarray):
+            a_append = append
+        else:
+            a_append = dpt.asarray(
+                append,
+                dtype=append_dtype,
+                usm_type=coerced_usm_type,
+                sycl_queue=exec_q,
+            )
+        if not append_shape:
+            append_shape = arr_shape[:axis] + (1,) + arr_shape[axis + 1 :]
+            a_append = dpt.broadcast_to(a_append, arr_shape)
+        return dpt.concat((arr, a_append), axis=axis)
+    else:
+        arr1 = arr
+    return arr1
+
+
+def diff(x, /, *, axis=-1, n=1, prepend=None, append=None):
+
+    if not isinstance(x, dpt.usm_ndarray):
+        raise TypeError(
+            "Expecting dpctl.tensor.usm_ndarray type, " f"got {type(x)}"
+        )
+    x_nd = x.ndim
+    axis = normalize_axis_index(operator.index(axis), x_nd)
+    n = operator.index(n)
+
+    arr = _concat_diff_input(x, axis, prepend, append)
+
+    # form slices and recurse
+    sl0 = tuple(
+        slice(None) if i != axis else slice(1, None) for i in range(x_nd)
+    )
+    sl1 = tuple(
+        slice(None) if i != axis else slice(None, -1) for i in range(x_nd)
+    )
+
+    for _ in range(n):
+        arr = arr[sl0] - arr[sl1]
+
+    return arr

From 442ba6f56dab4bb76d57591e10eceec0e2900d26 Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Mon, 29 Jul 2024 14:11:39 +0000
Subject: [PATCH 02/25] Move `_resolve_one_strong_one_weak_types` and
 `_resolve_one_strong_two_weak_types` to `_type_utils`

---
 dpctl/tensor/_clip.py              | 115 +----------------------------
 dpctl/tensor/_type_utils.py        | 108 +++++++++++++++++++++++++++
 dpctl/tensor/_utility_functions.py |  13 +++-
 3 files changed, 120 insertions(+), 116 deletions(-)

diff --git a/dpctl/tensor/_clip.py b/dpctl/tensor/_clip.py
index 594e713581..c8eea607f9 100644
--- a/dpctl/tensor/_clip.py
+++ b/dpctl/tensor/_clip.py
@@ -30,124 +30,15 @@
     _validate_dtype,
 )
 from dpctl.tensor._manipulation_functions import _broadcast_shape_impl
-from dpctl.tensor._type_utils import _can_cast, _to_device_supported_dtype
+from dpctl.tensor._type_utils import _can_cast
 from dpctl.utils import ExecutionPlacementError, SequentialOrderManager
 
 from ._type_utils import (
-    WeakComplexType,
-    WeakIntegralType,
-    _is_weak_dtype,
-    _strong_dtype_num_kind,
-    _weak_type_num_kind,
+    _resolve_one_strong_one_weak_types,
+    _resolve_one_strong_two_weak_types,
 )
 
 
-def _resolve_one_strong_two_weak_types(st_dtype, dtype1, dtype2, dev):
-    "Resolves weak data types per NEP-0050,"
-    "where the second and third arguments are"
-    "permitted to be weak types"
-    if _is_weak_dtype(st_dtype):
-        raise ValueError
-    if _is_weak_dtype(dtype1):
-        if _is_weak_dtype(dtype2):
-            kind_num1 = _weak_type_num_kind(dtype1)
-            kind_num2 = _weak_type_num_kind(dtype2)
-            st_kind_num = _strong_dtype_num_kind(st_dtype)
-
-            if kind_num1 > st_kind_num:
-                if isinstance(dtype1, WeakIntegralType):
-                    ret_dtype1 = dpt.dtype(ti.default_device_int_type(dev))
-                elif isinstance(dtype1, WeakComplexType):
-                    if st_dtype is dpt.float16 or st_dtype is dpt.float32:
-                        ret_dtype1 = dpt.complex64
-                    ret_dtype1 = _to_device_supported_dtype(dpt.complex128, dev)
-                else:
-                    ret_dtype1 = _to_device_supported_dtype(dpt.float64, dev)
-            else:
-                ret_dtype1 = st_dtype
-
-            if kind_num2 > st_kind_num:
-                if isinstance(dtype2, WeakIntegralType):
-                    ret_dtype2 = dpt.dtype(ti.default_device_int_type(dev))
-                elif isinstance(dtype2, WeakComplexType):
-                    if st_dtype is dpt.float16 or st_dtype is dpt.float32:
-                        ret_dtype2 = dpt.complex64
-                    ret_dtype2 = _to_device_supported_dtype(dpt.complex128, dev)
-                else:
-                    ret_dtype2 = _to_device_supported_dtype(dpt.float64, dev)
-            else:
-                ret_dtype2 = st_dtype
-
-            return ret_dtype1, ret_dtype2
-
-        max_dt_num_kind, max_dtype = max(
-            [
-                (_strong_dtype_num_kind(st_dtype), st_dtype),
-                (_strong_dtype_num_kind(dtype2), dtype2),
-            ]
-        )
-        dt1_kind_num = _weak_type_num_kind(dtype1)
-        if dt1_kind_num > max_dt_num_kind:
-            if isinstance(dtype1, WeakIntegralType):
-                return dpt.dtype(ti.default_device_int_type(dev)), dtype2
-            if isinstance(dtype1, WeakComplexType):
-                if max_dtype is dpt.float16 or max_dtype is dpt.float32:
-                    return dpt.complex64, dtype2
-                return (
-                    _to_device_supported_dtype(dpt.complex128, dev),
-                    dtype2,
-                )
-            return _to_device_supported_dtype(dpt.float64, dev), dtype2
-        else:
-            return max_dtype, dtype2
-    elif _is_weak_dtype(dtype2):
-        max_dt_num_kind, max_dtype = max(
-            [
-                (_strong_dtype_num_kind(st_dtype), st_dtype),
-                (_strong_dtype_num_kind(dtype1), dtype1),
-            ]
-        )
-        dt2_kind_num = _weak_type_num_kind(dtype2)
-        if dt2_kind_num > max_dt_num_kind:
-            if isinstance(dtype2, WeakIntegralType):
-                return dtype1, dpt.dtype(ti.default_device_int_type(dev))
-            if isinstance(dtype2, WeakComplexType):
-                if max_dtype is dpt.float16 or max_dtype is dpt.float32:
-                    return dtype1, dpt.complex64
-                return (
-                    dtype1,
-                    _to_device_supported_dtype(dpt.complex128, dev),
-                )
-            return dtype1, _to_device_supported_dtype(dpt.float64, dev)
-        else:
-            return dtype1, max_dtype
-    else:
-        # both are strong dtypes
-        # return unmodified
-        return dtype1, dtype2
-
-
-def _resolve_one_strong_one_weak_types(st_dtype, dtype, dev):
-    "Resolves one weak data type with one strong data type per NEP-0050"
-    if _is_weak_dtype(st_dtype):
-        raise ValueError
-    if _is_weak_dtype(dtype):
-        st_kind_num = _strong_dtype_num_kind(st_dtype)
-        kind_num = _weak_type_num_kind(dtype)
-        if kind_num > st_kind_num:
-            if isinstance(dtype, WeakIntegralType):
-                return dpt.dtype(ti.default_device_int_type(dev))
-            if isinstance(dtype, WeakComplexType):
-                if st_dtype is dpt.float16 or st_dtype is dpt.float32:
-                    return dpt.complex64
-                return _to_device_supported_dtype(dpt.complex128, dev)
-            return _to_device_supported_dtype(dpt.float64, dev)
-        else:
-            return st_dtype
-    else:
-        return dtype
-
-
 def _check_clip_dtypes(res_dtype, arg1_dtype, arg2_dtype, sycl_dev):
     "Checks if both types `arg1_dtype` and `arg2_dtype` can be"
     "cast to `res_dtype` according to the rule `safe`"
diff --git a/dpctl/tensor/_type_utils.py b/dpctl/tensor/_type_utils.py
index 364d2fc146..bebb1889f4 100644
--- a/dpctl/tensor/_type_utils.py
+++ b/dpctl/tensor/_type_utils.py
@@ -450,6 +450,112 @@ def _resolve_weak_types_all_py_ints(o1_dtype, o2_dtype, dev):
         return o1_dtype, o2_dtype
 
 
+def _resolve_one_strong_two_weak_types(st_dtype, dtype1, dtype2, dev):
+    "Resolves weak data types per NEP-0050,"
+    "where the second and third arguments are"
+    "permitted to be weak types"
+    if _is_weak_dtype(st_dtype):
+        raise ValueError
+    if _is_weak_dtype(dtype1):
+        if _is_weak_dtype(dtype2):
+            kind_num1 = _weak_type_num_kind(dtype1)
+            kind_num2 = _weak_type_num_kind(dtype2)
+            st_kind_num = _strong_dtype_num_kind(st_dtype)
+
+            if kind_num1 > st_kind_num:
+                if isinstance(dtype1, WeakIntegralType):
+                    ret_dtype1 = dpt.dtype(ti.default_device_int_type(dev))
+                elif isinstance(dtype1, WeakComplexType):
+                    if st_dtype is dpt.float16 or st_dtype is dpt.float32:
+                        ret_dtype1 = dpt.complex64
+                    ret_dtype1 = _to_device_supported_dtype(dpt.complex128, dev)
+                else:
+                    ret_dtype1 = _to_device_supported_dtype(dpt.float64, dev)
+            else:
+                ret_dtype1 = st_dtype
+
+            if kind_num2 > st_kind_num:
+                if isinstance(dtype2, WeakIntegralType):
+                    ret_dtype2 = dpt.dtype(ti.default_device_int_type(dev))
+                elif isinstance(dtype2, WeakComplexType):
+                    if st_dtype is dpt.float16 or st_dtype is dpt.float32:
+                        ret_dtype2 = dpt.complex64
+                    ret_dtype2 = _to_device_supported_dtype(dpt.complex128, dev)
+                else:
+                    ret_dtype2 = _to_device_supported_dtype(dpt.float64, dev)
+            else:
+                ret_dtype2 = st_dtype
+
+            return ret_dtype1, ret_dtype2
+
+        max_dt_num_kind, max_dtype = max(
+            [
+                (_strong_dtype_num_kind(st_dtype), st_dtype),
+                (_strong_dtype_num_kind(dtype2), dtype2),
+            ]
+        )
+        dt1_kind_num = _weak_type_num_kind(dtype1)
+        if dt1_kind_num > max_dt_num_kind:
+            if isinstance(dtype1, WeakIntegralType):
+                return dpt.dtype(ti.default_device_int_type(dev)), dtype2
+            if isinstance(dtype1, WeakComplexType):
+                if max_dtype is dpt.float16 or max_dtype is dpt.float32:
+                    return dpt.complex64, dtype2
+                return (
+                    _to_device_supported_dtype(dpt.complex128, dev),
+                    dtype2,
+                )
+            return _to_device_supported_dtype(dpt.float64, dev), dtype2
+        else:
+            return max_dtype, dtype2
+    elif _is_weak_dtype(dtype2):
+        max_dt_num_kind, max_dtype = max(
+            [
+                (_strong_dtype_num_kind(st_dtype), st_dtype),
+                (_strong_dtype_num_kind(dtype1), dtype1),
+            ]
+        )
+        dt2_kind_num = _weak_type_num_kind(dtype2)
+        if dt2_kind_num > max_dt_num_kind:
+            if isinstance(dtype2, WeakIntegralType):
+                return dtype1, dpt.dtype(ti.default_device_int_type(dev))
+            if isinstance(dtype2, WeakComplexType):
+                if max_dtype is dpt.float16 or max_dtype is dpt.float32:
+                    return dtype1, dpt.complex64
+                return (
+                    dtype1,
+                    _to_device_supported_dtype(dpt.complex128, dev),
+                )
+            return dtype1, _to_device_supported_dtype(dpt.float64, dev)
+        else:
+            return dtype1, max_dtype
+    else:
+        # both are strong dtypes
+        # return unmodified
+        return dtype1, dtype2
+
+
+def _resolve_one_strong_one_weak_types(st_dtype, dtype, dev):
+    "Resolves one weak data type with one strong data type per NEP-0050"
+    if _is_weak_dtype(st_dtype):
+        raise ValueError
+    if _is_weak_dtype(dtype):
+        st_kind_num = _strong_dtype_num_kind(st_dtype)
+        kind_num = _weak_type_num_kind(dtype)
+        if kind_num > st_kind_num:
+            if isinstance(dtype, WeakIntegralType):
+                return dpt.dtype(ti.default_device_int_type(dev))
+            if isinstance(dtype, WeakComplexType):
+                if st_dtype is dpt.float16 or st_dtype is dpt.float32:
+                    return dpt.complex64
+                return _to_device_supported_dtype(dpt.complex128, dev)
+            return _to_device_supported_dtype(dpt.float64, dev)
+        else:
+            return st_dtype
+    else:
+        return dtype
+
+
 class finfo_object:
     """
     `numpy.finfo` subclass which returns Python floating-point scalars for
@@ -838,6 +944,8 @@ def _default_accumulation_dtype_fp_types(inp_dt, q):
     "_acceptance_fn_divide",
     "_acceptance_fn_negative",
     "_acceptance_fn_subtract",
+    "_resolve_one_strong_one_weak_types",
+    "_resolve_one_strong_two_weak_types",
     "_resolve_weak_types",
     "_resolve_weak_types_all_py_ints",
     "_weak_type_num_kind",
diff --git a/dpctl/tensor/_utility_functions.py b/dpctl/tensor/_utility_functions.py
index 2ee06ba93e..9cf250cb1e 100644
--- a/dpctl/tensor/_utility_functions.py
+++ b/dpctl/tensor/_utility_functions.py
@@ -20,10 +20,6 @@
 import dpctl.tensor._tensor_impl as ti
 import dpctl.tensor._tensor_reductions_impl as tri
 import dpctl.utils as du
-from dpctl.tensor._clip import (
-    _resolve_one_strong_one_weak_types,
-    _resolve_one_strong_two_weak_types,
-)
 from dpctl.tensor._elementwise_common import (
     _get_dtype,
     _get_queue_usm_type,
@@ -32,6 +28,10 @@
 )
 
 from ._numpy_helper import normalize_axis_index, normalize_axis_tuple
+from ._type_utils import (
+    _resolve_one_strong_one_weak_types,
+    _resolve_one_strong_two_weak_types,
+)
 
 
 def _boolean_reduction(x, axis, keepdims, func):
@@ -159,6 +159,8 @@ def any(x, /, *, axis=None, keepdims=False):
 
 
 def _validate_diff_shape(sh1, sh2, axis):
+    """Utility for validating that two shapes `sh1` and `sh2`
+    are possible to concatenate along `axis`."""
     if not sh2:
         # scalars will always be accepted
         return True
@@ -173,6 +175,9 @@ def _validate_diff_shape(sh1, sh2, axis):
 
 
 def _concat_diff_input(arr, axis, prepend, append):
+    """Concatenates `arr`, `prepend` and, `append` along `axis`,
+    where `arr` is an array and `prepend` and `append` are
+    any mixture of arrays and scalars."""
     if prepend is not None and append is not None:
         q1, x_usm_type = arr.sycl_queue, arr.usm_type
         q2, prepend_usm_type = _get_queue_usm_type(prepend)

From e3432d7d93f76328ce113473e849afd945351f2d Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Mon, 29 Jul 2024 14:15:05 +0000
Subject: [PATCH 03/25] Implements `dpctl.tensor.count_nonzero`

---
 dpctl/tensor/__init__.py   |  2 ++
 dpctl/tensor/_reduction.py | 12 ++++++++++++
 2 files changed, 14 insertions(+)

diff --git a/dpctl/tensor/__init__.py b/dpctl/tensor/__init__.py
index 28e3cdde23..dff75b9c2c 100644
--- a/dpctl/tensor/__init__.py
+++ b/dpctl/tensor/__init__.py
@@ -176,6 +176,7 @@
 from ._reduction import (
     argmax,
     argmin,
+    count_nonzero,
     logsumexp,
     max,
     min,
@@ -374,4 +375,5 @@
     "cumulative_sum",
     "nextafter",
     "diff",
+    "count_nonzero",
 ]
diff --git a/dpctl/tensor/_reduction.py b/dpctl/tensor/_reduction.py
index cdc1cf0f26..cadab1dd89 100644
--- a/dpctl/tensor/_reduction.py
+++ b/dpctl/tensor/_reduction.py
@@ -773,3 +773,15 @@ def argmin(x, /, *, axis=None, keepdims=False, out=None):
             default array index data type for the device of ``x``.
     """
     return _search_over_axis(x, axis, keepdims, out, tri._argmin_over_axis)
+
+
+def count_nonzero(x, /, *, axis=None, keepdims=False, out=None):
+    if x.dtype != dpt.bool:
+        x = dpt.astype(x, dpt.bool, copy=False)
+    return sum(
+        x,
+        axis=axis,
+        dtype=ti.default_device_index_type(x.sycl_device),
+        keepdims=keepdims,
+        out=None,
+    )

From eaaac83a006c627ce956e671eda9564cb9e1fd48 Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Mon, 29 Jul 2024 14:16:26 +0000
Subject: [PATCH 04/25] Adds docstrings for `diff` and `count_nonzero`

---
 dpctl/tensor/_reduction.py         | 31 +++++++++++++++++++++++
 dpctl/tensor/_utility_functions.py | 40 +++++++++++++++++++++++++++++-
 2 files changed, 70 insertions(+), 1 deletion(-)

diff --git a/dpctl/tensor/_reduction.py b/dpctl/tensor/_reduction.py
index cadab1dd89..4089f13374 100644
--- a/dpctl/tensor/_reduction.py
+++ b/dpctl/tensor/_reduction.py
@@ -776,6 +776,37 @@ def argmin(x, /, *, axis=None, keepdims=False, out=None):
 
 
 def count_nonzero(x, /, *, axis=None, keepdims=False, out=None):
+    """
+    Counts the number of elements in the input array ``x`` which are non-zero.
+
+    Args:
+        x (usm_ndarray):
+            input array.
+        axis (Optional[int, Tuple[int, ...]]):
+            axis or axes along which to count. If a tuple of unique integers,
+            the number of non-zero values are computed over multiple axes.
+            If ``None``, the number of non-zero values is computed over the
+            entire array.
+            Default: ``None``.
+        keepdims (Optional[bool]):
+            if ``True``, the reduced axes (dimensions) are included in the
+            result as singleton dimensions, so that the returned array remains
+            compatible with the input arrays according to Array Broadcasting
+            rules. Otherwise, if ``False``, the reduced axes are not included
+            in the returned array. Default: ``False``.
+        out (Optional[usm_ndarray]):
+            the array into which the result is written.
+            The data type of ``out`` must match the expected shape and data
+            type.
+            If ``None`` then a new array is returned. Default: ``None``.
+
+    Returns:
+        usm_ndarray:
+            an array containing the count of non-zero values. If the sum was
+            computed over the entire array, a zero-dimensional array is
+            returned. The returned array will have the default array index data
+            type.
+    """
     if x.dtype != dpt.bool:
         x = dpt.astype(x, dpt.bool, copy=False)
     return sum(
diff --git a/dpctl/tensor/_utility_functions.py b/dpctl/tensor/_utility_functions.py
index 9cf250cb1e..cde91bc6f6 100644
--- a/dpctl/tensor/_utility_functions.py
+++ b/dpctl/tensor/_utility_functions.py
@@ -407,6 +407,43 @@ def _concat_diff_input(arr, axis, prepend, append):
 
 
 def diff(x, /, *, axis=-1, n=1, prepend=None, append=None):
+    """
+    Calculates the `n`-th discrete forward difference of `x` along `axis`.
+
+    Args:
+        x (usm_ndarray):
+            input array.
+        axis (int):
+            axis along which to compute the difference. A valid axis must be on
+            the interval `[-N, N)`, where `N` is the rank (number of
+            dimensions) of `x`.
+            Default: `-1`
+        n (int):
+            number of times to recursively compute the difference.
+            Default: `1`.
+        prepend (Union[usm_ndarray, bool, int, float, complex]):
+            value or values to prepend to the specified axis before taking the
+            difference.
+            Must have the same shape as `x` except along `axis`, which can have
+            any shape.
+            Default: `None`.
+        append (Union[usm_ndarray, bool, int, float, complex]):
+            value or values to append to the specified axis before taking the
+            difference.
+            Must have the same shape as `x` except along `axis`, which can have
+            any shape.
+            Default: `None`.
+
+    Returns:
+        usm_ndarray:
+            an array containing the `n`-th differences. The array will have the
+            same shape as `x`, except along `axis`, which will have shape
+
+            - prepend.shape[axis] + x.shape[axis] + append.shape[axis] - n
+
+            The data type of the returned array is determined by the Type
+            Promotion Rules.
+    """
 
     if not isinstance(x, dpt.usm_ndarray):
         raise TypeError(
@@ -417,7 +454,8 @@ def diff(x, /, *, axis=-1, n=1, prepend=None, append=None):
     n = operator.index(n)
 
     arr = _concat_diff_input(x, axis, prepend, append)
-
+    if n == 0:
+        return arr
     # form slices and recurse
     sl0 = tuple(
         slice(None) if i != axis else slice(1, None) for i in range(x_nd)

From 6c8a61f3034c9637d3f257b1e3753c9a564c512c Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Mon, 29 Jul 2024 22:21:43 +0000
Subject: [PATCH 05/25] Fix typo in `count_nonzero` caught in review

---
 dpctl/tensor/_reduction.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/dpctl/tensor/_reduction.py b/dpctl/tensor/_reduction.py
index 4089f13374..4b8403309f 100644
--- a/dpctl/tensor/_reduction.py
+++ b/dpctl/tensor/_reduction.py
@@ -814,5 +814,5 @@ def count_nonzero(x, /, *, axis=None, keepdims=False, out=None):
         axis=axis,
         dtype=ti.default_device_index_type(x.sycl_device),
         keepdims=keepdims,
-        out=None,
+        out=out,
     )

From b280036a82cbdb9f8f1e2f41bd31fc2565fa1cbc Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Mon, 29 Jul 2024 22:27:54 +0000
Subject: [PATCH 06/25] `diff` now uses not_equal when the input is boolean

---
 dpctl/tensor/_utility_functions.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/dpctl/tensor/_utility_functions.py b/dpctl/tensor/_utility_functions.py
index cde91bc6f6..9e4b431420 100644
--- a/dpctl/tensor/_utility_functions.py
+++ b/dpctl/tensor/_utility_functions.py
@@ -464,7 +464,8 @@ def diff(x, /, *, axis=-1, n=1, prepend=None, append=None):
         slice(None) if i != axis else slice(None, -1) for i in range(x_nd)
     )
 
+    diff_op = dpt.not_equal if x.dtype == dpt.bool else dpt.subtract
     for _ in range(n):
-        arr = arr[sl0] - arr[sl1]
+        arr = diff_op(arr[sl0], arr[sl1])
 
     return arr

From 9fb94d570ce3c796f4c15388fa9a94f2955f9622 Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Tue, 30 Jul 2024 16:33:54 +0000
Subject: [PATCH 07/25] Rewrite `diff` to juggle between two larger temporary
 allocations when n > 2

Saves some temporary memory allocations
---
 dpctl/tensor/_utility_functions.py | 17 +++++++++++++++--
 1 file changed, 15 insertions(+), 2 deletions(-)

diff --git a/dpctl/tensor/_utility_functions.py b/dpctl/tensor/_utility_functions.py
index 9e4b431420..d8bf47542a 100644
--- a/dpctl/tensor/_utility_functions.py
+++ b/dpctl/tensor/_utility_functions.py
@@ -465,7 +465,20 @@ def diff(x, /, *, axis=-1, n=1, prepend=None, append=None):
     )
 
     diff_op = dpt.not_equal if x.dtype == dpt.bool else dpt.subtract
-    for _ in range(n):
+    if n > 1:
+        arr_tmp0 = diff_op(arr[sl0], arr[sl1])
+        arr_tmp1 = diff_op(arr_tmp0[sl0], arr_tmp0[sl1])
+        n = n - 2
+        if n > 0:
+            sl3 = tuple(
+                slice(None) if i != axis else slice(None, -2)
+                for i in range(x_nd)
+            )
+            for _ in range(n):
+                arr_tmp0_sliced = arr_tmp0[sl3]
+                diff_op(arr_tmp1[sl0], arr_tmp1[sl1], out=arr_tmp0_sliced)
+                arr_tmp0, arr_tmp1 = arr_tmp1, arr_tmp0_sliced
+        arr = arr_tmp1
+    else:
         arr = diff_op(arr[sl0], arr[sl1])
-
     return arr

From 234d81e9b992ac3d14c12cdf9911f38587b1bac3 Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Tue, 30 Jul 2024 17:12:54 +0000
Subject: [PATCH 08/25] Adds a basic test for `count_nonzero`

---
 dpctl/tests/test_usm_ndarray_reductions.py | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/dpctl/tests/test_usm_ndarray_reductions.py b/dpctl/tests/test_usm_ndarray_reductions.py
index 4e66ae29a7..f7b08ba369 100644
--- a/dpctl/tests/test_usm_ndarray_reductions.py
+++ b/dpctl/tests/test_usm_ndarray_reductions.py
@@ -21,6 +21,7 @@
 from numpy.testing import assert_allclose
 
 import dpctl.tensor as dpt
+from dpctl.tensor._tensor_impl import default_device_index_type
 from dpctl.tests.helper import get_queue_or_skip, skip_if_dtype_not_supported
 from dpctl.utils import ExecutionPlacementError
 
@@ -669,3 +670,21 @@ def test_reduction_out_kwarg_arg_validation():
             keepdims=True,
             out=dpt.empty_like(out_wrong_keepdims, dtype=ind_dt),
         )
+
+
+@pytest.mark.parametrize("dt", _all_dtypes)
+def test_count_nonzero(dt):
+    q = get_queue_or_skip()
+    skip_if_dtype_not_supported(dt, q)
+
+    expected_dt = default_device_index_type(q.sycl_device)
+
+    x = dpt.ones(10, dtype=dt, sycl_queue=q)
+    res = dpt.count_nonzero(x)
+    assert res == 10
+    assert x.dtype == expected_dt
+
+    x[3:6] = 0
+    res = dpt.count_nonzero(x)
+    assert res == 7
+    assert x.dtype == expected_dt

From 6664b18771e8b3c827be1fcaa15efab2ec5c754a Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Tue, 30 Jul 2024 17:28:44 +0000
Subject: [PATCH 09/25] Use `builtins.all` in `diff` utilities to avoid
 conflicts with `all` defined in the same file

---
 dpctl/tensor/_utility_functions.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/dpctl/tensor/_utility_functions.py b/dpctl/tensor/_utility_functions.py
index d8bf47542a..8e9cd574a6 100644
--- a/dpctl/tensor/_utility_functions.py
+++ b/dpctl/tensor/_utility_functions.py
@@ -14,6 +14,7 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
+import builtins
 import operator
 
 import dpctl.tensor as dpt
@@ -166,7 +167,7 @@ def _validate_diff_shape(sh1, sh2, axis):
         return True
     else:
         sh1_ndim = len(sh1)
-        if sh1_ndim == len(sh2) and all(
+        if sh1_ndim == len(sh2) and builtins.all(
             sh1[i] == sh2[i] for i in range(sh1_ndim) if i != axis
         ):
             return True
@@ -229,7 +230,7 @@ def _concat_diff_input(arr, axis, prepend, append):
         arr_shape = arr.shape
         prepend_shape = _get_shape(prepend)
         append_shape = _get_shape(append)
-        if not all(
+        if not builtins.all(
             isinstance(s, (tuple, list))
             for s in (
                 prepend_shape,

From d4b575f4181895b7e52568ba31859a58215ad05d Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Tue, 30 Jul 2024 18:03:02 +0000
Subject: [PATCH 10/25] Adds `test_tensor_diff` with tests for basic `diff` and
 axis keyword

---
 dpctl/tests/test_tensor_diff.py | 81 +++++++++++++++++++++++++++++++++
 1 file changed, 81 insertions(+)
 create mode 100644 dpctl/tests/test_tensor_diff.py

diff --git a/dpctl/tests/test_tensor_diff.py b/dpctl/tests/test_tensor_diff.py
new file mode 100644
index 0000000000..48848d4c68
--- /dev/null
+++ b/dpctl/tests/test_tensor_diff.py
@@ -0,0 +1,81 @@
+#                       Data Parallel Control (dpctl)
+#
+#  Copyright 2020-2024 Intel Corporation
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+import pytest
+
+import dpctl.tensor as dpt
+from dpctl.tests.helper import get_queue_or_skip, skip_if_dtype_not_supported
+
+_all_dtypes = [
+    "?",
+    "i1",
+    "u1",
+    "i2",
+    "u2",
+    "i4",
+    "u4",
+    "i8",
+    "u8",
+    "f2",
+    "f4",
+    "f8",
+    "c8",
+    "c16",
+]
+
+
+@pytest.mark.parametrize("dt", _all_dtypes)
+def test_diff_basic(dt):
+    q = get_queue_or_skip()
+    skip_if_dtype_not_supported(dt, q)
+
+    x = dpt.asarray([9, 12, 7, 17, 10, 18, 15, 9, 8, 8], dtype=dt, sycl_queue=q)
+    res = dpt.diff(x)
+    op = dpt.not_equal if x.dtype is dpt.bool else dpt.subtract
+    expected_res = op(x[1:], x[:-1])
+    if dpt.dtype(dt).kind in "fc":
+        assert dpt.allclose(res, expected_res)
+    else:
+        assert dpt.all(res == expected_res)
+
+    res = dpt.diff(x, n=5)
+    expected_res = x
+    for _ in range(5):
+        expected_res = op(expected_res[1:], expected_res[:-1])
+    if dpt.dtype(dt).kind in "fc":
+        assert dpt.allclose(res, expected_res)
+    else:
+        assert dpt.all(res == expected_res)
+
+
+def test_diff_axis():
+    get_queue_or_skip()
+
+    x = dpt.tile(
+        dpt.asarray([9, 12, 7, 17, 10, 18, 15, 9, 8, 8], dtype="i4"), (3, 4, 1)
+    )
+    x[:, ::2, :] = 0
+    res = dpt.diff(x, n=1, axis=1)
+    expected_res = dpt.subtract(x[:, 1:, :], x[:, :-1, :])
+    assert dpt.all(res == expected_res)
+
+    res = dpt.diff(x, n=3, axis=1)
+    expected_res = x
+    for _ in range(3):
+        expected_res = dpt.subtract(
+            expected_res[:, 1:, :], expected_res[:, :-1, :]
+        )
+    assert dpt.all(res == expected_res)

From e5b144c4a463b8b80141736bd9869ef01f9e2790 Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Tue, 30 Jul 2024 19:35:16 +0000
Subject: [PATCH 11/25] Fixes typo in `test_count_nonzero`

---
 dpctl/tests/test_usm_ndarray_reductions.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/dpctl/tests/test_usm_ndarray_reductions.py b/dpctl/tests/test_usm_ndarray_reductions.py
index f7b08ba369..b74ef372f4 100644
--- a/dpctl/tests/test_usm_ndarray_reductions.py
+++ b/dpctl/tests/test_usm_ndarray_reductions.py
@@ -682,9 +682,9 @@ def test_count_nonzero(dt):
     x = dpt.ones(10, dtype=dt, sycl_queue=q)
     res = dpt.count_nonzero(x)
     assert res == 10
-    assert x.dtype == expected_dt
+    assert res.dtype == expected_dt
 
     x[3:6] = 0
     res = dpt.count_nonzero(x)
     assert res == 7
-    assert x.dtype == expected_dt
+    assert res.dtype == expected_dt

From 2a6b1cf118de9ee82de5fbffb65d161039cbb710 Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Tue, 30 Jul 2024 19:42:02 +0000
Subject: [PATCH 12/25] Fixes another case of `all` name conflict

---
 dpctl/tensor/_utility_functions.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/dpctl/tensor/_utility_functions.py b/dpctl/tensor/_utility_functions.py
index 8e9cd574a6..c6e5c12a0f 100644
--- a/dpctl/tensor/_utility_functions.py
+++ b/dpctl/tensor/_utility_functions.py
@@ -260,7 +260,9 @@ def _concat_diff_input(arr, axis, prepend, append):
         arr_dtype = arr.dtype
         prepend_dtype = _get_dtype(prepend, sycl_dev)
         append_dtype = _get_dtype(append, sycl_dev)
-        if not all(_validate_dtype(o) for o in (prepend_dtype, append_dtype)):
+        if not builtins.all(
+            _validate_dtype(o) for o in (prepend_dtype, append_dtype)
+        ):
             raise ValueError("Operands have unsupported data types")
         prepend_dtype, append_dtype = _resolve_one_strong_two_weak_types(
             arr_dtype, prepend_dtype, append_dtype, sycl_dev

From 6ae2ac20e2bd2018ab86cda4c7862c419273a198 Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Tue, 30 Jul 2024 20:17:47 +0000
Subject: [PATCH 13/25] More tests for `diff`

---
 dpctl/tests/test_tensor_diff.py | 43 +++++++++++++++++++++++++++++++++
 1 file changed, 43 insertions(+)

diff --git a/dpctl/tests/test_tensor_diff.py b/dpctl/tests/test_tensor_diff.py
index 48848d4c68..4e3ff93dbb 100644
--- a/dpctl/tests/test_tensor_diff.py
+++ b/dpctl/tests/test_tensor_diff.py
@@ -17,6 +17,7 @@
 import pytest
 
 import dpctl.tensor as dpt
+from dpctl.tensor._type_utils import _to_device_supported_dtype
 from dpctl.tests.helper import get_queue_or_skip, skip_if_dtype_not_supported
 
 _all_dtypes = [
@@ -79,3 +80,45 @@ def test_diff_axis():
             expected_res[:, 1:, :], expected_res[:, :-1, :]
         )
     assert dpt.all(res == expected_res)
+
+
+def test_diff_prepend_append_type_promotion():
+    get_queue_or_skip()
+
+    dts = [
+        ("i1", "u1", "i8"),
+        ("i1", "u8", "u1"),
+        ("u4", "i4", "f4"),
+        ("i8", "c8", "u8"),
+    ]
+
+    for _dts in dts:
+        x = dpt.ones(10, dtype=_dts[1])
+        prepend = dpt.full(1, 2, dtype=_dts[0])
+        append = dpt.full(1, 3, dtype=_dts[2])
+
+        res = dpt.diff(x, prepend=prepend, append=append)
+        assert res.dtype == _to_device_supported_dtype(
+            dpt.result_type(prepend, x, append),
+            x.sycl_queue.sycl_device,
+        )
+
+        res = dpt.diff(x, prepend=prepend)
+        assert res.dtype == _to_device_supported_dtype(
+            dpt.result_type(prepend, x),
+            x.sycl_queue.sycl_device,
+        )
+
+        res = dpt.diff(x, append=append)
+        assert res.dtype == _to_device_supported_dtype(
+            dpt.result_type(x, append),
+            x.sycl_queue.sycl_device,
+        )
+
+
+def test_diff_0d():
+    get_queue_or_skip()
+
+    x = dpt.ones(())
+    with pytest.raises(ValueError):
+        dpt.diff(x)

From 316254bb442db909a51b7f3a209d1fa985ae4d75 Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Tue, 30 Jul 2024 20:27:26 +0000
Subject: [PATCH 14/25] Test `diff` with an empty array and axis

---
 dpctl/tests/test_tensor_diff.py | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/dpctl/tests/test_tensor_diff.py b/dpctl/tests/test_tensor_diff.py
index 4e3ff93dbb..d096cbe651 100644
--- a/dpctl/tests/test_tensor_diff.py
+++ b/dpctl/tests/test_tensor_diff.py
@@ -122,3 +122,22 @@ def test_diff_0d():
     x = dpt.ones(())
     with pytest.raises(ValueError):
         dpt.diff(x)
+
+
+def test_diff_empty_array():
+    get_queue_or_skip()
+
+    x = dpt.ones((3, 0, 5))
+    res = dpt.diff(x, axis=1)
+    assert res.shape == x.shape
+
+    res = dpt.diff(x, axis=0)
+    assert res.shape == (2, 0, 5)
+
+    append = dpt.ones((3, 2, 5))
+    res = dpt.diff(x, axis=1, append=append)
+    assert res.shape == (3, 1, 5)
+
+    prepend = dpt.ones((3, 2, 5))
+    res = dpt.diff(x, axis=1, prepend=prepend)
+    assert res.shape == (3, 1, 5)

From fdf9adae6f9a2bde63592ed70c34f936e00ac12a Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Tue, 30 Jul 2024 22:20:19 +0000
Subject: [PATCH 15/25] Fixes `diff` for Python scalar `append` or `prepend`

Typos caused this to result in incorrect shaped outputs
---
 dpctl/tensor/_utility_functions.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/dpctl/tensor/_utility_functions.py b/dpctl/tensor/_utility_functions.py
index c6e5c12a0f..0bd84694fb 100644
--- a/dpctl/tensor/_utility_functions.py
+++ b/dpctl/tensor/_utility_functions.py
@@ -287,10 +287,10 @@ def _concat_diff_input(arr, axis, prepend, append):
             )
         if not prepend_shape:
             prepend_shape = arr_shape[:axis] + (1,) + arr_shape[axis + 1 :]
-            a_prepend = dpt.broadcast_to(a_prepend, arr_shape)
+            a_prepend = dpt.broadcast_to(a_prepend, prepend_shape)
         if not append_shape:
             append_shape = arr_shape[:axis] + (1,) + arr_shape[axis + 1 :]
-            a_append = dpt.broadcast_to(a_append, arr_shape)
+            a_append = dpt.broadcast_to(a_append, append_shape)
         return dpt.concat((a_prepend, arr, a_append), axis=axis)
     elif prepend is not None:
         q1, x_usm_type = arr.sycl_queue, arr.usm_type
@@ -347,7 +347,7 @@ def _concat_diff_input(arr, axis, prepend, append):
             )
         if not prepend_shape:
             prepend_shape = arr_shape[:axis] + (1,) + arr_shape[axis + 1 :]
-            a_prepend = dpt.broadcast_to(a_prepend, arr_shape)
+            a_prepend = dpt.broadcast_to(a_prepend, prepend_shape)
         return dpt.concat((a_prepend, arr), axis=axis)
     elif append is not None:
         q1, x_usm_type = arr.sycl_queue, arr.usm_type
@@ -402,7 +402,7 @@ def _concat_diff_input(arr, axis, prepend, append):
             )
         if not append_shape:
             append_shape = arr_shape[:axis] + (1,) + arr_shape[axis + 1 :]
-            a_append = dpt.broadcast_to(a_append, arr_shape)
+            a_append = dpt.broadcast_to(a_append, append_shape)
         return dpt.concat((arr, a_append), axis=axis)
     else:
         arr1 = arr

From e38234aa33115cae7a860b93dc7c0d097584a7e8 Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Tue, 30 Jul 2024 22:40:44 +0000
Subject: [PATCH 16/25] Corrects typo in `diff`

---
 dpctl/tensor/_utility_functions.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/dpctl/tensor/_utility_functions.py b/dpctl/tensor/_utility_functions.py
index 0bd84694fb..d0e6b8f9bc 100644
--- a/dpctl/tensor/_utility_functions.py
+++ b/dpctl/tensor/_utility_functions.py
@@ -280,7 +280,7 @@ def _concat_diff_input(arr, axis, prepend, append):
             a_append = append
         else:
             a_append = dpt.asarray(
-                prepend,
+                append,
                 dtype=append_dtype,
                 usm_type=coerced_usm_type,
                 sycl_queue=exec_q,

From c3bdf9d67c0a3d3ffb51089a1b4937a963e2a86d Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Tue, 30 Jul 2024 22:57:37 +0000
Subject: [PATCH 17/25] Test `diff` with Python scalars for `prepend` and
 `append` and no-op case

---
 dpctl/tests/test_tensor_diff.py | 53 +++++++++++++++++++++++++++++++++
 1 file changed, 53 insertions(+)

diff --git a/dpctl/tests/test_tensor_diff.py b/dpctl/tests/test_tensor_diff.py
index d096cbe651..ecde3ff402 100644
--- a/dpctl/tests/test_tensor_diff.py
+++ b/dpctl/tests/test_tensor_diff.py
@@ -141,3 +141,56 @@ def test_diff_empty_array():
     prepend = dpt.ones((3, 2, 5))
     res = dpt.diff(x, axis=1, prepend=prepend)
     assert res.shape == (3, 1, 5)
+
+
+def test_diff_no_op():
+    get_queue_or_skip()
+
+    x = dpt.ones(10, dtype="i4")
+    res = dpt.diff(x, n=0)
+    assert dpt.all(x == res)
+
+    res = dpt.diff(dpt.reshape(x, (2, 5)), n=0, axis=0)
+    assert dpt.all(x == res)
+
+
+@pytest.mark.parametrize("sh,axis", [((1,), 0), ((3, 4, 5), 1)])
+def test_diff_prepend_append_py_scalars(sh, axis):
+    get_queue_or_skip()
+
+    arrs = [
+        dpt.ones(sh, dtype="?"),
+        dpt.ones(sh, dtype="i4"),
+        dpt.ones(sh, dtype="f4"),
+        dpt.ones(sh, dtype="c8"),
+    ]
+
+    py_zeros = [
+        False,
+        0,
+        0.0,
+        complex(0, 0),
+    ]
+
+    py_ones = [
+        True,
+        1,
+        1.0,
+        complex(1, 0),
+    ]
+
+    for zero, one, arr in zip(py_zeros, py_ones, arrs):
+        n = 1
+        r = dpt.diff(arr, n=n, axis=axis, prepend=zero, append=one)
+        assert isinstance(r, dpt.usm_ndarray)
+        assert all(
+            r.shape[i] == arr.shape[i] for i in range(arr.ndim) if i != axis
+        )
+        assert r.shape[axis] == arr.shape[axis] + 2 - n
+
+        r = dpt.diff(arr, n=n, axis=axis, prepend=zero)
+        assert isinstance(r, dpt.usm_ndarray)
+        assert all(
+            r.shape[i] == arr.shape[i] for i in range(arr.ndim) if i != axis
+        )
+        assert r.shape[axis] == arr.shape[axis] + 1 - n

From dbed59a6ac55b2841da9e98b7dde715a68d52a73 Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Wed, 31 Jul 2024 00:29:41 +0000
Subject: [PATCH 18/25] Corrects test_diff_no_op

---
 dpctl/tests/test_tensor_diff.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/dpctl/tests/test_tensor_diff.py b/dpctl/tests/test_tensor_diff.py
index ecde3ff402..e11ebb9548 100644
--- a/dpctl/tests/test_tensor_diff.py
+++ b/dpctl/tests/test_tensor_diff.py
@@ -150,7 +150,8 @@ def test_diff_no_op():
     res = dpt.diff(x, n=0)
     assert dpt.all(x == res)
 
-    res = dpt.diff(dpt.reshape(x, (2, 5)), n=0, axis=0)
+    x = dpt.reshape(x, (2, 5))
+    res = dpt.diff(x, n=0, axis=0)
     assert dpt.all(x == res)
 
 

From 33e437aa375374cf2510de872ed0b8efcc7119f6 Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Wed, 31 Jul 2024 17:01:08 +0000
Subject: [PATCH 19/25] Make diff test for appended Python scalars more
 thorough and more efficient

Also adds a correctness check
---
 dpctl/tests/test_tensor_diff.py | 83 +++++++++++++++++++--------------
 1 file changed, 47 insertions(+), 36 deletions(-)

diff --git a/dpctl/tests/test_tensor_diff.py b/dpctl/tests/test_tensor_diff.py
index e11ebb9548..037c96ce0a 100644
--- a/dpctl/tests/test_tensor_diff.py
+++ b/dpctl/tests/test_tensor_diff.py
@@ -159,39 +159,50 @@ def test_diff_no_op():
 def test_diff_prepend_append_py_scalars(sh, axis):
     get_queue_or_skip()
 
-    arrs = [
-        dpt.ones(sh, dtype="?"),
-        dpt.ones(sh, dtype="i4"),
-        dpt.ones(sh, dtype="f4"),
-        dpt.ones(sh, dtype="c8"),
-    ]
-
-    py_zeros = [
-        False,
-        0,
-        0.0,
-        complex(0, 0),
-    ]
-
-    py_ones = [
-        True,
-        1,
-        1.0,
-        complex(1, 0),
-    ]
-
-    for zero, one, arr in zip(py_zeros, py_ones, arrs):
-        n = 1
-        r = dpt.diff(arr, n=n, axis=axis, prepend=zero, append=one)
-        assert isinstance(r, dpt.usm_ndarray)
-        assert all(
-            r.shape[i] == arr.shape[i] for i in range(arr.ndim) if i != axis
-        )
-        assert r.shape[axis] == arr.shape[axis] + 2 - n
-
-        r = dpt.diff(arr, n=n, axis=axis, prepend=zero)
-        assert isinstance(r, dpt.usm_ndarray)
-        assert all(
-            r.shape[i] == arr.shape[i] for i in range(arr.ndim) if i != axis
-        )
-        assert r.shape[axis] == arr.shape[axis] + 1 - n
+    n = 1
+
+    arr = dpt.ones(sh, dtype="i4")
+    zero = 0
+
+    # first and last elements along axis
+    # will be checked for correctness
+    sl1 = [slice(None)] * arr.ndim
+    sl1[axis] = slice(1)
+    sl1 = tuple(sl1)
+
+    sl2 = [slice(None)] * arr.ndim
+    sl2[axis] = slice(-1, None, None)
+    sl2 = tuple(sl2)
+
+    r = dpt.diff(arr, axis=axis, prepend=zero, append=zero)
+    assert isinstance(r, dpt.usm_ndarray)
+    assert all(r.shape[i] == arr.shape[i] for i in range(arr.ndim) if i != axis)
+    assert r.shape[axis] == arr.shape[axis] + 2 - n
+    assert dpt.all(r[sl1] == 1)
+    assert dpt.all(r[sl2] == -1)
+
+    r = dpt.diff(arr, axis=axis, prepend=zero)
+    assert isinstance(r, dpt.usm_ndarray)
+    assert all(r.shape[i] == arr.shape[i] for i in range(arr.ndim) if i != axis)
+    assert r.shape[axis] == arr.shape[axis] + 1 - n
+    assert dpt.all(r[sl1] == 1)
+
+    r = dpt.diff(arr, axis=axis, append=zero)
+    assert isinstance(r, dpt.usm_ndarray)
+    assert all(r.shape[i] == arr.shape[i] for i in range(arr.ndim) if i != axis)
+    assert r.shape[axis] == arr.shape[axis] + 1 - n
+    assert dpt.all(r[sl2] == -1)
+
+    r = dpt.diff(arr, axis=axis, prepend=dpt.asarray(zero), append=zero)
+    assert isinstance(r, dpt.usm_ndarray)
+    assert all(r.shape[i] == arr.shape[i] for i in range(arr.ndim) if i != axis)
+    assert r.shape[axis] == arr.shape[axis] + 2 - n
+    assert dpt.all(r[sl1] == 1)
+    assert dpt.all(r[sl2] == -1)
+
+    r = dpt.diff(arr, axis=axis, prepend=zero, append=dpt.asarray(zero))
+    assert isinstance(r, dpt.usm_ndarray)
+    assert all(r.shape[i] == arr.shape[i] for i in range(arr.ndim) if i != axis)
+    assert r.shape[axis] == arr.shape[axis] + 2 - n
+    assert dpt.all(r[sl1] == 1)
+    assert dpt.all(r[sl2] == -1)

From 52a8d2337dc90152d4973d806f7e2c99c65269e5 Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Wed, 31 Jul 2024 18:07:53 +0000
Subject: [PATCH 20/25] `diff` input validation tests

---
 dpctl/tests/test_tensor_diff.py | 153 ++++++++++++++++++++++++++++++++
 1 file changed, 153 insertions(+)

diff --git a/dpctl/tests/test_tensor_diff.py b/dpctl/tests/test_tensor_diff.py
index 037c96ce0a..fd90052eb5 100644
--- a/dpctl/tests/test_tensor_diff.py
+++ b/dpctl/tests/test_tensor_diff.py
@@ -14,11 +14,15 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
+from math import prod
+
 import pytest
+from numpy.testing import assert_raises_regex
 
 import dpctl.tensor as dpt
 from dpctl.tensor._type_utils import _to_device_supported_dtype
 from dpctl.tests.helper import get_queue_or_skip, skip_if_dtype_not_supported
+from dpctl.utils import ExecutionPlacementError
 
 _all_dtypes = [
     "?",
@@ -206,3 +210,152 @@ def test_diff_prepend_append_py_scalars(sh, axis):
     assert r.shape[axis] == arr.shape[axis] + 2 - n
     assert dpt.all(r[sl1] == 1)
     assert dpt.all(r[sl2] == -1)
+
+
+def test_tensor_diff_append_prepend_arrays():
+    get_queue_or_skip()
+
+    n = 1
+    axis = 0
+
+    sz = 5
+    arr = dpt.arange(sz, 2 * sz, dtype="i4")
+    prepend = dpt.arange(sz, dtype="i4")
+    append = dpt.arange(2 * sz, 3 * sz, dtype="i4")
+    const_diff = 1
+
+    r = dpt.diff(arr, axis=axis, prepend=prepend, append=append)
+    assert all(r.shape[i] == arr.shape[i] for i in range(arr.ndim) if i != axis)
+    assert (
+        r.shape[axis]
+        == arr.shape[axis] + prepend.shape[axis] + append.shape[axis] - n
+    )
+    assert dpt.all(r == const_diff)
+
+    r = dpt.diff(arr, axis=axis, prepend=prepend)
+    assert all(r.shape[i] == arr.shape[i] for i in range(arr.ndim) if i != axis)
+    assert r.shape[axis] == arr.shape[axis] + prepend.shape[axis] - n
+    assert dpt.all(r == const_diff)
+
+    r = dpt.diff(arr, axis=axis, append=append)
+    assert all(r.shape[i] == arr.shape[i] for i in range(arr.ndim) if i != axis)
+    assert r.shape[axis] == arr.shape[axis] + append.shape[axis] - n
+    assert dpt.all(r == const_diff)
+
+    sh = (3, 4, 5)
+    sz = prod(sh)
+    arr = dpt.reshape(dpt.arange(sz, 2 * sz, dtype="i4"), sh)
+    prepend = dpt.reshape(dpt.arange(sz, dtype="i4"), sh)
+    append = dpt.reshape(dpt.arange(2 * sz, 3 * sz, dtype="i4"), sh)
+    const_diff = prod(sh[axis + 1 :])
+
+    r = dpt.diff(arr, axis=axis, prepend=prepend, append=append)
+    assert all(r.shape[i] == arr.shape[i] for i in range(arr.ndim) if i != axis)
+    assert (
+        r.shape[axis]
+        == arr.shape[axis] + prepend.shape[axis] + append.shape[axis] - n
+    )
+    assert dpt.all(r == const_diff)
+
+    r = dpt.diff(arr, axis=axis, prepend=prepend)
+    assert all(r.shape[i] == arr.shape[i] for i in range(arr.ndim) if i != axis)
+    assert r.shape[axis] == arr.shape[axis] + prepend.shape[axis] - n
+    assert dpt.all(r == const_diff)
+
+    r = dpt.diff(arr, axis=axis, append=append)
+    assert all(r.shape[i] == arr.shape[i] for i in range(arr.ndim) if i != axis)
+    assert r.shape[axis] == arr.shape[axis] + append.shape[axis] - n
+    assert dpt.all(r == const_diff)
+
+
+def test_diff_wrong_append_prepend_shape():
+    get_queue_or_skip()
+
+    arr = dpt.ones((3, 4, 5), dtype="i4")
+    arr_bad_sh = dpt.ones(2, dtype="i4")
+
+    assert_raises_regex(
+        ValueError,
+        "`diff` argument `prepend` with shape.*is invalid"
+        " for first input with shape.*",
+        dpt.diff,
+        arr,
+        prepend=arr_bad_sh,
+        append=arr_bad_sh,
+    )
+
+    assert_raises_regex(
+        ValueError,
+        "`diff` argument `append` with shape.*is invalid"
+        " for first input with shape.*",
+        dpt.diff,
+        arr,
+        prepend=arr,
+        append=arr_bad_sh,
+    )
+
+    assert_raises_regex(
+        ValueError,
+        "`diff` argument `prepend` with shape.*is invalid"
+        " for first input with shape.*",
+        dpt.diff,
+        arr,
+        prepend=arr_bad_sh,
+    )
+
+    assert_raises_regex(
+        ValueError,
+        "`diff` argument `append` with shape.*is invalid"
+        " for first input with shape.*",
+        dpt.diff,
+        arr,
+        append=arr_bad_sh,
+    )
+
+
+def test_diff_compute_follows_data():
+    q1 = get_queue_or_skip()
+    q2 = get_queue_or_skip()
+    q3 = get_queue_or_skip()
+
+    ar1 = dpt.ones(1, dtype="i4", sycl_queue=q1)
+    ar2 = dpt.ones(1, dtype="i4", sycl_queue=q2)
+    ar3 = dpt.ones(1, dtype="i4", sycl_queue=q3)
+
+    assert_raises_regex(
+        ExecutionPlacementError,
+        "Execution placement can not be unambiguously inferred from input "
+        "arguments",
+        dpt.diff,
+        ar1,
+        prepend=ar2,
+        append=ar3,
+    )
+
+    assert_raises_regex(
+        ExecutionPlacementError,
+        "Execution placement can not be unambiguously inferred from input "
+        "arguments",
+        dpt.diff,
+        ar1,
+        prepend=ar2,
+    )
+
+    assert_raises_regex(
+        ExecutionPlacementError,
+        "Execution placement can not be unambiguously inferred from input "
+        "arguments",
+        dpt.diff,
+        ar1,
+        append=ar2,
+    )
+
+
+def test_diff_input_validation():
+    bad_in = dict()
+    assert_raises_regex(
+        TypeError,
+        "Expecting dpctl.tensor.usm_ndarray type, got.*",
+        dpt.diff,
+        bad_in,
+    )

From 440fb775cfb6adf532fbaa1c3f5ace560e0ab714 Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Wed, 31 Jul 2024 13:14:32 -0700
Subject: [PATCH 21/25] Refactor diff tests and improve coverage

---
 dpctl/tests/test_tensor_diff.py | 141 +++++++++++++++-----------------
 1 file changed, 66 insertions(+), 75 deletions(-)

diff --git a/dpctl/tests/test_tensor_diff.py b/dpctl/tests/test_tensor_diff.py
index fd90052eb5..5f00e9b956 100644
--- a/dpctl/tests/test_tensor_diff.py
+++ b/dpctl/tests/test_tensor_diff.py
@@ -48,22 +48,18 @@ def test_diff_basic(dt):
     skip_if_dtype_not_supported(dt, q)
 
     x = dpt.asarray([9, 12, 7, 17, 10, 18, 15, 9, 8, 8], dtype=dt, sycl_queue=q)
-    res = dpt.diff(x)
     op = dpt.not_equal if x.dtype is dpt.bool else dpt.subtract
-    expected_res = op(x[1:], x[:-1])
-    if dpt.dtype(dt).kind in "fc":
-        assert dpt.allclose(res, expected_res)
-    else:
-        assert dpt.all(res == expected_res)
 
-    res = dpt.diff(x, n=5)
-    expected_res = x
-    for _ in range(5):
-        expected_res = op(expected_res[1:], expected_res[:-1])
-    if dpt.dtype(dt).kind in "fc":
-        assert dpt.allclose(res, expected_res)
-    else:
-        assert dpt.all(res == expected_res)
+    # test both n=2 and n>2 branches
+    for n in [1, 2, 5]:
+        res = dpt.diff(x, n=n)
+        expected_res = x
+        for _ in range(n):
+            expected_res = op(expected_res[1:], expected_res[:-1])
+        if dpt.dtype(dt).kind in "fc":
+            assert dpt.allclose(res, expected_res)
+        else:
+            assert dpt.all(res == expected_res)
 
 
 def test_diff_axis():
@@ -73,17 +69,15 @@ def test_diff_axis():
         dpt.asarray([9, 12, 7, 17, 10, 18, 15, 9, 8, 8], dtype="i4"), (3, 4, 1)
     )
     x[:, ::2, :] = 0
-    res = dpt.diff(x, n=1, axis=1)
-    expected_res = dpt.subtract(x[:, 1:, :], x[:, :-1, :])
-    assert dpt.all(res == expected_res)
-
-    res = dpt.diff(x, n=3, axis=1)
-    expected_res = x
-    for _ in range(3):
-        expected_res = dpt.subtract(
-            expected_res[:, 1:, :], expected_res[:, :-1, :]
-        )
-    assert dpt.all(res == expected_res)
+
+    for n in [1, 2, 3]:
+        res = dpt.diff(x, n=3, axis=1)
+        expected_res = x
+        for _ in range(3):
+            expected_res = dpt.subtract(
+                expected_res[:, 1:, :], expected_res[:, :-1, :]
+            )
+        assert dpt.all(res == expected_res)
 
 
 def test_diff_prepend_append_type_promotion():
@@ -179,33 +173,28 @@ def test_diff_prepend_append_py_scalars(sh, axis):
     sl2 = tuple(sl2)
 
     r = dpt.diff(arr, axis=axis, prepend=zero, append=zero)
-    assert isinstance(r, dpt.usm_ndarray)
     assert all(r.shape[i] == arr.shape[i] for i in range(arr.ndim) if i != axis)
     assert r.shape[axis] == arr.shape[axis] + 2 - n
     assert dpt.all(r[sl1] == 1)
     assert dpt.all(r[sl2] == -1)
 
     r = dpt.diff(arr, axis=axis, prepend=zero)
-    assert isinstance(r, dpt.usm_ndarray)
     assert all(r.shape[i] == arr.shape[i] for i in range(arr.ndim) if i != axis)
     assert r.shape[axis] == arr.shape[axis] + 1 - n
     assert dpt.all(r[sl1] == 1)
 
     r = dpt.diff(arr, axis=axis, append=zero)
-    assert isinstance(r, dpt.usm_ndarray)
     assert all(r.shape[i] == arr.shape[i] for i in range(arr.ndim) if i != axis)
     assert r.shape[axis] == arr.shape[axis] + 1 - n
     assert dpt.all(r[sl2] == -1)
 
     r = dpt.diff(arr, axis=axis, prepend=dpt.asarray(zero), append=zero)
-    assert isinstance(r, dpt.usm_ndarray)
     assert all(r.shape[i] == arr.shape[i] for i in range(arr.ndim) if i != axis)
     assert r.shape[axis] == arr.shape[axis] + 2 - n
     assert dpt.all(r[sl1] == 1)
     assert dpt.all(r[sl2] == -1)
 
     r = dpt.diff(arr, axis=axis, prepend=zero, append=dpt.asarray(zero))
-    assert isinstance(r, dpt.usm_ndarray)
     assert all(r.shape[i] == arr.shape[i] for i in range(arr.ndim) if i != axis)
     assert r.shape[axis] == arr.shape[axis] + 2 - n
     assert dpt.all(r[sl1] == 1)
@@ -218,54 +207,36 @@ def test_tensor_diff_append_prepend_arrays():
     n = 1
     axis = 0
 
-    sz = 5
-    arr = dpt.arange(sz, 2 * sz, dtype="i4")
-    prepend = dpt.arange(sz, dtype="i4")
-    append = dpt.arange(2 * sz, 3 * sz, dtype="i4")
-    const_diff = 1
-
-    r = dpt.diff(arr, axis=axis, prepend=prepend, append=append)
-    assert all(r.shape[i] == arr.shape[i] for i in range(arr.ndim) if i != axis)
-    assert (
-        r.shape[axis]
-        == arr.shape[axis] + prepend.shape[axis] + append.shape[axis] - n
-    )
-    assert dpt.all(r == const_diff)
-
-    r = dpt.diff(arr, axis=axis, prepend=prepend)
-    assert all(r.shape[i] == arr.shape[i] for i in range(arr.ndim) if i != axis)
-    assert r.shape[axis] == arr.shape[axis] + prepend.shape[axis] - n
-    assert dpt.all(r == const_diff)
-
-    r = dpt.diff(arr, axis=axis, append=append)
-    assert all(r.shape[i] == arr.shape[i] for i in range(arr.ndim) if i != axis)
-    assert r.shape[axis] == arr.shape[axis] + append.shape[axis] - n
-    assert dpt.all(r == const_diff)
-
-    sh = (3, 4, 5)
-    sz = prod(sh)
-    arr = dpt.reshape(dpt.arange(sz, 2 * sz, dtype="i4"), sh)
-    prepend = dpt.reshape(dpt.arange(sz, dtype="i4"), sh)
-    append = dpt.reshape(dpt.arange(2 * sz, 3 * sz, dtype="i4"), sh)
-    const_diff = prod(sh[axis + 1 :])
+    for sh in [(5,), (3, 4, 5)]:
+        sz = prod(sh)
+        arr = dpt.reshape(dpt.arange(sz, 2 * sz, dtype="i4"), sh)
+        prepend = dpt.reshape(dpt.arange(sz, dtype="i4"), sh)
+        append = dpt.reshape(dpt.arange(2 * sz, 3 * sz, dtype="i4"), sh)
+        const_diff = sz / sh[axis]
 
-    r = dpt.diff(arr, axis=axis, prepend=prepend, append=append)
-    assert all(r.shape[i] == arr.shape[i] for i in range(arr.ndim) if i != axis)
-    assert (
-        r.shape[axis]
-        == arr.shape[axis] + prepend.shape[axis] + append.shape[axis] - n
-    )
-    assert dpt.all(r == const_diff)
+        r = dpt.diff(arr, axis=axis, prepend=prepend, append=append)
+        assert all(
+            r.shape[i] == arr.shape[i] for i in range(arr.ndim) if i != axis
+        )
+        assert (
+            r.shape[axis]
+            == arr.shape[axis] + prepend.shape[axis] + append.shape[axis] - n
+        )
+        assert dpt.all(r == const_diff)
 
-    r = dpt.diff(arr, axis=axis, prepend=prepend)
-    assert all(r.shape[i] == arr.shape[i] for i in range(arr.ndim) if i != axis)
-    assert r.shape[axis] == arr.shape[axis] + prepend.shape[axis] - n
-    assert dpt.all(r == const_diff)
+        r = dpt.diff(arr, axis=axis, prepend=prepend)
+        assert all(
+            r.shape[i] == arr.shape[i] for i in range(arr.ndim) if i != axis
+        )
+        assert r.shape[axis] == arr.shape[axis] + prepend.shape[axis] - n
+        assert dpt.all(r == const_diff)
 
-    r = dpt.diff(arr, axis=axis, append=append)
-    assert all(r.shape[i] == arr.shape[i] for i in range(arr.ndim) if i != axis)
-    assert r.shape[axis] == arr.shape[axis] + append.shape[axis] - n
-    assert dpt.all(r == const_diff)
+        r = dpt.diff(arr, axis=axis, append=append)
+        assert all(
+            r.shape[i] == arr.shape[i] for i in range(arr.ndim) if i != axis
+        )
+        assert r.shape[axis] == arr.shape[axis] + append.shape[axis] - n
+        assert dpt.all(r == const_diff)
 
 
 def test_diff_wrong_append_prepend_shape():
@@ -332,6 +303,26 @@ def test_diff_compute_follows_data():
         append=ar3,
     )
 
+    assert_raises_regex(
+        ExecutionPlacementError,
+        "Execution placement can not be unambiguously inferred from input "
+        "arguments",
+        dpt.diff,
+        ar1,
+        prepend=ar2,
+        append=0,
+    )
+
+    assert_raises_regex(
+        ExecutionPlacementError,
+        "Execution placement can not be unambiguously inferred from input "
+        "arguments",
+        dpt.diff,
+        ar1,
+        prepend=0,
+        append=ar2,
+    )
+
     assert_raises_regex(
         ExecutionPlacementError,
         "Execution placement can not be unambiguously inferred from input "

From 16a71c355a021659923f31c55c6a57ff88ceade2 Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Thu, 1 Aug 2024 08:32:01 -0700
Subject: [PATCH 22/25] Fixes per review by @oleksandr-pavlyk

---
 dpctl/tests/test_tensor_diff.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/dpctl/tests/test_tensor_diff.py b/dpctl/tests/test_tensor_diff.py
index 5f00e9b956..765858d0cb 100644
--- a/dpctl/tests/test_tensor_diff.py
+++ b/dpctl/tests/test_tensor_diff.py
@@ -71,9 +71,9 @@ def test_diff_axis():
     x[:, ::2, :] = 0
 
     for n in [1, 2, 3]:
-        res = dpt.diff(x, n=3, axis=1)
+        res = dpt.diff(x, n=n, axis=1)
         expected_res = x
-        for _ in range(3):
+        for _ in range(n):
             expected_res = dpt.subtract(
                 expected_res[:, 1:, :], expected_res[:, :-1, :]
             )
@@ -90,10 +90,10 @@ def test_diff_prepend_append_type_promotion():
         ("i8", "c8", "u8"),
     ]
 
-    for _dts in dts:
-        x = dpt.ones(10, dtype=_dts[1])
-        prepend = dpt.full(1, 2, dtype=_dts[0])
-        append = dpt.full(1, 3, dtype=_dts[2])
+    for dt0, dt1, dt2 in dts:
+        x = dpt.ones(10, dtype=dt1)
+        prepend = dpt.full(1, 2, dtype=dt0)
+        append = dpt.full(1, 3, dtype=dt2)
 
         res = dpt.diff(x, prepend=prepend, append=append)
         assert res.dtype == _to_device_supported_dtype(

From 3215c9ff2cd547bc7075f72d2b845d1e29b496d8 Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Thu, 1 Aug 2024 09:04:28 -0700
Subject: [PATCH 23/25] More PR review changes

---
 dpctl/tests/test_tensor_diff.py | 65 +++++++--------------------------
 1 file changed, 14 insertions(+), 51 deletions(-)

diff --git a/dpctl/tests/test_tensor_diff.py b/dpctl/tests/test_tensor_diff.py
index 765858d0cb..ff2f104515 100644
--- a/dpctl/tests/test_tensor_diff.py
+++ b/dpctl/tests/test_tensor_diff.py
@@ -247,8 +247,7 @@ def test_diff_wrong_append_prepend_shape():
 
     assert_raises_regex(
         ValueError,
-        "`diff` argument `prepend` with shape.*is invalid"
-        " for first input with shape.*",
+        ".*shape.*is invalid.*",
         dpt.diff,
         arr,
         prepend=arr_bad_sh,
@@ -257,8 +256,7 @@ def test_diff_wrong_append_prepend_shape():
 
     assert_raises_regex(
         ValueError,
-        "`diff` argument `append` with shape.*is invalid"
-        " for first input with shape.*",
+        ".*shape.*is invalid.*",
         dpt.diff,
         arr,
         prepend=arr,
@@ -267,8 +265,7 @@ def test_diff_wrong_append_prepend_shape():
 
     assert_raises_regex(
         ValueError,
-        "`diff` argument `prepend` with shape.*is invalid"
-        " for first input with shape.*",
+        ".*shape.*is invalid.*",
         dpt.diff,
         arr,
         prepend=arr_bad_sh,
@@ -276,8 +273,7 @@ def test_diff_wrong_append_prepend_shape():
 
     assert_raises_regex(
         ValueError,
-        "`diff` argument `append` with shape.*is invalid"
-        " for first input with shape.*",
+        ".*shape.*is invalid.*",
         dpt.diff,
         arr,
         append=arr_bad_sh,
@@ -293,53 +289,20 @@ def test_diff_compute_follows_data():
     ar2 = dpt.ones(1, dtype="i4", sycl_queue=q2)
     ar3 = dpt.ones(1, dtype="i4", sycl_queue=q3)
 
-    assert_raises_regex(
-        ExecutionPlacementError,
-        "Execution placement can not be unambiguously inferred from input "
-        "arguments",
-        dpt.diff,
-        ar1,
-        prepend=ar2,
-        append=ar3,
-    )
+    with pytest.raises(ExecutionPlacementError):
+        dpt.diff(ar1, prepend=ar2, append=ar3)
 
-    assert_raises_regex(
-        ExecutionPlacementError,
-        "Execution placement can not be unambiguously inferred from input "
-        "arguments",
-        dpt.diff,
-        ar1,
-        prepend=ar2,
-        append=0,
-    )
+    with pytest.raises(ExecutionPlacementError):
+        dpt.diff(ar1, prepend=ar2, append=0)
 
-    assert_raises_regex(
-        ExecutionPlacementError,
-        "Execution placement can not be unambiguously inferred from input "
-        "arguments",
-        dpt.diff,
-        ar1,
-        prepend=0,
-        append=ar2,
-    )
+    with pytest.raises(ExecutionPlacementError):
+        dpt.diff(ar1, prepend=0, append=ar2)
 
-    assert_raises_regex(
-        ExecutionPlacementError,
-        "Execution placement can not be unambiguously inferred from input "
-        "arguments",
-        dpt.diff,
-        ar1,
-        prepend=ar2,
-    )
+    with pytest.raises(ExecutionPlacementError):
+        dpt.diff(ar1, prepend=ar2)
 
-    assert_raises_regex(
-        ExecutionPlacementError,
-        "Execution placement can not be unambiguously inferred from input "
-        "arguments",
-        dpt.diff,
-        ar1,
-        append=ar2,
-    )
+    with pytest.raises(ExecutionPlacementError):
+        dpt.diff(ar1, append=ar2)
 
 
 def test_diff_input_validation():

From 89c1339a1b7f49bbf046df91d7cdf9d802ab7287 Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Thu, 1 Aug 2024 15:21:50 -0700
Subject: [PATCH 24/25] Adds `count_nonzero`, `diff`, and `nextafter` to docs

---
 .../api_reference/dpctl/tensor.elementwise_functions.rst         | 1 +
 .../api_reference/dpctl/tensor.searching_functions.rst           | 1 +
 .../doc_sources/api_reference/dpctl/tensor.utility_functions.rst | 1 +
 3 files changed, 3 insertions(+)

diff --git a/docs/doc_sources/api_reference/dpctl/tensor.elementwise_functions.rst b/docs/doc_sources/api_reference/dpctl/tensor.elementwise_functions.rst
index 9c4706e356..8edf7d5845 100644
--- a/docs/doc_sources/api_reference/dpctl/tensor.elementwise_functions.rst
+++ b/docs/doc_sources/api_reference/dpctl/tensor.elementwise_functions.rst
@@ -64,6 +64,7 @@ function values computed for every element of input array(s).
     minimum
     multiply
     negative
+    nextafter
     not_equal
     positive
     pow
diff --git a/docs/doc_sources/api_reference/dpctl/tensor.searching_functions.rst b/docs/doc_sources/api_reference/dpctl/tensor.searching_functions.rst
index 368437d160..5a784c4c1c 100644
--- a/docs/doc_sources/api_reference/dpctl/tensor.searching_functions.rst
+++ b/docs/doc_sources/api_reference/dpctl/tensor.searching_functions.rst
@@ -10,6 +10,7 @@ Searching functions
 
     argmax
     argmin
+    count_nonzero
     nonzero
     searchsorted
     where
diff --git a/docs/doc_sources/api_reference/dpctl/tensor.utility_functions.rst b/docs/doc_sources/api_reference/dpctl/tensor.utility_functions.rst
index 121b82fc40..adbd3cf285 100644
--- a/docs/doc_sources/api_reference/dpctl/tensor.utility_functions.rst
+++ b/docs/doc_sources/api_reference/dpctl/tensor.utility_functions.rst
@@ -11,6 +11,7 @@ Utility functions
     all
     any
     allclose
+    diff
 
 Device object
 -------------

From 6520cc09b9a1265cc741f7e57e15c977bd2c47ad Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Thu, 1 Aug 2024 16:21:45 -0700
Subject: [PATCH 25/25] Change to `diff` docstring per PR review

---
 dpctl/tensor/_utility_functions.py | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/dpctl/tensor/_utility_functions.py b/dpctl/tensor/_utility_functions.py
index d0e6b8f9bc..642b3fec0b 100644
--- a/dpctl/tensor/_utility_functions.py
+++ b/dpctl/tensor/_utility_functions.py
@@ -440,10 +440,8 @@ def diff(x, /, *, axis=-1, n=1, prepend=None, append=None):
     Returns:
         usm_ndarray:
             an array containing the `n`-th differences. The array will have the
-            same shape as `x`, except along `axis`, which will have shape
-
-            - prepend.shape[axis] + x.shape[axis] + append.shape[axis] - n
-
+            same shape as `x`, except along `axis`, which will have shape:
+                prepend.shape[axis] + x.shape[axis] + append.shape[axis] - n
             The data type of the returned array is determined by the Type
             Promotion Rules.
     """