Repository: chainer/chainer

Issue:
Implement cupy.roll
I want `cupy.roll`.
http://docs.scipy.org/doc/numpy/reference/generated/numpy.roll.html



Discussion:

Resolving patch:
diff --git a/cupy/__init__.py b/cupy/__init__.py
--- a/cupy/__init__.py
+++ b/cupy/__init__.py
@@ -135,6 +135,8 @@
 tile = manipulation.tiling.tile
 repeat = manipulation.tiling.repeat
 
+roll = manipulation.rearrange.roll
+
 # -----------------------------------------------------------------------------
 # Binary operations
 # -----------------------------------------------------------------------------
diff --git a/cupy/manipulation/rearrange.py b/cupy/manipulation/rearrange.py
--- a/cupy/manipulation/rearrange.py
+++ b/cupy/manipulation/rearrange.py
@@ -1,5 +1,4 @@
-# flake8: NOQA
-# "flake8: NOQA" to suppress warning "H104  File contains nothing but comments"
+import cupy
 
 # TODO(okuta): Implement fliplr
 
@@ -7,7 +6,54 @@
 # TODO(okuta): Implement flipud
 
 
-# TODO(okuta): Implement roll
+def roll(a, shift, axis=None):
+    """Roll array elements along a given axis.
 
+    Args:
+        a (~cupy.ndarray): Array to be rolled.
+        shift (int): The number of places by which elements are shifted.
+        axis (int or None): The axis along which elements are shifted.
+            If ``axis`` is ``None``, the array is flattend before shifting,
+            and afther that it is reshaped to the original shape.
+
+    Returns:
+        ~cupy.ndarray: Output array.
+
+    .. seealso:: :func:`numpy.roll`
+
+    """
+    if axis is None:
+        if a.size == 0:
+            return a
+        size = a.size
+        ra = a.ravel()
+        shift %= size
+        res = cupy.empty((size,), a.dtype)
+        res[:shift] = ra[size - shift:]
+        res[shift:] = ra[:size - shift]
+        return res.reshape(a.shape)
+    else:
+        axis = int(axis)
+        if axis < 0:
+            axis += a.ndim
+        if not 0 <= axis < a.ndim:
+            raise ValueError('axis must be >= %d and < %d' % (-a.ndim, a.ndim))
+        size = a.shape[axis]
+        if size == 0:
+            return a
+        shift %= size
+        prev = (slice(None),) * axis
+        rest = (slice(None),) * (a.ndim - axis - 1)
+        # Roll only the dimensiont at the given axis
+        # ind1 is [:, ..., size-shift:, ..., :]
+        # ind2 is [:, ..., :size-shift, ..., :]
+        ind1 = prev + (slice(size - shift, None, None),) + rest
+        ind2 = prev + (slice(None, size - shift, None),) + rest
+        r_ind1 = prev + (slice(None, shift, None),) + rest
+        r_ind2 = prev + (slice(shift, None, None),) + rest
+        res = cupy.empty_like(a)
+        res[r_ind1] = a[ind1]
+        res[r_ind2] = a[ind2]
+        return res
 
 # TODO(okuta): Implement rot90


Test patch:
diff --git a/tests/cupy_tests/manipulation_tests/test_rearrange.py b/tests/cupy_tests/manipulation_tests/test_rearrange.py
--- a/tests/cupy_tests/manipulation_tests/test_rearrange.py
+++ b/tests/cupy_tests/manipulation_tests/test_rearrange.py
@@ -7,3 +7,63 @@
 class TestRearrange(unittest.TestCase):
 
     _multiprocess_can_split_ = True
+
+    @testing.for_all_dtypes()
+    @testing.numpy_cupy_array_equal()
+    def test_roll(self, xp, dtype):
+        x = xp.arange(10)
+        return xp.roll(x, 2)
+
+    @testing.for_all_dtypes()
+    @testing.numpy_cupy_array_equal()
+    def test_roll2(self, xp, dtype):
+        x = testing.shaped_arange((5, 2), xp, dtype)
+        return xp.roll(x, 1)
+
+    @testing.for_all_dtypes()
+    @testing.numpy_cupy_array_equal()
+    def test_roll_negative(self, xp, dtype):
+        x = testing.shaped_arange((5, 2), xp, dtype)
+        return xp.roll(x, -2)
+
+    @testing.for_all_dtypes()
+    @testing.numpy_cupy_array_equal()
+    def test_roll_with_axis(self, xp, dtype):
+        x = testing.shaped_arange((5, 2), xp, dtype)
+        return xp.roll(x, 1, axis=0)
+
+    @testing.for_all_dtypes()
+    @testing.numpy_cupy_array_equal()
+    def test_roll_with_negative_axis(self, xp, dtype):
+        x = testing.shaped_arange((5, 2), xp, dtype)
+        return xp.roll(x, 1, axis=-1)
+
+    @testing.for_all_dtypes()
+    @testing.numpy_cupy_array_equal()
+    def test_roll_double_shift(self, xp, dtype):
+        x = testing.shaped_arange((10,), xp, dtype)
+        return xp.roll(x, 35)
+
+    @testing.for_all_dtypes()
+    @testing.numpy_cupy_array_equal()
+    def test_roll_double_shift_with_axis(self, xp, dtype):
+        x = testing.shaped_arange((5, 2), xp, dtype)
+        return xp.roll(x, 11, axis=0)
+
+    @testing.for_all_dtypes()
+    @testing.numpy_cupy_array_equal()
+    def test_roll_zero_array(self, xp, dtype):
+        x = testing.shaped_arange((), xp, dtype)
+        return xp.roll(x, 5)
+
+    @testing.for_all_dtypes()
+    @testing.numpy_cupy_raises()
+    def test_roll_invalid_axis(self, xp, dtype):
+        x = testing.shaped_arange((5, 2), xp, dtype)
+        return xp.roll(x, 1, axis=2)
+
+    @testing.for_all_dtypes()
+    @testing.numpy_cupy_raises()
+    def test_roll_invalid_negative_axis(self, xp, dtype):
+        x = testing.shaped_arange((5, 2), xp, dtype)
+        return xp.roll(x, 1, axis=-3)