From a99f44e49b4d52421a3865a9180698e60fa9ffd2 Mon Sep 17 00:00:00 2001 From: Rohan Yadav Date: Mon, 12 Feb 2024 16:08:38 -0800 Subject: [PATCH] cunumeric: addressing review comments --- cunumeric/module.py | 24 +++++++++--------------- tests/integration/test_diff.py | 16 ++++++++-------- 2 files changed, 17 insertions(+), 23 deletions(-) diff --git a/cunumeric/module.py b/cunumeric/module.py index d51fe580b..c2f64aee7 100644 --- a/cunumeric/module.py +++ b/cunumeric/module.py @@ -6242,13 +6242,13 @@ def nansum( # Arithmetic operations -@add_boilerplate("a") +@add_boilerplate("a", "prepend", "append") def diff( a: ndarray, n: int = 1, axis: int = -1, - prepend: Any = None, - append: Any = None, + prepend: ndarray | None = None, + append: ndarray | None = None, ) -> ndarray: """ Calculate the n-th discrete difference along the given axis. @@ -6278,11 +6278,10 @@ def diff( 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. A notable exception is `datetime64`, which - results in a `timedelta64` output array. + `a` in most cases. See Also -------- - gradient, ediff1d, cumsum + numpy.diff Notes ----- Type is preserved for boolean arrays, so the result will contain @@ -6314,9 +6313,6 @@ def diff( [5, 1, 2]]) >>> np.diff(x, axis=0) array([[-1, 2, 0, -2]]) - >>> x = np.arange('1066-10-13', '1066-10-16', dtype=np.datetime64) - >>> np.diff(x) - array([1, 1], dtype='timedelta64[D]') Availability -------- @@ -6336,25 +6332,23 @@ def diff( combined = [] if prepend is not None: - prepend = np.asanyarray(prepend) if prepend.ndim == 0: shape = list(a.shape) shape[axis] = 1 - prepend = np.broadcast_to(prepend, tuple(shape)) + prepend = broadcast_to(prepend, tuple(shape)) combined.append(prepend) combined.append(a) if append is not None: - append = np.asanyarray(append) if append.ndim == 0: shape = list(a.shape) shape[axis] = 1 - append = np.broadcast_to(append, tuple(shape)) + append = broadcast_to(append, tuple(shape)) combined.append(append) if len(combined) > 1: - a = np.concatenate(combined, axis) + 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 @@ -6362,7 +6356,7 @@ def diff( if a.shape[axis] <= n: shape = list(a.shape) shape[axis] = 0 - return np.empty(shape=shape, dtype=a.dtype) + return empty(shape=shape, dtype=a.dtype) slice1l = [slice(None)] * nd slice2l = [slice(None)] * nd diff --git a/tests/integration/test_diff.py b/tests/integration/test_diff.py index 06d17e666..b1e62b7cd 100644 --- a/tests/integration/test_diff.py +++ b/tests/integration/test_diff.py @@ -17,7 +17,7 @@ import pytest from utils.comparisons import allclose -import cunumeric as cn +import cunumeric as num @pytest.mark.parametrize( @@ -41,22 +41,22 @@ ) def test_diff(args): shape, n, axis, prepend, append = args - num = np.random.random(shape) - cun = cn.array(num) + 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_num = np.diff(num, n=n, axis=axis, prepend=n_prepend, append=n_append) - res_cn = cn.diff(cun, n=n, axis=axis, prepend=prepend, append=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_num, res_cn) + assert allclose(res_np, res_cn) def test_diff_nzero(): - a = cn.ones(100) - ad = cn.diff(a, n=0) + a = num.ones(100) + ad = num.diff(a, n=0) assert a is ad