Add dpnp.roll implementation and update dpnp.rollaxis/dpnp.moveaxis (#1517)

* Add a new dpnp.roll function using dpctl.tensor.roll impl

* Update dpnp.rollaxis and add tests for it

* Update dpnp.moveaxis

* Update implemantations of rolls funcs and fix remarks

* Raise ValueError instead of numpy.AxisError

* Add support for axis=None and shift is tuple in dpnp.roll

Author: vlad-perevezentsev

.github/workflows/conda-package.yml

@@ -32,6 +32,8 @@ env:
       third_party/cupy/logic_tests/test_truth.py
       third_party/cupy/manipulation_tests/test_basic.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

@@ -35,6 +35,7 @@ Transpose-like operations
    :nosignatures:
 
    dpnp.moveaxis
+   dpnp.roll
    dpnp.rollaxis
    dpnp.swapaxes
    dpnp.ndarray.T

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 *
@@ -67,6 +68,7 @@
     "repeat",
     "reshape",
     "result_type",
+    "roll",
     "rollaxis",
     "shape",
     "squeeze",
@@ -692,7 +694,8 @@ def moveaxis(a, source, destination):
     Limitations
     -----------
     Parameters `a` is supported as either :class:`dpnp.ndarray`
-    or :class:`dpctl.tensor.usm_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`.
     Otherwise ``TypeError`` exception will be raised.
 
@@ -910,18 +913,81 @@ 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]])
+
+    """
+    if axis is None:
+        return roll(x.reshape(-1), shift, 0).reshape(x.shape)
+    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
@@ -943,19 +1009,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 ValueError(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,68 @@ 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.parametrize(
+        ("axis", "start"),
+        [
+            (-5, 0),
+            (0, -5),
+            (4, 0),
+            (0, 5),
+        ],
+    )
+    def test_exceptions(self, axis, start):
+        a = dpnp.arange(1 * 2 * 3 * 4).reshape(1, 2, 3, 4)
+        assert_raises(ValueError, dpnp.rollaxis, a, axis, start)
+
+    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

@@ -25,7 +25,6 @@
     {"shape": (5, 2), "shift": (2, 1, 3), "axis": 0},
     {"shape": (5, 2), "shift": (2, 1, 3), "axis": None},
 )
-@pytest.mark.skip("`roll` isn't supported yet")
 class TestRoll(unittest.TestCase):
     @testing.for_all_dtypes()
     @testing.numpy_cupy_array_equal()
@@ -38,12 +37,9 @@ def test_roll(self, xp, dtype):
     def test_roll_cupy_shift(self, xp, dtype):
         x = testing.shaped_arange(self.shape, xp, dtype)
         shift = self.shift
-        if xp is cupy:
-            shift = cupy.array(shift)
         return xp.roll(x, shift, axis=self.axis)
 
 
-@pytest.mark.skip("`roll` isn't supported yet")
 class TestRollTypeError(unittest.TestCase):
     def test_roll_invalid_shift(self):
         for xp in (numpy, cupy):
@@ -66,7 +62,6 @@ def test_roll_invalid_axis_type(self):
     {"shape": (5, 2), "shift": 1, "axis": -3},
     {"shape": (5, 2), "shift": 1, "axis": (1, -3)},
 )
-@pytest.mark.skip("`roll` isn't supported yet")
 class TestRollValueError(unittest.TestCase):
     def test_roll_invalid(self):
         for xp in (numpy, cupy):
@@ -78,8 +73,6 @@ def test_roll_invalid_cupy_shift(self):
         for xp in (numpy, cupy):
             x = testing.shaped_arange(self.shape, xp)
             shift = self.shift
-            if xp is cupy:
-                shift = cupy.array(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)