Add dpnp.roll implementation and update dpnp.rollaxis/dpnp.moveaxis

.github/workflows/conda-package.yml

@@ -25,6 +25,8 @@ env:
       test_usm_type.py
       third_party/cupy/linalg_tests/test_product.py
       third_party/cupy/manipulation_tests/test_join.py
+      third_party/cupy/manipulation_tests/test_rearrange.py
+      third_party/cupy/manipulation_tests/test_transpose.py
       third_party/cupy/math_tests/test_explog.py
       third_party/cupy/math_tests/test_misc.py
       third_party/cupy/math_tests/test_trigonometric.py

doc/reference/manipulation.rst

@@ -32,6 +32,7 @@ Transpose-like operations
    :nosignatures:
 
    dpnp.moveaxis
+   dpnp.roll
    dpnp.rollaxis
    dpnp.swapaxes
    dpnp.transpose

dpnp/dpnp_iface_manipulation.py

@@ -42,6 +42,7 @@
 
 import dpctl.tensor as dpt
 import numpy
+from numpy.core.numeric import normalize_axis_index
 
 import dpnp
 from dpnp.dpnp_algo import *
@@ -64,6 +65,7 @@
     "repeat",
     "reshape",
     "result_type",
+    "roll",
     "rollaxis",
     "shape",
     "squeeze",
@@ -473,8 +475,8 @@ def moveaxis(x, source, destination):
     Limitations
     -----------
     Parameters `x` is supported as either :class:`dpnp.ndarray`
-    or :class:`dpctl.tensor.usm_ndarray`.
-    Otherwise the function will be executed sequentially on CPU.
+    or :class:`dpctl.tensor.usm_ndarray`. Otherwise ``TypeError`` exception
+    will be raised.
     Input array data types are limited by supported DPNP :ref:`Data types`.
 
     See Also
@@ -493,13 +495,10 @@ def moveaxis(x, source, destination):
 
     """
 
-    if isinstance(x, dpnp_array) or isinstance(x, dpt.usm_ndarray):
-        dpt_array = x.get_array() if isinstance(x, dpnp_array) else x
-        return dpnp_array._create_from_usm_ndarray(
-            dpt.moveaxis(dpt_array, source, destination)
-        )
-
-    return call_origin(numpy.moveaxis, x, source, destination)
+    dpt_array = dpnp.get_usm_ndarray(x)
+    return dpnp_array._create_from_usm_ndarray(
+        dpt.moveaxis(dpt_array, source, destination)
+    )
 
 
 def ravel(x1, order="C"):
@@ -694,18 +693,80 @@ def result_type(*arrays_and_dtypes):
     return dpt.result_type(*usm_arrays_and_dtypes)
 
 
-def rollaxis(x1, axis, start=0):
+def roll(x, shift, axis=None):
+    """
+    Roll the elements of an array by a number of positions along a given axis.
+
+    Array elements that roll beyond the last position are re-introduced
+    at the first position. Array elements that roll beyond the first position
+    are re-introduced at the last position.
+
+    For full documentation refer to :obj:`numpy.roll`.
+
+    Returns
+    -------
+    dpnp.ndarray
+        An array with the same data type as `x`
+        and whose elements, relative to `x`, are shifted.
+
+    Limitations
+    -----------
+    Parameter `x` is supported either as :class:`dpnp.ndarray`
+    or :class:`dpctl.tensor.usm_ndarray`. Otherwise ``TypeError`` exception
+    will be raised.
+    Input array data types are limited by supported DPNP :ref:`Data types`.
+
+
+    See Also
+    --------
+    :obj:`dpnp.moveaxis` : Move array axes to new positions.
+    :obj:`dpnp.rollaxis` : Roll the specified axis backwards
+                       until it lies in a given position.
+
+    Examples
+    --------
+    >>> import dpnp as np
+    >>> x1 = np.arange(10)
+    >>> np.roll(x1, 2)
+    array([8, 9, 0, 1, 2, 3, 4, 5, 6, 7])
+
+    >>> np.roll(x1, -2)
+    array([2, 3, 4, 5, 6, 7, 8, 9, 0, 1])
+
+    >>> x2 = np.reshape(x1, (2, 5))
+    >>> np.roll(x2, 1, axis=0)
+    array([[5, 6, 7, 8, 9],
+           [0, 1, 2, 3, 4]])
+
+    >>> np.roll(x2, (2, 1), axis=(1, 0))
+    array([[8, 9, 5, 6, 7],
+           [3, 4, 0, 1, 2]])
+
+    """
+
+    dpt_array = dpnp.get_usm_ndarray(x)
+    return dpnp_array._create_from_usm_ndarray(
+        dpt.roll(dpt_array, shift=shift, axis=axis)
+    )
+
+
+def rollaxis(x, axis, start=0):
     """
     Roll the specified axis backwards, until it lies in a given position.
 
     For full documentation refer to :obj:`numpy.rollaxis`.
 
