Merge branch 'branch-22.05' into bryanv/pytest_integration

cunumeric/array.py

@@ -1934,10 +1934,25 @@ def copy(self, order="C"):
     # currently offset option is implemented only for the case of number of
     # axes=2. This restriction can be lifted in the future if there is a
     # use case of having arbitrary number of offsets
-    def _diag_helper(self, offset=0, axes=None, extract=True):
+    def _diag_helper(
+        self,
+        offset=0,
+        axes=None,
+        extract=True,
+        trace=False,
+        out=None,
+        dtype=None,
+    ):
         # _diag_helper can be used only for arrays with dim>=1
         if self.ndim < 1:
             raise ValueError("_diag_helper is implemented for dim>=1")
+        # out should be passed only for Trace
+        if out is not None and not trace:
+            raise ValueError("_diag_helper supports out only for trace=True")
+        # dtype should be passed only for Trace
+        if dtype is not None and not trace:
+            raise ValueError("_diag_helper supports dtype only for trace=True")
+
         elif self.ndim == 1:
             if axes is not None:
                 raise ValueError(
@@ -1948,10 +1963,7 @@ def _diag_helper(self, offset=0, axes=None, extract=True):
             out = ndarray((m, m), dtype=self.dtype, inputs=(self,))
             diag_size = self.shape[0]
             out._thunk._diag_helper(
-                self._thunk,
-                offset=offset,
-                naxes=0,
-                extract=False,
+                self._thunk, offset=offset, naxes=0, extract=False, trace=False
             )
         else:
             N = len(axes)
@@ -2004,14 +2016,36 @@ def _diag_helper(self, offset=0, axes=None, extract=True):
 
             tr_shape = tuple(a.shape[i] for i in range(a.ndim - N))
             # calculate shape of the output array
-            out_shape = tr_shape + (diag_size,)
-            out = ndarray(shape=out_shape, dtype=self.dtype, inputs=(self,))
+            if trace:
+                if N != 2:
+                    raise ValueError(
+                        " exactly 2 axes should be passed to trace"
+                    )
+                if self.ndim == 2:
+                    out_shape = (1,)
+                elif self.ndim > 2:
+                    out_shape = tr_shape
+                else:
+                    raise ValueError(
+                        "dimension of the array for trace operation:"
+                        " should be >=2"
+                    )
+            else:
+                out_shape = tr_shape + (diag_size,)
+
+            if out is not None:
+                if out.shape != out_shape:
+                    raise ValueError("output array has wrong shape")
+                a = a._maybe_convert(out.dtype, (a, out))
+            else:
+                out = ndarray(
+                    shape=out_shape, dtype=self.dtype, inputs=(self,)
+                )
+            if out is None and dtype is not None:
+                a = a._maybe_convert(dtype, (a,))
 
             out._thunk._diag_helper(
-                a._thunk,
-                offset=offset,
-                naxes=N,
-                extract=extract,
+                a._thunk, offset=offset, naxes=N, extract=extract, trace=trace
             )
         return out
 
@@ -2053,6 +2087,47 @@ def diagonal(
                 raise ValueError("Either axis1/axis2 or axes must be supplied")
         return self._diag_helper(offset=offset, axes=axes, extract=extract)
 
+    @add_boilerplate()
+    def trace(self, offset=0, axis1=None, axis2=None, dtype=None, out=None):
+        """a.trace(offset=0, axis1=None, axis2=None, dtype = None, out = None)
+
+        Return the sum along diagonals of the array.
+
+        Refer to :func:`cunumeric.trace` for full documentation.
+
+        See Also
+        --------
+        cunumeric.trace : equivalent function
+
+        Availability
+        --------
+        Multiple GPUs, Multiple CPUs
+
+        """
+        if self.ndim < 2:
+            raise ValueError(
+                "trace operation can't be called on a array with DIM<2"
+            )
+
+        axes = []
+        if (axis1 is None) and (axis2 is None):
+            # default values for axis
+            axes = (0, 1)
+        elif (axis1 is None) or (axis2 is None):
+            raise ValueError("both axes should be passed")
+        else:
+            axes = (axis1, axis2)
+
+        res = self._diag_helper(
+            offset=offset, axes=axes, trace=True, out=out, dtype=dtype
+        )
+
+        # for 2D arrays we must return scalar
+        if self.ndim == 2:
+            res = res[0]
+
+        return res
+
     @add_boilerplate("rhs")
     def dot(self, rhs, out=None):
         """a.dot(rhs, out=None)

cunumeric/deferred.py

@@ -1357,6 +1357,7 @@ def _diag_helper(
         offset,
         naxes,
         extract,
+        trace,
     ):
         # fill output array with 0
         self.fill(np.array(0, dtype=self.dtype))
@@ -1370,10 +1371,18 @@ def _diag_helper(
                 # get slice of the original array by the offset
                 if offset > 0:
                     matrix = matrix.slice(start + 1, slice(offset, None))
-                if matrix.shape[n - 1] < matrix.shape[n]:
-                    diag = diag.promote(start + 1, matrix.shape[ndim - 1])
+                if trace:
+                    if matrix.ndim == 2:
+                        diag = diag.promote(0, matrix.shape[0])
+                        diag = diag.project(1, 0).promote(1, matrix.shape[1])
+                    else:
+                        for i in range(0, naxes):
+                            diag = diag.promote(start, matrix.shape[-i - 1])
                 else:
-                    diag = diag.promote(start, matrix.shape[ndim - 2])
+                    if matrix.shape[n - 1] < matrix.shape[n]:
+                        diag = diag.promote(start + 1, matrix.shape[ndim - 1])
+                    else:
+                        diag = diag.promote(start, matrix.shape[ndim - 2])
             else:
                 # promote output to the shape of the input  array
                 for i in range(1, naxes):

cunumeric/eager.py

@@ -533,29 +533,23 @@ def choose(self, *args, rhs):
             choices = tuple(c.array for c in args)
             self.array[:] = np.choose(rhs.array, choices, mode="raise")
 
-    def _diag_helper(
-        self,
-        rhs,
-        offset,
-        naxes,
-        extract,
-    ):
+    def _diag_helper(self, rhs, offset, naxes, extract, trace):
         self.check_eager_args(rhs)
         if self.deferred is not None:
-            self.deferred._diag_helper(
-                rhs,
-                offset,
-                naxes,
-                extract,
-            )
+            self.deferred._diag_helper(rhs, offset, naxes, extract, trace)
         else:
-            if (naxes == 2) and extract:
+            if (naxes == 2) and extract and not trace:
                 ndims = rhs.array.ndim
                 self.array[:] = np.diagonal(
                     rhs.array, offset=offset, axis1=ndims - 2, axis2=ndims - 1
                 )
             elif (naxes < 2) and not extract:
                 self.array[:] = np.diag(rhs.array, offset)
+            elif (naxes >= 2) and trace:
+                ndim = rhs.array.ndim
+                self.array[:] = np.trace(
+                    rhs.array, offset=offset, axis1=ndim - 2, axis2=ndim - 1
+                )
             else:  # naxes>2
                 ndims = rhs.array.ndim
                 axes = tuple(range(ndims - naxes, ndims))

cunumeric/module.py

@@ -3009,6 +3009,64 @@ def einsum(expr, *operands, out=None, optimize=False):
     return operands[0]
 
 
+@add_boilerplate("a")
+def trace(a, offset=0, axis1=None, axis2=None, dtype=None, out=None):
+    """
+    Return the sum along diagonals of the array.
+
+    If a is 2-D, the sum along its diagonal with the given offset is
+    returned, i.e., the sum of elements a[i,i+offset] for all i.
+    If a has more than two dimensions, then the axes specified by axis1
+    and axis2 are used to determine the 2-D sub-arrays whose traces
+    are returned. The shape of the resulting array is the same as that
+    of a with axis1 and axis2 removed.
+
+    Parameters
+    ----------
+    a : array_like
+        Input array, from which the diagonals are taken.
+    offset : int, optional
+        Offset of the diagonal from the main diagonal. Can be both
+        positive and negative. Defaults to 0.
+    axis1, axis2 : int, optional
+        Axes to be used as the first and second axis of the 2-D sub-arrays
+        from which the diagonals should be taken. Defaults are the
+        first two axes of a.
+    dtype : dtype, optional
+        Determines the data-type of the returned array and of the
+        accumulator where the elements are summed. If dtype has the value
+        None and a is of integer type of precision less than the default
+        integer precision, then the default integer precision is used.
+        Otherwise, the precision is the same as that of a.
+
+    out : ndarray, optional
+        Array into which the output is placed. Its type is preserved and
+        it must be of the right shape to hold the output.
+
+    Returns
+    -------
+    sum_along_diagonals : ndarray
+        If a is 2-D, the sum along the diagonal is returned. If a has
+        larger dimensions, then an array of sums along diagonals is returned.
+
+    Raises
+    ------
+    ValueError
+        If the dimension of `a` is less than 2.
+
+    See Also
+    --------
+    numpy.diagonal
+
+    Availability
+    --------
+    Multiple GPUs, Multiple CPUs
+    """
+    return a.trace(
+        offset=offset, axis1=axis1, axis2=axis2, dtype=dtype, out=out
+    )
+
+
 #################
 # Logic functions
 #################

cunumeric/thunk.py

@@ -143,13 +143,7 @@ def choose(self, *args, rhs):
         ...
 
     @abstractmethod
-    def _diag_helper(
-        self,
-        rhs,
-        offset,
-        naxes,
-        extract,
-    ):
+    def _diag_helper(self, rhs, offset, naxes, extract, trace):
         ...
 
     @abstractmethod

tests/integration/test_partition.py

@@ -159,21 +159,10 @@ def generate_random(shape, datatype):
 
 
 CASES = [
-    ((2, 5, 7), np.uint8),
-    ((8, 5), np.uint16),
-    ((22, 5, 7), np.uint32),
     ((220,), np.uint32),
     ((2, 5, 7), np.int8),
-    ((8, 5), np.int16),
-    ((22, 5, 7), np.int32),
-    ((2, 5, 7), np.int64),
     ((8, 5), np.float32),
-    ((8, 5), np.float64),
-    ((22, 5, 7), np.double),
-    ((220,), np.double),
-    ((2, 5, 7), np.complex64),
     ((2, 5, 7), np.complex128),
-    ((220,), np.complex128),
 ]
 
 

tests/trace.py

@@ -0,0 +1,87 @@
+# 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 random
+from itertools import permutations
+
+import numpy as np
+from test_tools.generators import mk_seq_array
+
+import cunumeric as num
+from legate.core import LEGATE_MAX_DIM
+
+
+def test():
+    # --------------------------------------------------------------
+    # TRACE
+    # --------------------------------------------------------------
+
+    a = np.arange(8).reshape((2, 4))
+    a_num = num.array(a)
+    res = np.trace(a)
+    res_num = num.trace(a_num)
+    assert np.array_equal(res, res_num)
+
+    res = np.trace(a, dtype=float)
+    res_num = num.trace(a_num, dtype=float)
+    assert np.array_equal(res, res_num)
+
+    a = np.arange(8).reshape((2, 2, 2))
+    a_num = num.array(a)
+    res = np.trace(a)
+    res_num = num.trace(a_num)
+    assert np.array_equal(res, res_num)
+
+    res = np.trace(a, offset=1)
+    res_num = num.trace(a_num, offset=1)
+    assert np.array_equal(res, res_num)
+
+    res = np.trace(a, offset=1, axis1=1, axis2=2)
+    res_num = num.trace(a_num, offset=1, axis1=1, axis2=2)
+    assert np.array_equal(res, res_num)
+
+    out = np.array([1, 2], dtype=float)
+    out_num = num.array(out)
+    np.trace(a, out=out)
+    num.trace(a_num, out=out_num)
+    assert np.array_equal(out, out_num)
+
+    np.trace(a, dtype=int, out=out)
+    num.trace(a_num, dtype=int, out=out_num)
+    assert np.array_equal(out, out_num)
+
+    a = np.arange(24).reshape((2, 2, 2, 3))
+    a_num = num.array(a)
+    res = np.trace(a)
+    res_num = num.trace(a_num)
+    assert np.array_equal(res, res_num)
+
+    for ndim in range(2, LEGATE_MAX_DIM + 1):
+        a_shape = tuple(random.randint(1, 9) for i in range(ndim))
+        np_array = mk_seq_array(np, a_shape)
+        num_array = mk_seq_array(num, a_shape)
+
+        # test trace
+        for axes in permutations(range(ndim), 2):
+            diag_size = min(a_shape[axes[0]], a_shape[axes[1]]) - 1
+            for offset in range(-diag_size + 1, diag_size):
+                assert np.array_equal(
+                    np_array.trace(offset, axes[0], axes[1]),
+                    num_array.trace(offset, axes[0], axes[1]),
+                )
+
+
+if __name__ == "__main__":
+    test()