fixing small bugs + code clean-up for Diagonal

cunumeric/array.py

@@ -1144,26 +1144,34 @@ def cumsum(self, axis=None, dtype=None, out=None):
     # 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):
-        N = len(axes)
         # diag_helper can be used only for arrays with dim>=1
         if self.ndim < 1:
-            raise RuntimeError("diag_helper is implemented for dim>=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,
-                diag_size=diag_size,
                 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"
+                    "Dimension of input array shouldn't be less "
                     "than number of axes"
                 )
             # pack the axes that are not going to change
@@ -1187,14 +1195,17 @@ def diag_helper(self, offset=0, axes=None, extract=True):
             # more than 2 axes provided:
             elif N > 2:
                 # offsets are supported only when naxes=2
-                assert offset == 0
+                if offset != 0:
+                    raise ValueError(
+                        "offset supported for number of axes == 2"
+                    )
                 # sort axes along which diagonal is calculated by size
-                axes.sort(reverse=True, key=lambda i: a.shape[i])
+                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)
-                sizes = tuple(a.shape[i] for i in axes)
-                diag_size = max(0, min(sizes))
+                diag_size = a.shape[a.ndim - 1]
             elif N < 2:
                 raise ValueError(
                     "number of axes passed to the diag_helper"
@@ -1210,18 +1221,35 @@ def diag_helper(self, offset=0, axes=None, extract=True):
 
             out._thunk.diag_helper(
                 a._thunk,
-                diag_size=diag_size,
                 offset=offset,
                 naxes=N,
                 extract=extract,
                 stacklevel=3,
             )
         return out
 
-    def diagonal(self, offset=0, axis1=0, axis2=1, extract=True):
-        return self.diag_helper(
-            offset=offset, axes=(axis1, axis2), extract=extract
-        )
+    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

@@ -1270,7 +1270,6 @@ def choose(
     def diag_helper(
         self,
         rhs,
-        diag_size,
         offset,
         naxes,
         extract,
@@ -1287,7 +1286,7 @@ def diag_helper(
             diag = self.base
             matrix = rhs.base
             ndim = rhs.ndim
-            start = self.ndim - naxes + 1
+            start = matrix.ndim - naxes
             n = ndim - 1
             if naxes == 2:
                 # get slice of the original array by the offset
@@ -1305,7 +1304,6 @@ def diag_helper(
             matrix = self.base
             diag = rhs.base
             ndim = self.ndim
-            n = ndim - 1
             # get slice of the original array by the offset
             if offset > 0:
                 matrix = matrix.slice(1, slice(offset, None))
@@ -1329,51 +1327,11 @@ def diag_helper(
             task.add_input(matrix)
             task.add_alignment(diag, matrix)
 
-        task.add_scalar_arg(diag_size, ty.int32)
-        task.add_scalar_arg(offset, ty.int32)
         task.add_scalar_arg(naxes, ty.int32)
         task.add_scalar_arg(extract, bool)
 
         task.execute()
 
-    @profile
-    @auto_convert([1])
-    @shadow_debug("diagonal", [1])
-    def diagonal(
-        self,
-        rhs,
-        offset=0,
-        axis1=0,
-        axis2=1,
-        extract=True,
-        stacklevel=0,
-        callsite=None,
-    ):
-        axes = (axis1, axis2)
-        if rhs.ndim < 2:
-            raise RuntimeError("diagonal works for dim>=2")
-
-        tr = tuple(ax for ax in range(rhs.ndim) if ax not in axes)
-        #        for ax in range(self.ndim):
-        #            if ax not in axes:
-        #                tr = tr + (ax,)
-
-        if offset >= 0:
-            a = self.transpose(tr + (axes[0], axes[1]))
-        else:
-            a = self.transpose(tr + (axes[1], axes[0]))
-            offset = -offset
-        diag_size = max(0, min(a.shape[-2], a.shape[-1] - offset))
-
-        self.diag_helper(
-            rhs=rhs,
-            diag_size=diag_size,
-            offset=offset,
-            naxes=2,
-            extract=extract,
-            stacklevel=3,
-        )
-
     # Create an identity array with the ones offset from the diagonal by k
     @profile
     @shadow_debug("eye", [])

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,21 +445,9 @@ 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):
-        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))
-        else:
-            self.array[:] = np.diag(rhs.array, k)
-            self.runtime.profile_callsite(stacklevel + 1, False)
-
     def diag_helper(
         self,
         rhs,
-        diag_size,
         offset,
         naxes,
         extract,
@@ -448,34 +461,23 @@ def diag_helper(
         if self.deferred is not None:
             self.deferred.diag_helper(
                 rhs,
-                diag_size,
                 offset,
                 naxes,
                 extract,
                 stacklevel=(stacklevel + 1),
             )
         else:
-            # this should be called only if number of axes <=2
-            if naxes == 2:
+            if (naxes == 2) and extract:
                 ndims = rhs.array.ndim
                 self.array[:] = np.diagonal(
                     rhs.array, offset=offset, axis1=ndims - 2, axis2=ndims - 1
                 )
-            if naxes < 2:
+            elif (naxes < 2) and not extract:
                 self.array[:] = np.diag(rhs.array, offset)
-            self.runtime.profile_callsite(stacklevel + 1, False)
-
-    def diagonal(self, rhs, offset, axis1, axis2, extract, stacklevel):
-        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.diagonal(
-                rhs, offset, axis1, axis2, extract, stacklevel=(stacklevel + 1)
-            )
-        else:
-            self.array[:] = np.diagonal(rhs.array, offset, axis1, axis2)
+            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,14 +101,15 @@ def choose(
     ):
         raise NotImplementedError("Implement in derived classes")
 
-    def diag(self, rhs, extract, k, stacklevel):
-        """Fill in or extract a diagonal from a matrix
-
-        :meta private:
-        """
-        raise NotImplementedError("Implement in derived classes")
-
-    def diagonal(self, rhs, offset, axis1, axis2, extract, stacklevel):
+    def diag_helper(
+        self,
+        rhs,
+        offset,
+        naxes,
+        extract,
+        stacklevel=0,
+        callsite=None,
+    ):
         """Fill in or extract a diagonal from array
 
         :meta private:

cunumeric/module.py

@@ -208,8 +208,8 @@ def choose(a, choices, out=None, mode="raise"):
 @copy_docstring(np.diag)
 def diag(v, k=0):
     array = ndarray.convert_to_cunumeric_ndarray(v)
-    if array._thunk.scalar:
-        return array.copy()
+    if array.ndim == 0:
+        raise ValueError("Input must be 1- or 2-d")
     elif array.ndim == 1:
         return array.diagonal(offset=k, axis1=0, axis2=1, extract=False)
     elif array.ndim == 2:
@@ -218,11 +218,31 @@ def diag(v, k=0):
         raise ValueError("diag requires 1- or 2-D array, use diagonal instead")
 
 
-@copy_docstring(np.diagonal)
-def diagonal(a, offset=0, axis1=0, axis2=1, extract=True):
+def diagonal(a, offset=0, axis1=None, axis2=None, extract=True, axes=None):
+    """
+    Return specified diagonals.
+
+    See description in numpy
+    ------------------------
+    https://numpy.org/doc/stable/reference/generated/numpy.diag.html
+    https://numpy.org/doc/stable/reference/generated/numpy.diagonal.html#numpy.diagonal
+    https://numpy.org/doc/stable/reference/generated/numpy.ndarray.diagonal.html
+
+
+    cuNumeric implementation differences:
+    ------------------------------------
+    - It never returns a view
+    - support 1D arrays (similar to `diag`)
+    - support extra arguments:
+         -- extract: used to create diagonal from 1D array vs extracting
+           the diagonal from"
+         -- axes: list of axes for diagonal ( size of the list should be in
+           between 2 and size of the array)
+
+    """
     array = ndarray.convert_to_cunumeric_ndarray(a)
     return array.diagonal(
-        offset=offset, axis1=axis1, axis2=axis2, extract=extract
+        offset=offset, axis1=axis1, axis2=axis2, extract=extract, axes=axes
     )
 
 
@@ -693,23 +713,19 @@ def _contract(expr, a, b=None, out=None, stacklevel=1):
 
     # Handle duplicate modes on inputs
     c_a_modes = Counter(a_modes)
-    for mode in c_a_modes:
-        count = c_a_modes[mode]
+    for (mode, count) in c_a_modes.items():
         if count > 1:
             axes = [i for (i, m) in enumerate(a_modes) if m == mode]
             a = a.diag_helper(axes=axes)
             # diagonal is stored on last axis
-            for _ in range(count - 1):
-                a_modes.remove(mode)
+            a_modes = [m for m in a_modes if m != mode] + [mode]
     c_b_modes = Counter(b_modes)
-    for mode in c_b_modes:
-        count = c_b_modes[mode]
+    for (mode, count) in c_b_modes.items():
         if count > 1:
             axes = [i for (i, m) in enumerate(b_modes) if m == mode]
             b = b.diag_helper(axes=axes)
             # diagonal is stored on last axis
-            for _ in range(count - 1):
-                b_modes.remove(mode)
+            b_modes = [m for m in b_modes if m != mode] + [mode]
     # Drop modes corresponding to singleton dimensions. This handles cases of
     # broadcasting.
     for dim in reversed(range(a.ndim)):

cunumeric/thunk.py

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