cunumeric: implement np.diff (#636)

* cunumeric: implement `np.diff`

This commit implements some necessary cases of `np.diff`.

Signed-off-by: Rohan Yadav <[email protected]>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* cunumeric: fix missing import

* cunumeric: addressing review comments

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* cunumeric: appease the typechecker

---------

Signed-off-by: Rohan Yadav <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Rohan Yadav <[email protected]>

cunumeric/module.py

@@ -42,9 +42,9 @@
 
 from cunumeric.coverage import is_implemented
 
-from ._ufunc.comparison import maximum, minimum
+from ._ufunc.comparison import maximum, minimum, not_equal
 from ._ufunc.floating import floor, isnan
-from ._ufunc.math import add, multiply
+from ._ufunc.math import add, multiply, subtract
 from ._unary_red_utils import get_non_nan_unary_red_code
 from .array import (
     add_boilerplate,
@@ -6242,6 +6242,136 @@ def nansum(
 # Arithmetic operations
 
 
+@add_boilerplate("a", "prepend", "append")
+def diff(
+    a: ndarray,
+    n: int = 1,
+    axis: int = -1,
+    prepend: ndarray | None = None,
+    append: ndarray | None = None,
+) -> ndarray:
+    """
+    Calculate the n-th discrete difference along the given axis.
+    The first difference is given by ``out[i] = a[i+1] - a[i]`` along
+    the given axis, higher differences are calculated by using `diff`
+    recursively.
+    Parameters
+    ----------
+    a : array_like
+        Input array
+    n : int, optional
+        The number of times values are differenced. If zero, the input
+        is returned as-is.
+    axis : int, optional
+        The axis along which the difference is taken, default is the
+        last axis.
+    prepend, append : array_like, optional
+        Values to prepend or append to `a` along axis prior to
+        performing the difference.  Scalar values are expanded to
+        arrays with length 1 in the direction of axis and the shape
+        of the input array in along all other axes.  Otherwise the
+        dimension and shape must match `a` except along axis.
+    Returns
+    -------
+    diff : ndarray
+        The n-th differences. The shape of the output is the same as `a`
+        except along `axis` where the dimension is smaller by `n`. The
+        type of the output is the same as the type of the difference
+        between any two elements of `a`. This is the same as the type of
+        `a` in most cases.
+    See Also
+    --------
+    numpy.diff
+    Notes
+    -----
+    Type is preserved for boolean arrays, so the result will contain
+    `False` when consecutive elements are the same and `True` when they
+    differ.
+    For unsigned integer arrays, the results will also be unsigned. This
+    should not be surprising, as the result is consistent with
+    calculating the difference directly:
+    >>> u8_arr = np.array([1, 0], dtype=np.uint8)
+    >>> np.diff(u8_arr)
+    array([255], dtype=uint8)
+    >>> u8_arr[1,...] - u8_arr[0,...]
+    255
+    If this is not desirable, then the array should be cast to a larger
+    integer type first:
+    >>> i16_arr = u8_arr.astype(np.int16)
+    >>> np.diff(i16_arr)
+    array([-1], dtype=int16)
+    Examples
+    --------
+    >>> x = np.array([1, 2, 4, 7, 0])
+    >>> np.diff(x)
+    array([ 1,  2,  3, -7])
+    >>> np.diff(x, n=2)
+    array([  1,   1, -10])
+    >>> x = np.array([[1, 3, 6, 10], [0, 5, 6, 8]])
+    >>> np.diff(x)
+    array([[2, 3, 4],
+           [5, 1, 2]])
+    >>> np.diff(x, axis=0)
+    array([[-1,  2,  0, -2]])
+
+    Availability
+    --------
+    Multiple GPUs, Multiple CPUs
+    """
+    if n == 0:
+        return a
+    if n < 0:
+        raise ValueError("order must be non-negative but got " + repr(n))
+
+    nd = a.ndim
+    if nd == 0:
+        raise ValueError(
+            "diff requires input that is at least one dimensional"
+        )
+    axis = normalize_axis_index(axis, nd)
+
+    combined = []
+    if prepend is not None:
+        if prepend.ndim == 0:
+            shape = list(a.shape)
+            shape[axis] = 1
+            prepend = broadcast_to(prepend, tuple(shape))
+        combined.append(prepend)
+
+    combined.append(a)
+
+    if append is not None:
+        if append.ndim == 0:
+            shape = list(a.shape)
+            shape[axis] = 1
+            append = broadcast_to(append, tuple(shape))
+        combined.append(append)
+
+    if len(combined) > 1:
+        a = concatenate(combined, axis)
+
+    # Diffing with n > shape results in an empty array. We have
+    # to handle this case explicitly as our slicing routines raise
+    # an exception with out-of-bounds slices, while NumPy's dont.
+    if a.shape[axis] <= n:
+        shape = list(a.shape)
+        shape[axis] = 0
+        return empty(shape=tuple(shape), dtype=a.dtype)
+
+    slice1l = [slice(None)] * nd
+    slice2l = [slice(None)] * nd
+    slice1l[axis] = slice(1, None)
+    slice2l[axis] = slice(None, -1)
+    slice1 = tuple(slice1l)
+    slice2 = tuple(slice2l)
+
+    op = not_equal if a.dtype == np.bool_ else subtract
+    for _ in range(n):
+        a = op(a[slice1], a[slice2])
+
+    return a
+
+
 # Handling complex numbers
 
 

tests/integration/test_diff.py

@@ -0,0 +1,66 @@
+# Copyright 2022 NVIDIA 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 numpy as np
+import pytest
+from utils.comparisons import allclose
+
+import cunumeric as num
+
+
+@pytest.mark.parametrize(
+    "args",
+    [
+        ((100,), 1, -1, None, None),
+        ((100,), 2, -1, None, None),
+        ((100,), 3, -1, None, None),
+        ((100,), 2, 0, None, None),
+        ((10, 10), 2, -1, None, None),
+        ((10, 10), 2, 0, None, None),
+        ((10, 10), 2, 1, None, None),
+        ((100,), 3, -1, [1.0, 2.0], None),
+        ((100,), 3, -1, None, [1.0, 2.0]),
+        ((100,), 3, -1, [1.0, 2.0], [1.0, 2.0]),
+        ((5,), 5, -1, None, None),
+        ((5,), 6, 0, None, None),
+        ((5, 5), 5, 1, None, None),
+        ((5, 5), 6, 1, None, None),
+    ],
+)
+def test_diff(args):
+    shape, n, axis, prepend, append = args
+    nparr = np.random.random(shape)
+    cnarr = num.array(nparr)
+
+    # We are not adopting the np._NoValue default arguments
+    # for this function, as no special behavior is needed on None.
+    n_prepend = np._NoValue if prepend is None else prepend
+    n_append = np._NoValue if append is None else append
+    res_np = np.diff(nparr, n=n, axis=axis, prepend=n_prepend, append=n_append)
+    res_cn = num.diff(cnarr, n=n, axis=axis, prepend=prepend, append=append)
+
+    assert allclose(res_np, res_cn)
+
+
+def test_diff_nzero():
+    a = num.ones(100)
+    ad = num.diff(a, n=0)
+    assert a is ad
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(sys.argv))