Update dpnp.diagonal()

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,164 @@ 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 unless `a` is a `matrix`, in which case
+        a 1-D array rather than a (2-D) `matrix` is returned in order to
+        maintain backward compatibility.
+
+        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.
+
+    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
+
+    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")
+
+    if axis1 < axis2:
+        min_axis, max_axis = axis1, axis2
+    else:
+        min_axis, max_axis = axis2, axis1
 
-    return call_origin(numpy.diagonal, x1, offset, axis1, axis2)
+    # get list of the order of axes removing the two target axes
+    axes_order = list(range(a_ndim))
+    del axes_order[max_axis]
+    del axes_order[min_axis]
+
+    # 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
+    elif -n < offset < 0:
+        out_shape = a_shape[:-2] + (min(n + offset, m),)
+        out_strides = a_straides[:-2] + (st_n + st_m,)
+        out_offset = a_element_offset - st_n * 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):

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", [-1, 0, 1])
+    @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())

tests/test_sycl_queue.py

@@ -245,7 +245,6 @@ def test_array_creation_follow_device_2d_array(func, args, kwargs, device):
     assert_sycl_queue_equal(y.sycl_queue, x.sycl_queue)
 
 
-@pytest.mark.skip("muted until the issue reported by SAT-5969 is resolved")
 @pytest.mark.parametrize(
     "func, args, kwargs",
     [
@@ -292,7 +291,6 @@ def test_array_creation_cross_device(func, args, kwargs, device_x, device_y):
     assert_sycl_queue_equal(y.sycl_queue, x.to_device(device_y).sycl_queue)
 
 
-@pytest.mark.skip("muted until the issue reported by SAT-5969 is resolved")
 @pytest.mark.parametrize(
     "func, args, kwargs",
     [
@@ -2101,3 +2099,21 @@ def test_histogram(weights, device):
     edges_queue = result_edges.sycl_queue
     assert_sycl_queue_equal(hist_queue, iv.sycl_queue)
     assert_sycl_queue_equal(edges_queue, iv.sycl_queue)
+
+
+@pytest.mark.parametrize(
+    "device",
+    valid_devices,
+    ids=[device.filter_string for device in valid_devices],
+)
+def test_diagonal(device):
+    a = numpy.arange(12).reshape(3, 2, 2)
+    a_dp = dpnp.array(a, device=device)
+
+    expected = numpy.diagonal(a)
+    result = dpnp.diagonal(a_dp)
+
+    assert_dtype_allclose(result, expected)
+
+    result_queue = result.sycl_queue
+    assert_sycl_queue_equal(result_queue, a_dp.sycl_queue)

tests/test_usm_type.py

@@ -1228,3 +1228,11 @@ def test_histogram(usm_type_v, usm_type_w):
     assert w.usm_type == usm_type_w
     assert hist.usm_type == du.get_coerced_usm_type([usm_type_v, usm_type_w])
     assert edges.usm_type == du.get_coerced_usm_type([usm_type_v, usm_type_w])
+
+
+@pytest.mark.parametrize("usm_type", list_of_usm_types, ids=list_of_usm_types)
+def test_diagonal(usm_type):
+    a = dp.arange(12, usm_type=usm_type).reshape(3, 2, 2)
+
+    res = dp.diagonal(a)
+    assert res.usm_type == usm_type

tests/third_party/cupy/core_tests/test_ndarray_copy_and_view.py

@@ -384,7 +384,6 @@ def test_astype_boolean_view(self, xp):
         return a.astype(numpy.int8)
 
 
-@pytest.mark.usefixtures("allow_fall_back_on_numpy")
 class TestArrayDiagonal:
     @testing.for_all_dtypes()
     @testing.numpy_cupy_array_equal()

tests/third_party/cupy/indexing_tests/test_indexing.py

@@ -99,7 +99,6 @@ def test_diagonal(self, xp, dtype):
         a = testing.shaped_arange((3, 4, 5), xp, dtype)
         return a.diagonal(1, 2, 0)
 
-    @pytest.mark.usefixtures("allow_fall_back_on_numpy")
     @testing.for_all_dtypes()
     @testing.numpy_cupy_array_equal()
     def test_external_diagonal(self, xp, dtype):