Merge pull request #1732 from IntelPython/diff-count-nonzero-array-api

ndgrigorian · web-flow · commit 1ecd8a810d2a · 2024-08-01T17:44:22.000-07:00
Implements `dpctl.tensor.count_nonzero` and `dpctl.tensor.diff`
diff --git 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
@@ -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
@@ -11,6 +11,7 @@ Utility functions
     all
     any
     allclose
+    diff
 
 Device object
 -------------
diff --git 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__
@@ -176,6 +176,7 @@
 from ._reduction import (
     argmax,
     argmin,
+    count_nonzero,
     logsumexp,
     max,
     min,
@@ -373,4 +374,6 @@
     "cumulative_prod",
     "cumulative_sum",
     "nextafter",
+    "diff",
+    "count_nonzero",
 ]
diff --git 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/_reduction.py b/dpctl/tensor/_reduction.py
@@ -773,3 +773,46 @@ 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):
+    """
+    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(
+        x,
+        axis=axis,
+        dtype=ti.default_device_index_type(x.sycl_device),
+        keepdims=keepdims,
+        out=out,
+    )
diff --git 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
diff --git a/dpctl/tests/test_tensor_diff.py b/dpctl/tests/test_tensor_diff.py
diff --git a/dpctl/tests/test_usm_ndarray_reductions.py b/dpctl/tests/test_usm_ndarray_reductions.py
