Repository: cupy/cupy

Issue:
Columnar matrix slice doesn't work properly with ravel
When creating a view into a 2d array using `[:, n:]` notation, the resulting view doesn't work properly with `ravel` or functions that depend on it.

```
import cupy as cp
a = cp.zeros((3, 3))
b = a[:, 1:]
b.ravel()[0] = 1
print(b)
b[0, 0] = 2
print(a)
```

``` output
[[0. 0.]
 [0. 0.]
 [0. 0.]]
[[0. 2. 0.]
 [0. 0. 0.]
 [0. 0. 0.]]
```

As you can see, `b[0, 0]` still works as expected, but the array produced by `ravel` is a copy. This is not the case for the array produced with a row-wise slice:

```
c = a[1:, :]
c.ravel()[0] = 3
print(a)
```

``` output
array([[0., 2., 0.],
       [3., 0., 0.],
       [0., 0., 0.]])
```

This affects me because I'm trying to fill out a tridiagonal matrix using `cupy`, but the upper diagonal can't be properly referenced:
```
d = cp.zeros((3, 3))
cp.fill_diagonal(d, 3)
cp.fill_diagonal(d[1:], 2)
cp.fill_diagonal(d[:, 1:], 1)
print(d)
```

``` output
[[3. 0. 0.]
 [2. 3. 0.]
 [0. 2. 3.]]
```

```
In [182]: cp.show_config()
CuPy Version          : 7.1.0
CUDA Root             : /usr/local/cuda
CUDA Build Version    : 10000
CUDA Driver Version   : 10000
CUDA Runtime Version  : 10000
cuBLAS Version        : 10000
cuFFT Version         : 10000
cuRAND Version        : 10000
cuSOLVER Version      : (10, 0, 0)
cuSPARSE Version      : 10000
NVRTC Version         : (10, 0)
cuDNN Build Version   : 7605
cuDNN Version         : 7600
NCCL Build Version    : 2406
NCCL Runtime Version  : 2406
```


Discussion:
Thanks for your report. We've got that the issue is not the way of `cupy.ndarray.ravel` but `cupy.fill_diagonal` uses `ravel` in its implementation as opposed to `numpy.fill_diagonal` that uses `numpy.ndarray.flat`.

The value of `numpy.ndarray.flat` is a `numpy.flatiter` instance. It is supposed to act as Python's  built-in iterator object and the concept is essentially serial, so providing the same feature in CuPy would be fairly inefficient. Please let us consider how to address this issue. A possible option would be providing `cupy.flatiter` that worked in parallel with giving some restriction compared to `numpy.flatiter`.

I think that creating a new type of `Indexer` could work

Resolving patch:
diff --git a/cupy/indexing/insert.py b/cupy/indexing/insert.py
--- a/cupy/indexing/insert.py
+++ b/cupy/indexing/insert.py
@@ -1,6 +1,12 @@
+import string
+
 import numpy
 
 import cupy
+from cupy import util
+from cupy.core import _carray
+from cupy.core import _scalar
+from cupy.cuda import device
 
 
 def place(arr, mask, vals):
@@ -69,6 +75,33 @@ def put(a, ind, v, mode='wrap'):
 # TODO(okuta): Implement putmask
 
 
