WIP: adding logic for DIAGONAL

cunumeric/array.py

@@ -1139,12 +1139,117 @@ def cumsum(self, axis=None, dtype=None, out=None):
         )
         return self.convert_to_cunumeric_ndarray(numpy_array, stacklevel=3)
 
-    @unimplemented
-    def diagonal(self, offset=0, axis1=0, axis2=1):
-        numpy_array = self.__array__(stacklevel=3).diagonal(
-            offset=offset, axis1=axis1, axis2=axis2
-        )
-        return self.convert_to_cunumeric_ndarray(numpy_array, stacklevel=3)
+    # diagonal helper. Will return diagonal for arbitrary number of axes;
+    # 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):
+        # diag_helper can be used only for arrays with dim>=1
+        if self.ndim < 1:
+            raise ValueError("diag_helper is implemented for dim>=1")
+        elif self.ndim == 1:
+            if axes is not None:
+                raise ValueError(
+                    "Axes shouldn't be specified when getting "
+                    "diagonal for 1D array"
+                )
+            m = self.shape[0] + np.abs(offset)
+            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,
+                stacklevel=3,
+            )
+        else:
+            N = len(axes)
+            if len(axes) != len(set(axes)):
+                raise ValueError(
+                    "axes passed to diag_helper should be all different"
+                )
+            if self.ndim < N:
+                raise ValueError(
+                    "Dimension of input array shouldn't be less "
+                    "than number of axes"
+                )
+            # pack the axes that are not going to change
+            transpose_axes = tuple(
+                ax for ax in range(self.ndim) if ax not in axes
+            )
+            # only 2 axes provided, we transpose based on the offset
+            if N == 2:
+                if offset >= 0:
+                    a = self.transpose(transpose_axes + (axes[0], axes[1]))
+                else:
+                    a = self.transpose(transpose_axes + (axes[1], axes[0]))
+                    offset = -offset
+
+                if offset >= a.shape[self.ndim - 1]:
+                    raise ValueError(
+                        "'offset' for diag or diagonal must be in range"
+                    )
+
+                diag_size = max(0, min(a.shape[-2], a.shape[-1] - offset))
+            # more than 2 axes provided:
+            elif N > 2:
+                # offsets are supported only when naxes=2
+                if offset != 0:
+                    raise ValueError(
+                        "offset supported for number of axes == 2"
+                    )
+                # sort axes along which diagonal is calculated by size
+                axes = sorted(axes, reverse=True, key=lambda i: self.shape[i])
+                axes = tuple(axes)
+                # transpose a so axes for which diagonal is calculated are at
+                #  at the end
+                a = self.transpose(transpose_axes + axes)
+                diag_size = a.shape[a.ndim - 1]
+            elif N < 2:
+                raise ValueError(
+                    "number of axes passed to the diag_helper"
+                    " should be more than 1"
+                )
+
+            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, stacklevel=3, inputs=(self)
+            )
+
+            out._thunk.diag_helper(
+                a._thunk,
+                offset=offset,
+                naxes=N,
+                extract=extract,
+                stacklevel=3,
+            )
+        return out
+
+    def diagonal(
+        self, offset=0, axis1=None, axis2=None, extract=True, axes=None
+    ):
+        if self.ndim == 1:
+            if extract is True:
+                raise ValueError("extract can be true only for Ndim >=2")
+            axes = None
+        else:
+            if type(axis1) == int and type(axis2) == int:
+                if axes is not None:
+                    raise ValueError(
+                        "Either axis1/axis2 or axes must be supplied"
+                    )
+                axes = (axis1, axis2)
+            # default values for axes
+            elif (axis1 is None) and (axis2 is None) and (axes is None):
+                axes = (0, 1)
+            elif (axes is not None) and (
+                (axis1 is not None) or (axis2 is not None)
+            ):
+                raise ValueError("Either axis1/axis2 or axes must be supplied")
+        return self.diag_helper(offset=offset, axes=axes, extract=extract)
 
     def dot(self, rhs, out=None, stacklevel=1):
         rhs_array = self.convert_to_cunumeric_ndarray(

cunumeric/deferred.py

@@ -1266,63 +1266,70 @@ def choose(
     # Create or extract a diagonal from a matrix
     @profile
     @auto_convert([1])
-    @shadow_debug("diag", [1])
-    def diag(self, rhs, extract, k, stacklevel=0, callsite=None):
-        if extract:
-            matrix_array = rhs
-            diag_array = self
-        else:
-            matrix_array = self
-            diag_array = rhs
-
-        assert diag_array.ndim == 1
-        assert matrix_array.ndim == 2
-        assert diag_array.shape[0] <= min(
-            matrix_array.shape[0], matrix_array.shape[1]
+    @shadow_debug("diag_helper", [1])
+    def diag_helper(
+        self,
+        rhs,
+        offset,
+        naxes,
+        extract,
+        stacklevel=0,
+        callsite=None,
+    ):
+        # fill output array with 0
+        self.fill(
+            np.array(0, dtype=self.dtype),
+            stacklevel=(stacklevel + 1),
+            callsite=callsite,
         )
-        assert rhs.dtype == self.dtype
-
-        # Issue a fill operation to get the output initialized
         if extract:
-            diag_array.fill(
-                np.array(0, dtype=diag_array.dtype),
-                stacklevel=(stacklevel + 1),
-                callsite=callsite,
-            )
-        else:
-            matrix_array.fill(
-                np.array(0, dtype=matrix_array.dtype),
-                stacklevel=(stacklevel + 1),
-                callsite=callsite,
-            )
-
-        matrix = matrix_array.base
-        diag = diag_array.base
-
-        if k > 0:
-            matrix = matrix.slice(1, slice(k, None))
-        elif k < 0:
-            matrix = matrix.slice(0, slice(-k, None))
-
-        if matrix.shape[0] < matrix.shape[1]:
-            diag = diag.promote(1, matrix.shape[1])
+            diag = self.base
+            matrix = rhs.base
+            ndim = rhs.ndim
+            start = matrix.ndim - naxes
+            n = ndim - 1
+            if naxes == 2:
+                # 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])
+                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):
+                    diag = diag.promote(start, matrix.shape[-i - 1])
         else:
-            diag = diag.promote(0, matrix.shape[0])
+            matrix = self.base
+            diag = rhs.base
+            ndim = self.ndim
+            # get slice of the original array by the offset
+            if offset > 0:
+                matrix = matrix.slice(1, slice(offset, None))
+            elif offset < 0:
+                matrix = matrix.slice(0, slice(-offset, None))
+
+            if matrix.shape[0] < matrix.shape[1]:
+                diag = diag.promote(1, matrix.shape[1])
+            else:
+                diag = diag.promote(0, matrix.shape[0])
 
         task = self.context.create_task(CuNumericOpCode.DIAG)
 
         if extract:
             task.add_reduction(diag, ReductionOp.ADD)
             task.add_input(matrix)
+            task.add_alignment(matrix, diag)
         else:
             task.add_output(matrix)
-            task.add_input(matrix)
             task.add_input(diag)
+            task.add_input(matrix)
+            task.add_alignment(diag, matrix)
 
+        task.add_scalar_arg(naxes, ty.int32)
         task.add_scalar_arg(extract, bool)
 
-        task.add_alignment(matrix, diag)
-
         task.execute()
 
     # Create an identity array with the ones offset from the diagonal by k

cunumeric/eager.py

@@ -21,6 +21,31 @@
 from .thunk import NumPyThunk
 
 
+def eye_reference(shape, dtype, axes):
+    n = min(shape[ax] for ax in axes)
+    res = np.zeros(shape, dtype=dtype)
+    for i in range(n):
+        sl = tuple(
+            i if ax in axes else slice(None) for ax in range(len(shape))
+        )
+        res[sl] = 1
+    return res
+
+
+def diagonal_reference(a, axes):
+    transpose_axes = tuple(ax for ax in range(a.ndim) if ax not in axes)
+    axes = sorted(axes, reverse=False, key=lambda i: a.shape[i])
+    axes = tuple(axes)
+    a = a.transpose(transpose_axes + axes)
+    diff = a.ndim - len(axes)
+    axes = tuple((diff + ax) for ax in range(0, len(axes)))
+    eye = eye_reference(a.shape, a.dtype, axes)
+    res = a * eye
+    for ax in list(reversed(sorted(axes)))[:-1]:
+        res = res.sum(axis=ax)
+    return res
+
+
 class EagerArray(NumPyThunk):
     """This is an eager thunk for describing NumPy computations.
     It is backed by a standard NumPy array that stores the result
@@ -420,15 +445,39 @@ def choose(
             self.array[:] = np.choose(rhs.array, choices, mode="raise")
             self.runtime.profile_callsite(stacklevel + 1, False)
 
-    def diag(self, rhs, extract, k, stacklevel):
+    def diag_helper(
+        self,
+        rhs,
+        offset,
+        naxes,
+        extract,
+        stacklevel=0,
+        callsite=None,
+    ):
         if self.shadow:
             rhs = self.runtime.to_eager_array(rhs, stacklevel=(stacklevel + 1))
         elif self.deferred is None:
             self.check_eager_args((stacklevel + 1), rhs)
         if self.deferred is not None:
-            self.deferred.diag(rhs, extract, k, stacklevel=(stacklevel + 1))
+            self.deferred.diag_helper(
+                rhs,
+                offset,
+                naxes,
+                extract,
+                stacklevel=(stacklevel + 1),
+            )
         else:
-            self.array[:] = np.diag(rhs.array, k)
+            if (naxes == 2) and extract:
+                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)
+            else:  # naxes>2
+                ndims = rhs.array.ndim
+                axes = tuple(range(ndims - naxes, ndims))
+                self.array = diagonal_reference(rhs.array, axes)
             self.runtime.profile_callsite(stacklevel + 1, False)
 
     def eye(self, k, stacklevel):

cunumeric/lazy.py

@@ -101,8 +101,16 @@ def choose(
     ):
         raise NotImplementedError("Implement in derived classes")
 
-    def diag(self, rhs, extract, k, stacklevel):
-        """Fill in or extract a diagonal from a matrix
+    def diag_helper(
+        self,
+        rhs,
+        offset,
+        naxes,
+        extract,
+        stacklevel=0,
+        callsite=None,
+    ):
+        """Fill in or extract a diagonal from array
 
         :meta private:
         """