add support for axes as list in dpnp.ndarray.transpose (#1770)

* support axes as list

* fix a bug

* update description

* fix docstring when bullet is used

Author: vtavana

dpnp/dpnp_array.py

@@ -1317,21 +1317,31 @@ def transpose(self, *axes):
 
         For full documentation refer to :obj:`numpy.ndarray.transpose`.
 
+        Parameters
+        ----------
+        axes : None, tuple or list of ints, n ints, optional
+            * ``None`` or no argument: reverses the order of the axes.
+            * tuple or list of ints: `i` in the `j`-th place in the tuple/list
+              means that the array’s `i`-th axis becomes the transposed
+              array’s `j`-th axis.
+            * n ints: same as an n-tuple/n-list of the same ints (this form is
+              intended simply as a “convenience” alternative to the tuple form).
+
         Returns
         -------
-        y : dpnp.ndarray
+        out : dpnp.ndarray
             View of the array with its axes suitably permuted.
 
         See Also
         --------
-            :obj:`dpnp.transpose` : Equivalent function.
-            :obj:`dpnp.ndarray.ndarray.T` : Array property returning the array transposed.
-            :obj:`dpnp.ndarray.reshape` : Give a new shape to an array without changing its data.
+        :obj:`dpnp.transpose` : Equivalent function.
+        :obj:`dpnp.ndarray.ndarray.T` : Array property returning the array transposed.
+        :obj:`dpnp.ndarray.reshape` : Give a new shape to an array without changing its data.
 
         Examples
         --------
-        >>> import dpnp as dp
-        >>> a = dp.array([[1, 2], [3, 4]])
+        >>> import dpnp as np
+        >>> a = np.array([[1, 2], [3, 4]])
         >>> a
         array([[1, 2],
                [3, 4]])
@@ -1342,7 +1352,7 @@ def transpose(self, *axes):
         array([[1, 3],
                [2, 4]])
 
-        >>> a = dp.array([1, 2, 3, 4])
+        >>> a = np.array([1, 2, 3, 4])
         >>> a
         array([1, 2, 3, 4])
         >>> a.transpose()
@@ -1355,7 +1365,7 @@ def transpose(self, *axes):
             return self
 
         axes_len = len(axes)
-        if axes_len == 1 and isinstance(axes[0], tuple):
+        if axes_len == 1 and isinstance(axes[0], (tuple, list)):
             axes = axes[0]
 
         res = self.__new__(dpnp_array)

dpnp/dpnp_iface_linearalgebra.py

@@ -587,12 +587,12 @@ def tensordot(a, b, axes=2):
         Second input array. Both inputs `a` and `b` can not be scalars
         at the same time.
     axes : int or (2,) array_like
-        * integer_like
-          If an int `N`, sum over the last `N` axes of `a` and the first `N`
-          axes of `b` in order. The sizes of the corresponding axes must match.
-        * (2,) array_like
-          Or, a list of axes to be summed over, first sequence applying to `a`,
-          second to `b`. Both elements array_like must be of the same length.
+        * integer_like: If an int `N`, sum over the last `N` axes of `a` and
+          the first `N` axes of `b` in order. The sizes of the corresponding
+          axes must match.
+        * (2,) array_like: A list of axes to be summed over, first sequence
+          applying to `a`, second to `b`. Both elements array_like must be of
+          the same length.
 
     Returns
     -------

dpnp/dpnp_iface_manipulation.py

@@ -1861,12 +1861,13 @@ def transpose(a, axes=None):
     ----------
     a : {dpnp.ndarray, usm_ndarray}
         Input array.
-    axes : tuple or list of ints, optional
+    axes : None, tuple or list of ints, optional
         If specified, it must be a tuple or list which contains a permutation
         of [0, 1, ..., N-1] where N is the number of axes of `a`.
         The `i`'th axis of the returned array will correspond to the axis
-        numbered ``axes[i]`` of the input. If not specified, defaults to
-        ``range(a.ndim)[::-1]``, which reverses the order of the axes.
+        numbered ``axes[i]`` of the input. If not specified or ``None``,
+        defaults to ``range(a.ndim)[::-1]``, which reverses the order of
+        the axes.
 
     Returns
     -------
@@ -1881,35 +1882,35 @@ def transpose(a, axes=None):
 
     Examples
     --------
-    >>> import dpnp as dp
-    >>> a = dp.array([[1, 2], [3, 4]])
+    >>> import dpnp as np
+    >>> a = np.array([[1, 2], [3, 4]])
     >>> a
     array([[1, 2],
            [3, 4]])
-    >>> dp.transpose(a)
+    >>> np.transpose(a)
     array([[1, 3],
            [2, 4]])
 
-    >>> a = dp.array([1, 2, 3, 4])
+    >>> a = np.array([1, 2, 3, 4])
     >>> a
     array([1, 2, 3, 4])
-    >>> dp.transpose(a)
+    >>> np.transpose(a)
     array([1, 2, 3, 4])
 
-    >>> a = dp.ones((1, 2, 3))
-    >>> dp.transpose(a, (1, 0, 2)).shape
+    >>> a = np.ones((1, 2, 3))
+    >>> np.transpose(a, (1, 0, 2)).shape
     (2, 1, 3)
 
-    >>> a = dp.ones((2, 3, 4, 5))
-    >>> dp.transpose(a).shape
+    >>> a = np.ones((2, 3, 4, 5))
+    >>> np.transpose(a).shape
     (5, 4, 3, 2)
 
     """
 
     if isinstance(a, dpnp_array):
         array = a
     elif isinstance(a, dpt.usm_ndarray):
-        array = dpnp_array._create_from_usm_ndarray(a.get_array())
+        array = dpnp_array._create_from_usm_ndarray(a)
     else:
         raise TypeError(
             f"An array must be any of supported type, but got {type(a)}"

dpnp/dpnp_iface_mathematical.py

@@ -2802,22 +2802,22 @@ def sum(
         Data type of the returned array. If ``None``, the default data
         type is inferred from the "kind" of the input array data type.
             * If `a` has a real-valued floating-point data type,
-                the returned array will have the default real-valued
-                floating-point data type for the device where input
-                array `a` is allocated.
+              the returned array will have the default real-valued
+              floating-point data type for the device where input
+              array `a` is allocated.
             * If `a` has signed integral data type, the returned array
-                will have the default signed integral type for the device
-                where input array `a` is allocated.
+              will have the default signed integral type for the device
+              where input array `a` is allocated.
             * If `a` has unsigned integral data type, the returned array
-                will have the default unsigned integral type for the device
-                where input array `a` is allocated.
+              will have the default unsigned integral type for the device
+              where input array `a` is allocated.
             * If `a` has a complex-valued floating-point data type,
-                the returned array will have the default complex-valued
-                floating-pointer data type for the device where input
-                array `a` is allocated.
+              the returned array will have the default complex-valued
+              floating-pointer data type for the device where input
+              array `a` is allocated.
             * If `a` has a boolean data type, the returned array will
-                have the default signed integral type for the device
-                where input array `a` is allocated.
+              have the default signed integral type for the device
+              where input array `a` is allocated.
         If the data type (either specified or resolved) differs from the
         data type of `a`, the input array elements are cast to the
         specified data type before computing the sum.

dpnp/dpnp_iface_nanfunctions.py

@@ -717,22 +717,22 @@ def nansum(
         Data type of the returned array. If ``None``, the default data
         type is inferred from the "kind" of the input array data type.
             * If `a` has a real-valued floating-point data type,
-                the returned array will have the default real-valued
-                floating-point data type for the device where input
-                array `a` is allocated.
+              the returned array will have the default real-valued
+              floating-point data type for the device where input
+              array `a` is allocated.
             * If `a` has signed integral data type, the returned array
-                will have the default signed integral type for the device
-                where input array `a` is allocated.
+              will have the default signed integral type for the device
+              where input array `a` is allocated.
             * If `a` has unsigned integral data type, the returned array
-                will have the default unsigned integral type for the device
-                where input array `a` is allocated.
+              will have the default unsigned integral type for the device
+              where input array `a` is allocated.
             * If `a` has a complex-valued floating-point data type,
-                the returned array will have the default complex-valued
-                floating-pointer data type for the device where input
-                array `a` is allocated.
+              the returned array will have the default complex-valued
+              floating-pointer data type for the device where input
+              array `a` is allocated.
             * If `a` has a boolean data type, the returned array will
-                have the default signed integral type for the device
-                where input array `a` is allocated.
+              have the default signed integral type for the device
+              where input array `a` is allocated.
         If the data type (either specified or resolved) differs from the
         data type of `a`, the input array elements are cast to the
         specified data type before computing the sum.

dpnp/dpnp_iface_trigonometric.py

@@ -1355,14 +1355,14 @@ def logsumexp(x, axis=None, out=None, dtype=None, keepdims=False):
         Data type of the returned array. If ``None``, the default data
         type is inferred from the "kind" of the input array data type.
             * If `x` has a real-valued floating-point data type,
-                the returned array will have the default real-valued
-                floating-point data type for the device where input
-                array `x` is allocated.
+              the returned array will have the default real-valued
+              floating-point data type for the device where input
+              array `x` is allocated.
             * If `x` has a boolean or integral data type, the returned array
-                will have the default floating point data type for the device
-                where input array `x` is allocated.
+              will have the default floating point data type for the device
+              where input array `x` is allocated.
             * If `x` has a complex-valued floating-point data type,
-                an error is raised.
+              an error is raised.
         If the data type (either specified or resolved) differs from the
         data type of `x`, the input array elements are cast to the
         specified data type before computing the result. Default: ``None``.

tests/test_manipulation.py

@@ -115,7 +115,7 @@ def test_unique(array):
 
 
 class TestTranspose:
-    @pytest.mark.parametrize("axes", [(0, 1), (1, 0)])
+    @pytest.mark.parametrize("axes", [(0, 1), (1, 0), [0, 1]])
     def test_2d_with_axes(self, axes):
         na = numpy.array([[1, 2], [3, 4]])
         da = dpnp.array(na)
@@ -124,7 +124,22 @@ def test_2d_with_axes(self, axes):
         result = dpnp.transpose(da, axes)
         assert_array_equal(expected, result)
 
-    @pytest.mark.parametrize("axes", [(1, 0, 2), ((1, 0, 2),)])
+        # ndarray
+        expected = na.transpose(axes)
+        result = da.transpose(axes)
+        assert_array_equal(expected, result)
+
+    @pytest.mark.parametrize(
+        "axes",
+        [
+            (1, 0, 2),
+            [1, 0, 2],
+            ((1, 0, 2),),
+            ([1, 0, 2],),
+            [(1, 0, 2)],
+            [[1, 0, 2]],
+        ],
+    )
     def test_3d_with_packed_axes(self, axes):
         na = numpy.ones((1, 2, 3))
         da = dpnp.array(na)
@@ -133,10 +148,27 @@ def test_3d_with_packed_axes(self, axes):
         result = da.transpose(*axes)
         assert_array_equal(expected, result)
 
+        # ndarray
+        expected = na.transpose(*axes)
+        result = da.transpose(*axes)
+        assert_array_equal(expected, result)
+
     @pytest.mark.parametrize("shape", [(10,), (2, 4), (5, 3, 7), (3, 8, 4, 1)])
     def test_none_axes(self, shape):
         na = numpy.ones(shape)
         da = dpnp.ones(shape)
 
+        assert_array_equal(numpy.transpose(na), dpnp.transpose(da))
+        assert_array_equal(numpy.transpose(na, None), dpnp.transpose(da, None))
+
+        # ndarray
         assert_array_equal(na.transpose(), da.transpose())
         assert_array_equal(na.transpose(None), da.transpose(None))
+
+    def test_ndarray_axes_n_int(self):
+        na = numpy.ones((1, 2, 3))
+        da = dpnp.array(na)
+
+        expected = na.transpose(1, 0, 2)
+        result = da.transpose(1, 0, 2)
+        assert_array_equal(expected, result)

tests/third_party/cupy/manipulation_tests/test_transpose.py

@@ -64,20 +64,38 @@ def test_moveaxis_invalid2_2(self):
             with pytest.raises(numpy.AxisError):
                 xp.moveaxis(a, [0, -4], [1, 2])
 
+    def test_moveaxis_invalid2_3(self):
+        for xp in (numpy, cupy):
+            a = testing.shaped_arange((2, 3, 4), xp)
+            with pytest.raises(numpy.AxisError):
+                xp.moveaxis(a, -4, 0)
+
     # len(source) != len(destination)
-    def test_moveaxis_invalid3(self):
+    def test_moveaxis_invalid3_1(self):
         for xp in (numpy, cupy):
             a = testing.shaped_arange((2, 3, 4), xp)
             with pytest.raises(ValueError):
                 xp.moveaxis(a, [0, 1, 2], [1, 2])
 
+    def test_moveaxis_invalid3_2(self):
+        for xp in (numpy, cupy):
+            a = testing.shaped_arange((2, 3, 4), xp)
+            with pytest.raises(ValueError):
+                xp.moveaxis(a, 0, [1, 2])
+
     # len(source) != len(destination)
-    def test_moveaxis_invalid4(self):
+    def test_moveaxis_invalid4_1(self):
         for xp in (numpy, cupy):
             a = testing.shaped_arange((2, 3, 4), xp)
             with pytest.raises(ValueError):
                 xp.moveaxis(a, [0, 1], [1, 2, 0])
 
+    def test_moveaxis_invalid4_2(self):
+        for xp in (numpy, cupy):
+            a = testing.shaped_arange((2, 3, 4), xp)
+            with pytest.raises(ValueError):
+                xp.moveaxis(a, [0, 1], 1)
+
     # Use the same axis twice
     def test_moveaxis_invalid5_1(self):
         for xp in (numpy, cupy):