+_fill_diagonal_template = string.Template(r'''
+#include <cupy/complex.cuh>
+#include <cupy/carray.cuh>
+extern "C" __global__
+void cupy_fill_diagonal(CArray<${type}, ${a_ndim}> a,
+                        CIndexer<${a_ndim}> a_ind,
+                        int start,
+                        int stop,
+                        int step,
+                        CArray<${type}, ${val_ndim}> val,
+                        CIndexer<${val_ndim}> val_ind) {
+    int n = (stop - start) / step + 1;
+    CUPY_FOR(i, n) {
+        a_ind.set(start + i * step);
+        val_ind.set(i % val_ind.size());
+        a[a_ind.get()] = val[val_ind.get()];
+    }
+}''')
+
+
+@util.memoize(for_each_device=True)
+def _fill_diagonal_kernel(type, a_ndim, val_ndim):
+    code = _fill_diagonal_template.substitute(
+        type=type, a_ndim=a_ndim, val_ndim=val_ndim)
+    return cupy.RawKernel(code, 'cupy_fill_diagonal')
+
+
 def fill_diagonal(a, val, wrap=False):
     """Fills the main diagonal of the given array of any dimensionality.
 
@@ -97,7 +130,7 @@ def fill_diagonal(a, val, wrap=False):
     # The followings are imported from the original numpy
     if a.ndim < 2:
         raise ValueError('array must be at least 2-d')
-    end = None
+    end = a.size
     if a.ndim == 2:
         step = a.shape[1] + 1
         if not wrap:
@@ -107,6 +140,21 @@ def fill_diagonal(a, val, wrap=False):
             raise ValueError('All dimensions of input must be of equal length')
         step = 1 + numpy.cumprod(a.shape[:-1]).sum()
 
-    # Since the current cupy does not support a.flat,
-    # we use a.ravel() instead of a.flat
-    a.ravel()[:end:step] = val
+    val = cupy.asarray(val, dtype=a.dtype)
+
+    dev_id = device.get_device_id()
+    for arr in [a, val]:
+        if arr.data.device_id != dev_id:
+            raise ValueError(
+                'Array device must be same as the current '
+                'device: array device = %d while current = %d'
+                % (arr.data.device_id, dev_id))
+
+    typename = _scalar.get_typename(a.dtype)
+    fill_diagonal_kernel = _fill_diagonal_kernel(typename, a.ndim, val.ndim)
+
+    size = end // step + 1
+    a_ind = _carray.Indexer(a.shape)
+    val_ind = _carray.Indexer(val.shape)
+    fill_diagonal_kernel.kernel.linear_launch(
+        size, (a, a_ind, 0, end, step, val, val_ind))


Test patch:
diff --git a/tests/cupy_tests/indexing_tests/test_insert.py b/tests/cupy_tests/indexing_tests/test_insert.py
--- a/tests/cupy_tests/indexing_tests/test_insert.py
+++ b/tests/cupy_tests/indexing_tests/test_insert.py
@@ -1,6 +1,7 @@
 import unittest
 
 import numpy
+import pytest
 
 from cupy import testing
 
@@ -112,21 +113,40 @@ def test_put_mode_error(self, xp, dtype):
 
 @testing.parameterize(*testing.product({
     'shape': [(3, 3), (2, 2, 2), (3, 5), (5, 3)],
-    'val': [1, 0],
+    'val': [1, 0, (2,), (2, 2)],
     'wrap': [True, False],
 }))
 @testing.gpu
 class TestFillDiagonal(unittest.TestCase):
 
+    def _compute_val(self, xp):
+        if type(self.val) is int:
+            return self.val
+        else:
+            return xp.arange(numpy.prod(self.val)).reshape(self.val)
+
     @testing.for_all_dtypes()
     @testing.numpy_cupy_array_equal()
     def test_fill_diagonal(self, xp, dtype):
         a = testing.shaped_arange(self.shape, xp, dtype)
-        xp.fill_diagonal(a, val=self.val, wrap=self.wrap)
+        val = self._compute_val(xp)
+        xp.fill_diagonal(a, val=val, wrap=self.wrap)
+        return a
+
+    @testing.for_all_dtypes()
+    @testing.numpy_cupy_array_equal()
+    def test_columnar_slice(self, xp, dtype):  # see cupy#2970
+        if self.shape == (2, 2, 2):
+            pytest.skip(
+                'The length of each dimension must be the same after slicing')
+        a = testing.shaped_arange(self.shape, xp, dtype)
+        val = self._compute_val(xp)
+        xp.fill_diagonal(a[:, 1:], val=val, wrap=self.wrap)
         return a
 
     @testing.for_all_dtypes()
     @testing.numpy_cupy_raises()
     def test_1darray(self, xp, dtype):
         a = testing.shaped_arange(5, xp, dtype)
-        xp.fill_diagonal(a, val=self.val, wrap=self.wrap)
+        val = self._compute_val(xp)
+        xp.fill_diagonal(a, val=val, wrap=self.wrap)