indexing ops: diag indices (#339)

* add indexing tests

Author: Alexander-Makaryev

dpnp/backend_indexing.pyx

@@ -39,10 +39,22 @@ from dpnp.dpnp_iface_counting import count_nonzero
 
 
 __all__ += [
+    "dpnp_diag_indices",
     "dpnp_nonzero",
 ]
 
 
+cpdef tuple dpnp_diag_indices(n, ndim):
+    cdef dparray res_item = dpnp.arange(n, dtype=dpnp.int64)
+
+    # yes, all are the same item
+    result = []
+    for i in range(ndim):
+        result.append(res_item)
+
+    return tuple(result)
+
+
 cpdef tuple dpnp_nonzero(dparray in_array1):
     res_count = in_array1.ndim
 

dpnp/dpnp_iface_indexing.py

@@ -49,10 +49,106 @@
 
 
 __all__ = [
+    "diag_indices",
+    "diag_indices_from",
     "nonzero",
 ]
 
 
+def diag_indices(n, ndim=2):
+    """
+    Return the indices to access the main diagonal of an array.
+
+    This returns a tuple of indices that can be used to access the main
+    diagonal of an array `a` with ``a.ndim >= 2`` dimensions and shape
+    (n, n, ..., n). For ``a.ndim = 2`` this is the usual diagonal, for
+    ``a.ndim > 2`` this is the set of indices to access ``a[i, i, ..., i]``
+    for ``i = [0..n-1]``.
+
+    For full documentation refer to :obj:`numpy.diag_indices`.
+
+    See also
+    --------
+    :obj:`diag_indices_from` : Return the indices to access the main
+                               diagonal of an n-dimensional array.
+
+    Examples
+    --------
+    Create a set of indices to access the diagonal of a (4, 4) array:
+
+    >>> import dpnp as np
+    >>> di = np.diag_indices(4)
+    >>> di
+    (array([0, 1, 2, 3]), array([0, 1, 2, 3]))
+    >>> a = np.arange(16).reshape(4, 4)
+    >>> a
+    array([[ 0,  1,  2,  3],
+           [ 4,  5,  6,  7],
+           [ 8,  9, 10, 11],
+           [12, 13, 14, 15]])
+    >>> a[di] = 100
+    >>> a
+    array([[100,   1,   2,   3],
+           [  4, 100,   6,   7],
+           [  8,   9, 100,  11],
+           [ 12,  13,  14, 100]])
+
+    Now, we create indices to manipulate a 3-D array:
+
+    >>> d3 = np.diag_indices(2, 3)
+    >>> d3
+    (array([0, 1]), array([0, 1]), array([0, 1]))
+
+    And use it to set the diagonal of an array of zeros to 1:
+
+    >>> a = np.zeros((2, 2, 2), dtype=int)
+    >>> a[d3] = 1
+    >>> a
+    array([[[1, 0],
+            [0, 0]],
+           [[0, 0],
+            [0, 1]]])
+
+    """
+
+    if not use_origin_backend():
+        return dpnp_diag_indices(n, ndim)
+
+    return call_origin(numpy.diag_indices, n, ndim)
+
+
+def diag_indices_from(arr):
+    """
+    Return the indices to access the main diagonal of an n-dimensional array.
+
+    For full documentation refer to :obj:`numpy.diag_indices_from`.
+
+    See also
+    --------
+    :obj:`diag_indices` : Return the indices to access the main
+                          diagonal of an array.
+
+    """
+
+    is_a_dparray = isinstance(arr, dparray)
+
+    if (not use_origin_backend(arr) and is_a_dparray):
+        # original limitation
+        if not arr.ndim >= 2:
+            checker_throw_value_error("diag_indices_from", "arr.ndim", arr.ndim, "at least 2-d")
+
+        # original limitation
+        # For more than d=2, the strided formula is only valid for arrays with
+        # all dimensions equal, so we check first.
+        if not numpy.alltrue(numpy.diff(arr.shape) == 0):  # TODO: replace alltrue and diff funcs with dpnp own ones
+            checker_throw_value_error("diag_indices_from", "arr.shape", arr.shape,
+                                      "All dimensions of input must be of equal length")
+
+        return dpnp_diag_indices(arr.shape[0], arr.ndim)
+
+    return call_origin(numpy.diag_indices_from, arr)
+
+
 def nonzero(a):
     """
     Return the indices of the elements that are non-zero.