Docstrings for diff and count_nonzero

Author: ndgrigorian

dpctl/tensor/_reduction.py

@@ -775,6 +775,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(

dpctl/tensor/_utility_functions.py

@@ -409,6 +409,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(
@@ -419,7 +456,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)