Update dpnp.diagonal() (#1817)

* Update dpnp.diagonal to support offset and axes

* Update cupy tests for dpnp.diagonal()

* Update dpnp tests for dpnp.diagonal()

* Remove pytest.skip for a closed issue

* Extend test cases in test_diagonal_errors

* Update dpnp.ndarray.diagonal

* Update docstrings for dpnp.diagonal

* Address remarks

* Simplify getting axes_order

* Add cupy test_indexing.py to test scope

* Remove TODO from dpnp_det/dpnp_slogdet

* Improve test coverage

* Update the calculation of result parameters

Author: vlad-perevezentsev

.github/workflows/conda-package.yml

@@ -38,6 +38,7 @@ env:
       test_umath.py
       test_usm_type.py
       third_party/cupy/core_tests
+      third_party/cupy/indexing_tests/test_indexing.py
       third_party/cupy/linalg_tests/test_decomposition.py
       third_party/cupy/linalg_tests/test_norms.py
       third_party/cupy/linalg_tests/test_product.py

dpnp/dpnp_algo/dpnp_algo_indexing.pxi

@@ -38,7 +38,6 @@ and the rest of the library
 __all__ += [
     "dpnp_choose",
     "dpnp_diag_indices",
-    "dpnp_diagonal",
     "dpnp_fill_diagonal",
     "dpnp_putmask",
     "dpnp_select",

dpnp/dpnp_array.py

@@ -744,15 +744,26 @@ def cumsum(self, axis=None, dtype=None, out=None):
 
     # 'data',
 
-    def diagonal(input, offset=0, axis1=0, axis2=1):
+    def diagonal(self, offset=0, axis1=0, axis2=1):
         """
         Return specified diagonals.
 
         Refer to :obj:`dpnp.diagonal` for full documentation.
 
+        See Also
+        --------
+        :obj:`dpnp.diagonal` : Equivalent function.
+
+        Examples
+        --------
+        >>> import dpnp as np
+        >>> a = np.arange(4).reshape(2,2)
+        >>> a.diagonal()
+        array([0, 3])
+
         """
 
-        return dpnp.diagonal(input, offset, axis1, axis2)
+        return dpnp.diagonal(self, offset=offset, axis1=axis1, axis2=axis2)
 
     def dot(self, b, out=None):
         """

dpnp/dpnp_iface_indexing.py

@@ -47,7 +47,6 @@
 from .dpnp_algo import (
     dpnp_choose,
     dpnp_diag_indices,
-    dpnp_diagonal,
     dpnp_fill_diagonal,
     dpnp_putmask,
     dpnp_select,
@@ -276,34 +275,156 @@ def diag_indices_from(x1):
     return call_origin(numpy.diag_indices_from, x1)
 
 
-def diagonal(x1, offset=0, axis1=0, axis2=1):
+def diagonal(a, offset=0, axis1=0, axis2=1):
     """
     Return specified diagonals.
 
     For full documentation refer to :obj:`numpy.diagonal`.
 
-    Limitations
-    -----------
-    Input array is supported as :obj:`dpnp.ndarray`.
-    Parameters `axis1` and `axis2` are supported only with default values.
-    Otherwise the function will be executed sequentially on CPU.
+    Parameters
+    ----------
+    a : {dpnp.ndarray, usm_ndarray}
+        Array from which the diagonals are taken.
+    offset : int, optional
+        Offset of the diagonal from the main diagonal. Can be positive or
+        negative. Defaults to main diagonal (``0``).
+    axis1 : int, optional
+        Axis to be used as the first axis of the 2-D sub-arrays from which
+        the diagonals should be taken. Defaults to first axis (``0``).
+    axis2 : int, optional
+        Axis to be used as the second axis of the 2-D sub-arrays from
+        which the diagonals should be taken. Defaults to second axis (``1``).
+
+    Returns
+    -------
+    array_of_diagonals : dpnp.ndarray
+        If `a` is 2-D, then a 1-D array containing the diagonal and of the
+        same type as `a` is returned.
+        If ``a.ndim > 2``, then the dimensions specified by `axis1` and `axis2`
+        are removed, and a new axis inserted at the end corresponding to the
+        diagonal.
+
+    See Also
+    --------
+    :obj:`dpnp.diag` : Extract a diagonal or construct a diagonal array.
+    :obj:`dpnp.diagflat` : Create a two-dimensional array
+                           with the flattened input as a diagonal.
+    :obj:`dpnp.trace` : Return the sum along diagonals of the array.
+
+    Examples
+    --------
+    >>> import dpnp as np
+    >>> a = np.arange(4).reshape(2,2)
+    >>> a
+    array([[0, 1],
+           [2, 3]])
+    >>> a.diagonal()
+    array([0, 3])
+    >>> a.diagonal(1)
+    array([1])
+
+    A 3-D example:
+
+    >>> a = np.arange(8).reshape(2,2,2)
+    >>> a
+    array([[[0, 1],
+            [2, 3]],
+           [[4, 5],
+            [6, 7]]])
+    >>> a.diagonal(0,  # Main diagonals of two arrays created by skipping
+    ...            0,  # across the outer(left)-most axis last and
+    ...            1)  # the "middle" (row) axis first.
+    array([[0, 6],
+           [1, 7]])
+
+    The sub-arrays whose main diagonals we just obtained; note that each
+    corresponds to fixing the right-most (column) axis, and that the
+    diagonals are "packed" in rows.
+
+    >>> a[:,:,0]  # main diagonal is [0 6]
+    array([[0, 2],
+           [4, 6]])
+    >>> a[:,:,1]  # main diagonal is [1 7]
+    array([[1, 3],
+           [5, 7]])
+
+    The anti-diagonal can be obtained by reversing the order of elements
+    using either `dpnp.flipud` or `dpnp.fliplr`.
+
+    >>> a = np.arange(9).reshape(3, 3)
+    >>> a
+    array([[0, 1, 2],
+           [3, 4, 5],
+           [6, 7, 8]])
+    >>> np.fliplr(a).diagonal()  # Horizontal flip
+    array([2, 4, 6])
+    >>> np.flipud(a).diagonal()  # Vertical flip
+    array([6, 4, 2])
+
+    Note that the order in which the diagonal is retrieved varies depending
+    on the flip function.
 
     """
 
-    x1_desc = dpnp.get_dpnp_descriptor(x1, copy_when_nondefault_queue=False)
-    if x1_desc:
-        if not isinstance(offset, int):
-            pass
-        elif offset < 0:
-            pass
-        elif axis1 != 0:
-            pass
-        elif axis2 != 1:
-            pass
-        else:
-            return dpnp_diagonal(x1_desc, offset).get_pyobj()
+    dpnp.check_supported_arrays_type(a)
+    a_ndim = a.ndim
 
-    return call_origin(numpy.diagonal, x1, offset, axis1, axis2)
+    if a_ndim < 2:
+        raise ValueError("diag requires an array of at least two dimensions")
+
+    if not isinstance(offset, int):
+        raise TypeError(
+            f"`offset` must be an integer data type, but got {type(offset)}"
+        )
+
+    axis1 = normalize_axis_index(axis1, a_ndim)
+    axis2 = normalize_axis_index(axis2, a_ndim)
+
+    if axis1 == axis2:
+        raise ValueError("`axis1` and `axis2` cannot be the same")
+
+    # get list of the order of all axes excluding the two target axes
+    axes_order = [i for i in range(a_ndim) if i not in [axis1, axis2]]
+
+    # transpose the input array to put the target axes at the end
+    # to simplify diagonal extraction
+    if offset >= 0:
+        a = dpnp.transpose(a, axes_order + [axis1, axis2])
+    else:
+        a = dpnp.transpose(a, axes_order + [axis2, axis1])
+        offset = -offset
+
+    a_shape = a.shape
+    a_straides = a.strides
+    n, m = a_shape[-2:]
+    st_n, st_m = a_straides[-2:]
+    # pylint: disable=W0212
+    a_element_offset = a.get_array()._element_offset
+
+    # Compute shape, strides and offset of the resulting diagonal array
+    # based on the input offset
+    if offset == 0:
+        out_shape = a_shape[:-2] + (min(n, m),)
+        out_strides = a_straides[:-2] + (st_n + st_m,)
+        out_offset = a_element_offset
+    elif 0 < offset < m:
+        out_shape = a_shape[:-2] + (min(n, m - offset),)
+        out_strides = a_straides[:-2] + (st_n + st_m,)
+        out_offset = a_element_offset + st_m * offset
+    else:
+        out_shape = a_shape[:-2] + (0,)
+        out_strides = a_straides[:-2] + (1,)
+        out_offset = a_element_offset
+
+    return dpnp_array._create_from_usm_ndarray(
+        dpt.usm_ndarray(
+            out_shape,
+            dtype=a.dtype,
+            buffer=a.get_array(),
+            strides=out_strides,
+            offset=out_offset,
+        )
+    )
 
 
 def extract(condition, x):

dpnp/linalg/dpnp_utils_linalg.py

@@ -953,11 +953,7 @@ def dpnp_det(a):
 
     lu, ipiv, dev_info = _lu_factor(a, res_type)
 
-    # Transposing 'lu' to swap the last two axes for compatibility
-    # with 'dpnp.diagonal' as it does not support 'axis1' and 'axis2' arguments.
-    # TODO: Replace with 'dpnp.diagonal(lu, axis1=-2, axis2=-1)' when supported.
-    lu_transposed = lu.transpose(-2, -1, *range(lu.ndim - 2))
-    diag = dpnp.diagonal(lu_transposed)
+    diag = dpnp.diagonal(lu, axis1=-2, axis2=-1)
 
     det = dpnp.prod(dpnp.abs(diag), axis=-1)
 
@@ -2112,11 +2108,7 @@ def dpnp_slogdet(a):
 
     lu, ipiv, dev_info = _lu_factor(a, res_type)
 
-    # Transposing 'lu' to swap the last two axes for compatibility
-    # with 'dpnp.diagonal' as it does not support 'axis1' and 'axis2' arguments.
-    # TODO: Replace with 'dpnp.diagonal(lu, axis1=-2, axis2=-1)' when supported.
-    lu_transposed = lu.transpose(-2, -1, *range(lu.ndim - 2))
-    diag = dpnp.diagonal(lu_transposed)
+    diag = dpnp.diagonal(lu, axis1=-2, axis2=-1)
 
     logdet = dpnp.log(dpnp.abs(diag)).sum(axis=-1)
 

tests/skipped_tests.tbl

@@ -148,14 +148,6 @@ tests/third_party/cupy/indexing_tests/test_indexing.py::TestIndexing::test_compr
 tests/third_party/cupy/indexing_tests/test_indexing.py::TestIndexing::test_compress_empty_1dim_no_axis
 tests/third_party/cupy/indexing_tests/test_indexing.py::TestIndexing::test_compress_no_axis
 tests/third_party/cupy/indexing_tests/test_indexing.py::TestIndexing::test_compress_no_bool
-tests/third_party/cupy/indexing_tests/test_indexing.py::TestIndexing::test_diagonal
-tests/third_party/cupy/indexing_tests/test_indexing.py::TestIndexing::test_diagonal_invalid1
-tests/third_party/cupy/indexing_tests/test_indexing.py::TestIndexing::test_diagonal_invalid2
-tests/third_party/cupy/indexing_tests/test_indexing.py::TestIndexing::test_diagonal_negative1
-tests/third_party/cupy/indexing_tests/test_indexing.py::TestIndexing::test_diagonal_negative2
-tests/third_party/cupy/indexing_tests/test_indexing.py::TestIndexing::test_diagonal_negative3
-tests/third_party/cupy/indexing_tests/test_indexing.py::TestIndexing::test_diagonal_negative4
-tests/third_party/cupy/indexing_tests/test_indexing.py::TestIndexing::test_diagonal_negative5
 tests/third_party/cupy/indexing_tests/test_indexing.py::TestSelect::test_select
 tests/third_party/cupy/indexing_tests/test_indexing.py::TestSelect::test_select_1D_choicelist
 tests/third_party/cupy/indexing_tests/test_indexing.py::TestSelect::test_select_choicelist_condlist_broadcast

tests/skipped_tests_gpu.tbl

@@ -210,14 +210,6 @@ tests/third_party/cupy/indexing_tests/test_indexing.py::TestSelect::test_select_
 tests/third_party/cupy/indexing_tests/test_indexing.py::TestSelect::test_select_odd_shaped_broadcastable_complex
 tests/third_party/cupy/indexing_tests/test_indexing.py::TestSelect::test_select_odd_shaped_non_broadcastable
 tests/third_party/cupy/indexing_tests/test_indexing.py::TestSelect::test_select_type_error_condlist
-tests/third_party/cupy/indexing_tests/test_indexing.py::TestIndexing::test_diagonal
-tests/third_party/cupy/indexing_tests/test_indexing.py::TestIndexing::test_diagonal_invalid1
-tests/third_party/cupy/indexing_tests/test_indexing.py::TestIndexing::test_diagonal_invalid2
-tests/third_party/cupy/indexing_tests/test_indexing.py::TestIndexing::test_diagonal_negative1
-tests/third_party/cupy/indexing_tests/test_indexing.py::TestIndexing::test_diagonal_negative2
-tests/third_party/cupy/indexing_tests/test_indexing.py::TestIndexing::test_diagonal_negative3
-tests/third_party/cupy/indexing_tests/test_indexing.py::TestIndexing::test_diagonal_negative4
-tests/third_party/cupy/indexing_tests/test_indexing.py::TestIndexing::test_diagonal_negative5
 
 tests/third_party/cupy/indexing_tests/test_insert.py::TestDiagIndices_param_0_{n=2, ndim=2}::test_diag_indices
 tests/third_party/cupy/indexing_tests/test_insert.py::TestDiagIndices_param_1_{n=2, ndim=3}::test_diag_indices

tests/test_indexing.py

@@ -322,48 +322,68 @@ def test_choose():
     assert_array_equal(expected, result)
 
 
-@pytest.mark.parametrize("arr_dtype", get_all_dtypes(no_bool=True))
-@pytest.mark.parametrize("offset", [0, 1], ids=["0", "1"])
-@pytest.mark.parametrize(
-    "array",
-    [
-        [[0, 0], [0, 0]],
-        [[1, 2], [1, 2]],
-        [[1, 2], [3, 4]],
-        [[0, 1, 2], [3, 4, 5], [6, 7, 8]],
-        [[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]],
-        [[[1, 2], [3, 4]], [[1, 2], [2, 1]], [[1, 3], [3, 1]]],
-        [
-            [[[1, 2], [3, 4]], [[1, 2], [2, 1]]],
-            [[[1, 3], [3, 1]], [[0, 1], [1, 3]]],
-        ],
-        [
-            [[[1, 2, 3], [3, 4, 5]], [[1, 2, 3], [2, 1, 0]]],
-            [[[1, 3, 5], [3, 1, 0]], [[0, 1, 2], [1, 3, 4]]],
+class TestDiagonal:
+    @pytest.mark.parametrize("dtype", get_all_dtypes(no_bool=True))
+    @pytest.mark.parametrize("offset", [-3, -1, 0, 1, 3])
+    @pytest.mark.parametrize(
+        "shape",
+        [(2, 2), (3, 3), (2, 5), (3, 2, 2), (2, 2, 2, 2), (2, 2, 2, 3)],
+        ids=[
+            "(2,2)",
+            "(3,3)",
+            "(2,5)",
+            "(3,2,2)",
+            "(2,2,2,2)",
+            "(2,2,2,3)",
         ],
+    )
+    def test_diagonal_offset(self, shape, dtype, offset):
+        a = numpy.arange(numpy.prod(shape), dtype=dtype).reshape(shape)
+        a_dp = dpnp.array(a)
+        expected = numpy.diagonal(a, offset)
+        result = dpnp.diagonal(a_dp, offset)
+        assert_array_equal(expected, result)
+
+    @pytest.mark.parametrize("dtype", get_all_dtypes(no_bool=True))
+    @pytest.mark.parametrize(
+        "shape, axis_pairs",
         [
-            [[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]],
-            [[[13, 14, 15], [16, 17, 18]], [[19, 20, 21], [22, 23, 24]]],
+            ((3, 4), [(0, 1), (1, 0)]),
+            ((3, 4, 5), [(0, 1), (1, 2), (0, 2)]),
+            ((4, 3, 5, 2), [(0, 1), (1, 2), (2, 3), (0, 3)]),
         ],
-    ],
-    ids=[
-        "[[0, 0], [0, 0]]",
-        "[[1, 2], [1, 2]]",
-        "[[1, 2], [3, 4]]",
-        "[[0, 1, 2], [3, 4, 5], [6, 7, 8]]",
-        "[[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]",
-        "[[[1, 2], [3, 4]], [[1, 2], [2, 1]], [[1, 3], [3, 1]]]",
-        "[[[[1, 2], [3, 4]], [[1, 2], [2, 1]]], [[[1, 3], [3, 1]], [[0, 1], [1, 3]]]]",
-        "[[[[1, 2, 3], [3, 4, 5]], [[1, 2, 3], [2, 1, 0]]], [[[1, 3, 5], [3, 1, 0]], [[0, 1, 2], [1, 3, 4]]]]",
-        "[[[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]], [[[13, 14, 15], [16, 17, 18]], [[19, 20, 21], [22, 23, 24]]]]",
-    ],
-)
-def test_diagonal(array, arr_dtype, offset):
-    a = numpy.array(array, dtype=arr_dtype)
-    ia = dpnp.array(a)
-    expected = numpy.diagonal(a, offset)
-    result = dpnp.diagonal(ia, offset)
-    assert_array_equal(expected, result)
+    )
+    def test_diagonal_axes(self, shape, axis_pairs, dtype):
+        a = numpy.arange(numpy.prod(shape), dtype=dtype).reshape(shape)
+        a_dp = dpnp.array(a)
+        for axis1, axis2 in axis_pairs:
+            expected = numpy.diagonal(a, axis1=axis1, axis2=axis2)
+            result = dpnp.diagonal(a_dp, axis1=axis1, axis2=axis2)
+            assert_array_equal(expected, result)
+
+    def test_diagonal_errors(self):
+        a = dpnp.arange(12).reshape(3, 4)
+
+        # unsupported type
+        a_np = dpnp.asnumpy(a)
+        assert_raises(TypeError, dpnp.diagonal, a_np)
+
+        # a.ndim < 2
+        a_ndim_1 = a.flatten()
+        assert_raises(ValueError, dpnp.diagonal, a_ndim_1)
+
+        # unsupported type `offset`
+        assert_raises(TypeError, dpnp.diagonal, a, offset=1.0)
+        assert_raises(TypeError, dpnp.diagonal, a, offset=[0])
+
+        # axes are out of bounds
+        assert_raises(numpy.AxisError, a.diagonal, axis1=0, axis2=5)
+        assert_raises(numpy.AxisError, a.diagonal, axis1=5, axis2=0)
+        assert_raises(numpy.AxisError, a.diagonal, axis1=5, axis2=5)
+
+        # same axes
+        assert_raises(ValueError, a.diagonal, axis1=1, axis2=1)
+        assert_raises(ValueError, a.diagonal, axis1=1, axis2=-1)
 
 
 @pytest.mark.parametrize("arr_dtype", get_all_dtypes())