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

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

@@ -64,6 +64,7 @@
     "repeat",
     "reshape",
     "result_type",
+    "roll",
     "rollaxis",
     "shape",
     "squeeze",
@@ -493,8 +494,8 @@ 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
+    if dpnp.is_supported_array_type(x):
+        dpt_array = dpnp.get_usm_ndarray(x)
         return dpnp_array._create_from_usm_ndarray(
             dpt.moveaxis(dpt_array, source, destination)
         )
@@ -694,19 +695,84 @@ 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`
+        containing elements are shifted relative to `x`.
+
+    Limitations
+    -----------
+    Parameter `x` is supported either as :class:`dpnp.ndarray`
+    or :class:`dpctl.tensor.usm_ndarray`.
+    Input array data types are limited by supported DPNP :ref:`Data types`.
+    Otherwise the function will be executed sequentially on CPU.
+
+    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]])
+
+    """
+
+    if dpnp.is_supported_array_type(x):
+        dpt_array = dpnp.get_usm_ndarray(x)
+        return dpnp_array._create_from_usm_ndarray(
+            dpt.roll(dpt_array, shift=shift, axis=axis)
+        )
+
+    return call_origin(numpy.roll, x, 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`.
+    Parameter `start` is limited by `-x.ndim <= start <= x.ndim`.
     Input array data types are limited by supported DPNP :ref:`Data types`.
+    Otherwise the function will be executed sequentially on CPU.
 
     See Also
     --------
@@ -727,19 +793,22 @@ 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:
+    if dpnp.is_supported_array_type(x):
+        if start < -x.ndim or start > x.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
+            axis_norm = axis + x.ndim if axis < 0 else axis
+            start_norm = start + x.ndim if start < 0 else start
+            destination = (
+                start_norm - 1 if start_norm > axis_norm else start_norm
+            )
 
-            return dpnp.moveaxis(x1_desc.get_pyobj(), axis, destination)
+            dpt_array = dpnp.get_usm_ndarray(x)
+            return dpnp.moveaxis(
+                dpt_array, source=axis, destination=destination
+            )
 
-    return call_origin(numpy.rollaxis, x1, axis, start)
+    return call_origin(numpy.rollaxis, x, axis=axis, start=start)
 
 
 def shape(a):

tests/test_arraymanipulation.py

@@ -605,3 +605,63 @@ 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),
+    ]
+
+    @pytest.mark.usefixtures("allow_fall_back_on_numpy")
+    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,76 @@
+import unittest
+
+import numpy
+import pytest
+
+import dpnp
+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},
+)
+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, dpnp):
+            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, dpnp):
+            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, dpnp):
+            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, dpnp):
+            x = testing.shaped_arange(self.shape, xp)
+            shift = self.shift
+            with pytest.raises(ValueError):
+                xp.roll(x, shift, axis=self.axis)