+    Returns
+    -------
+    dpnp.ndarray
+        An array with the same data type as `x` where the specified axis
+        has been repositioned to the desired position.
+
     Limitations
     -----------
-    Input array is supported as :obj:`dpnp.ndarray`.
-    Parameter ``axis`` is supported as integer only.
-    Parameter ``start`` is limited by ``-a.ndim <= start <= a.ndim``.
-    Otherwise the function will be executed sequentially on CPU.
+    Parameter `x` is supported either as :class:`dpnp.ndarray`
+    or :class:`dpctl.tensor.usm_ndarray`. Otherwise ``TypeError`` exception
+    will be raised.
     Input array data types are limited by supported DPNP :ref:`Data types`.
 
     See Also
@@ -727,19 +788,19 @@ def rollaxis(x1, axis, start=0):
 
     """
 
-    x1_desc = dpnp.get_dpnp_descriptor(x1, copy_when_nondefault_queue=False)
-    if x1_desc:
-        if not isinstance(axis, int):
-            pass
-        elif start < -x1_desc.ndim or start > x1_desc.ndim:
-            pass
-        else:
-            start_norm = start + x1_desc.ndim if start < 0 else start
-            destination = start_norm - 1 if start_norm > axis else start_norm
-
-            return dpnp.moveaxis(x1_desc.get_pyobj(), axis, destination)
-
-    return call_origin(numpy.rollaxis, x1, axis, start)
+    n = x.ndim
+    axis = normalize_axis_index(axis, n)
+    if start < 0:
+        start += n
+    msg = "'%s' arg requires %d <= %s < %d, but %d was passed in"
+    if not (0 <= start < n + 1):
+        raise numpy.AxisError(msg % ("start", -n, "start", n + 1, start))
+    if axis < start:
+        start -= 1
+    if axis == start:
+        return x
+    dpt_array = dpnp.get_usm_ndarray(x)
+    return dpnp.moveaxis(dpt_array, source=axis, destination=start)
 
 
 def shape(a):

tests/test_arraymanipulation.py

@@ -605,3 +605,62 @@ def test_2D_array2(self):
     def test_generator(self):
         with assert_warns(FutureWarning):
             dpnp.vstack((numpy.arange(3) for _ in range(2)))
+
+
+class TestRollaxis:
+    data = [
+        (0, 0),
+        (0, 1),
+        (0, 2),
+        (0, 3),
+        (0, 4),
+        (1, 0),
+        (1, 1),
+        (1, 2),
+        (1, 3),
+        (1, 4),
+        (2, 0),
+        (2, 1),
+        (2, 2),
+        (2, 3),
+        (2, 4),
+        (3, 0),
+        (3, 1),
+        (3, 2),
+        (3, 3),
+        (3, 4),
+    ]
+
+    def test_exceptions(self):
+        a = dpnp.arange(1 * 2 * 3 * 4).reshape(1, 2, 3, 4)
+        assert_raises(numpy.AxisError, dpnp.rollaxis, a, -5, 0)
+        assert_raises(numpy.AxisError, dpnp.rollaxis, a, 0, -5)
+        assert_raises(numpy.AxisError, dpnp.rollaxis, a, 4, 0)
+        assert_raises(numpy.AxisError, dpnp.rollaxis, a, 0, 5)
+
+    def test_results(self):
+        np_a = numpy.arange(1 * 2 * 3 * 4).reshape(1, 2, 3, 4)
+        dp_a = dpnp.array(np_a)
+        for i, j in self.data:
+            # positive axis, positive start
+            res = dpnp.rollaxis(dp_a, axis=i, start=j)
+            exp = numpy.rollaxis(np_a, axis=i, start=j)
+            assert res.shape == exp.shape
+
+            # negative axis, positive start
+            ip = i + 1
+            res = dpnp.rollaxis(dp_a, axis=-ip, start=j)
+            exp = numpy.rollaxis(np_a, axis=-ip, start=j)
+            assert res.shape == exp.shape
+
+            # positive axis, negative start
+            jp = j + 1 if j < 4 else j
+            res = dpnp.rollaxis(dp_a, axis=i, start=-jp)
+            exp = numpy.rollaxis(np_a, axis=i, start=-jp)
+            assert res.shape == exp.shape
+
+            # negative axis, negative start
+            ip = i + 1
+            jp = j + 1 if j < 4 else j
+            res = dpnp.rollaxis(dp_a, axis=-ip, start=-jp)
+            exp = numpy.rollaxis(np_a, axis=-ip, start=-jp)

tests/third_party/cupy/manipulation_tests/test_rearrange.py

@@ -0,0 +1,78 @@
+import unittest
+
+import numpy
+import pytest
+
+import dpnp as cupy
+from tests.third_party.cupy import testing
+
+
+@testing.parameterize(
+    {"shape": (10,), "shift": 2, "axis": None},
+    {"shape": (5, 2), "shift": 1, "axis": None},
+    {"shape": (5, 2), "shift": -2, "axis": None},
+    {"shape": (5, 2), "shift": 1, "axis": 0},
+    {"shape": (5, 2), "shift": 1, "axis": -1},
+    {"shape": (10,), "shift": 35, "axis": None},
+    {"shape": (5, 2), "shift": 11, "axis": 0},
+    {"shape": (), "shift": 5, "axis": None},
+    {"shape": (5, 2), "shift": 1, "axis": (0, 1)},
+    {"shape": (5, 2), "shift": 1, "axis": (0, 0)},
+    {"shape": (5, 2), "shift": 50, "axis": 0},
+    {"shape": (5, 2), "shift": (2, 1), "axis": (0, 1)},
+    {"shape": (5, 2), "shift": (2, 1), "axis": (0, -1)},
+    {"shape": (5, 2), "shift": (2, 1), "axis": (1, -1)},
+    {"shape": (5, 2), "shift": (2, 1, 3), "axis": 0},
+    # TODO: remove when dpctl is fixed
+    # {'shape': (5, 2), 'shift': (2, 1, 3), 'axis': None},
+)
+class TestRoll(unittest.TestCase):
+    @testing.for_all_dtypes()
+    @testing.numpy_cupy_array_equal()
+    def test_roll(self, xp, dtype):
+        x = testing.shaped_arange(self.shape, xp, dtype)
+        return xp.roll(x, self.shift, axis=self.axis)
+
+    @testing.for_all_dtypes()
+    @testing.numpy_cupy_array_equal()
+    def test_roll_cupy_shift(self, xp, dtype):
+        x = testing.shaped_arange(self.shape, xp, dtype)
+        shift = self.shift
+        return xp.roll(x, shift, axis=self.axis)
+
+
+class TestRollTypeError(unittest.TestCase):
+    def test_roll_invalid_shift(self):
+        for xp in (numpy, cupy):
+            x = testing.shaped_arange((5, 2), xp)
+            with pytest.raises(TypeError):
+                xp.roll(x, "0", axis=0)
+
+    def test_roll_invalid_axis_type(self):
+        for xp in (numpy, cupy):
+            x = testing.shaped_arange((5, 2), xp)
+            with pytest.raises(TypeError):
+                xp.roll(x, 2, axis="0")
+
+
+@testing.parameterize(
+    {"shape": (5, 2, 3), "shift": (2, 2, 2), "axis": (0, 1)},
+    {"shape": (5, 2), "shift": 1, "axis": 2},
+    {"shape": (5, 2), "shift": 1, "axis": -3},
+    {"shape": (5, 2, 2), "shift": (1, 0), "axis": (0, 1, 2)},
+    {"shape": (5, 2), "shift": 1, "axis": -3},
+    {"shape": (5, 2), "shift": 1, "axis": (1, -3)},
+)
+class TestRollValueError(unittest.TestCase):
+    def test_roll_invalid(self):
+        for xp in (numpy, cupy):
+            x = testing.shaped_arange(self.shape, xp)
+            with pytest.raises(ValueError):
+                xp.roll(x, self.shift, axis=self.axis)
+
+    def test_roll_invalid_cupy_shift(self):
+        for xp in (numpy, cupy):
+            x = testing.shaped_arange(self.shape, xp)
+            shift = self.shift
+            with pytest.raises(ValueError):
+                xp.roll(x, shift, axis=self.axis)

tests/third_party/cupy/manipulation_tests/test_transpose.py

@@ -103,7 +103,6 @@ def test_rollaxis(self, xp):
         a = testing.shaped_arange((2, 3, 4), xp)
         return xp.rollaxis(a, 2)
 
-    @pytest.mark.usefixtures("allow_fall_back_on_numpy")
     def test_rollaxis_failure(self):
         for xp in (numpy, cupy):
             a = testing.shaped_arange((2, 3, 4), xp)