Repository: cupy/cupy

Issue:
Sparse arrays do not support all of Scipy's slicing options
I encountered this behavior in the most recent (CuPy 8.0.0 at the time of writing this) branch when attempting to call `dask.array.from_array()` with a `csr_matrix`. It appears that Dask is attempting to do some slicing that is supported in Scipy's `csr_matrix` but is not supported by CuPy. 

The following is a reproducible example:

```
>>> import cupy as cp
>>> import dask.array
>>> 
>>> a = cp.sparse.random(100, 100, format='csr')
>>> b = dask.array.from_array(a)
>>> b.compute()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/share/workspace/dask/dask/base.py", line 165, in compute
    (result,) = compute(self, traverse=False, **kwargs)
  File "/share/workspace/dask/dask/base.py", line 436, in compute
    results = schedule(dsk, keys, **kwargs)
  File "/share/workspace/dask/dask/threaded.py", line 81, in get
    **kwargs
  File "/share/workspace/dask/dask/local.py", line 486, in get_async
    raise_exception(exc, tb)
  File "/share/workspace/dask/dask/local.py", line 316, in reraise
    raise exc
  File "/share/workspace/dask/dask/local.py", line 222, in execute_task
    result = _execute_task(task, data)
  File "/share/workspace/dask/dask/core.py", line 119, in _execute_task
    return func(*args2)
  File "/share/workspace/dask/dask/array/core.py", line 104, in getter
    c = a[b]
  File "/share/workspace/cupy/cupyx/scipy/sparse/compressed.py", line 509, in __getitem__
    raise ValueError('unsupported indexing')
ValueError: unsupported indexing
```

Here's the same code using Scipy:

```
>>> import scipy
>>> 
>>> a = scipy.sparse.random(100, 100, format='csr')
>>> b = dask.array.from_array(a)
>>> b.compute()
array(<100x100 sparse matrix of type '<class 'numpy.float64'>'
	with 100 stored elements in Compressed Sparse Row format>, dtype=object)
```


When using a CuPy `coo_matrix`:

```
>>> a = cp.sparse.random(100, 100, format='coo')
>>> b = dask.array.from_array(a)
>>> b.compute()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/share/workspace/dask/dask/base.py", line 165, in compute
    (result,) = compute(self, traverse=False, **kwargs)
  File "/share/workspace/dask/dask/base.py", line 436, in compute
    results = schedule(dsk, keys, **kwargs)
  File "/share/workspace/dask/dask/threaded.py", line 81, in get
    **kwargs
  File "/share/workspace/dask/dask/local.py", line 486, in get_async
    raise_exception(exc, tb)
  File "/share/workspace/dask/dask/local.py", line 316, in reraise
    raise exc
  File "/share/workspace/dask/dask/local.py", line 222, in execute_task
    result = _execute_task(task, data)
  File "/share/workspace/dask/dask/core.py", line 119, in _execute_task
    return func(*args2)
  File "/share/workspace/dask/dask/array/core.py", line 104, in getter
    c = a[b]
TypeError: 'coo_matrix' object is not subscriptable
```

As expected, the CuPy `csc_matrix` yields the same exception as CSR:

```
>>> a = cp.sparse.random(100, 100, format='csc')
>>> b = dask.array.from_array(a)
>>> b.compute()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/share/workspace/dask/dask/base.py", line 165, in compute
    (result,) = compute(self, traverse=False, **kwargs)
  File "/share/workspace/dask/dask/base.py", line 436, in compute
    results = schedule(dsk, keys, **kwargs)
  File "/share/workspace/dask/dask/threaded.py", line 81, in get
    **kwargs
  File "/share/workspace/dask/dask/local.py", line 486, in get_async
    raise_exception(exc, tb)
  File "/share/workspace/dask/dask/local.py", line 316, in reraise
    raise exc
  File "/share/workspace/dask/dask/local.py", line 222, in execute_task
    result = _execute_task(task, data)
  File "/share/workspace/dask/dask/core.py", line 119, in _execute_task
    return func(*args2)
  File "/share/workspace/dask/dask/array/core.py", line 104, in getter
    c = a[b]
  File "/share/workspace/cupy/cupyx/scipy/sparse/compressed.py", line 509, in __getitem__
    raise ValueError('unsupported indexing')
ValueError: unsupported indexing
```


Discussion:
Just to simplify this a bit, the issue can be reproduced using CuPy 7.3.0 alone.

```python
In [1]: import cupy as cp                                                       

In [2]: a = cp.sparse.random(100, 100, format='csr')                            

In [3]: a[0:100, 0:100]                                                         
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-3-a8ed163789e5> in <module>
----> 1 a[0:100, 0:100]

/datasets/jkirkham/miniconda/envs/rapids13dev/lib/python3.7/site-packages/cupyx/scipy/sparse/compressed.py in __getitem__(self, slices)
    247                 return self._get_major_slice(major)
    248 
--> 249         raise ValueError('unsupported indexing')
    250 
    251     def _get_single(self, major, minor):

ValueError: unsupported indexing
```
Should add indexing on the first axis works. Indexing on the second axis is problematic.

```python
In [1]: import cupy as cp                                                       

In [2]: a = cp.sparse.random(100, 100, format='csr')                            

In [3]: a[0:100]                                                                
Out[3]: <cupyx.scipy.sparse.csr.csr_matrix at 0x7fc7fb09cd10>

In [4]: a[:, 0:100]                                                             
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-4-fb183ec4eed3> in <module>
----> 1 a[:, 0:100]

/datasets/jkirkham/miniconda/envs/rapids13dev/lib/python3.7/site-packages/cupyx/scipy/sparse/compressed.py in __getitem__(self, slices)
    247                 return self._get_major_slice(major)
    248 
--> 249         raise ValueError('unsupported indexing')
    250 
    251     def _get_single(self, major, minor):

ValueError: unsupported indexing
```
@asi1024 as per our discussion during the PFN-NVIDIA meeting, would you be able to tell us here if there's any reason for this feature not to be allowed in CuPy? It seems that this is just something that nobody has had a chance to work on so far, but it would be great to have some confirmation.
I am currently a grad student working on my thesis. I am running into this same issue where the error is saying that the CuPy sparse array isn't subscript-able. I am wanting to show interest in working on this if it is appropriate. I haven't worked on any open source project, especially something like this. If you think this is something I could help out with, please point me in the right direction, if not that is fine too. 
@pentschev There is no reason for this indexing not to be supported.
Current implementation raises a `ValueError`, but it should raise `NotImplementedError`. Sorry for being confusing.
@crumleyc, we are needing this feature in Dask and RAPIDS currently. I think it would be great if you worked on this issue! Ideally, this change will be included in the 8.0.0 release.

As a first step, it will be helpful to familiarize yourself with Cupy’s [contribution guide](https://docs-cupy.chainer.org/en/stable/contribution.html), if you have not done so already. 

I tend to create a [pull request](https://docs-cupy.chainer.org/en/stable/contribution.html) early, so that I can solicit feedback as I work through changes. An additional benefit to this is that I am able to reference my changes directly when asking for help.
I have verified that the current implementation of #3486 fixes this issue. 

Resolving patch:
diff --git a/cupyx/scipy/sparse/_index.py b/cupyx/scipy/sparse/_index.py
--- a/cupyx/scipy/sparse/_index.py
+++ b/cupyx/scipy/sparse/_index.py
@@ -286,8 +286,6 @@ def _csr_row_slice(start_maj, step_maj, Ap, Aj, Ax, Bp):
         Bx : data array of output sparse matrix
     """
 
-    assert step_maj not in (-1, 0, 1)
-
     in_rows = cupy.arange(start_maj, start_maj + (Bp.size - 1) * step_maj,
                           step_maj, dtype=Bp.dtype)
     offsetsB = Ap[in_rows] - Bp[:-1]
@@ -301,6 +299,49 @@ def _csr_row_slice(start_maj, step_maj, Ap, Aj, Ax, Bp):
     return Bj, Bx
 
 
+def _csr_sample_values(n_row, n_col,
+                       Ap, Aj, Ax,
+                       Bi, Bj):
+    """Populate data array for a set of rows and columns
+    Args
+        n_row : total number of rows in input array
+        n_col : total number of columns in input array
+        Ap : indptr array for input sparse matrix
+        Aj : indices array for input sparse matrix
+        Ax : data array for input sparse matrix
+        Bi : array of rows to extract from input sparse matrix
+        Bj : array of columns to extract from input sparse matrix
+    Returns
+        Bx : data array for output sparse matrix
+    """
+
+    Bi[Bi < 0] += n_row
+    Bj[Bj < 0] += n_col
+
+    Bx = cupy.empty(Bi.size, dtype=Ax.dtype)
+    _csr_sample_values_kern(n_row, n_col,
+                            Ap, Aj, Ax,
+                            Bi, Bj, Bx, size=Bi.size)
+
+    return Bx
+
+
+_csr_sample_values_kern = core.ElementwiseKernel(
+    '''I n_row, I n_col, raw I Ap, raw I Aj, raw T Ax, raw I Bi, raw I Bj''',
+    'raw T Bx', '''
+    const I j = Bi[i]; // sample row
+    const I k = Bj[i]; // sample column
+    const I row_start = Ap[j];
+    const I row_end   = Ap[j+1];
+    T x = 0;
+    for(I jj = row_start; jj < row_end; jj++) {
+        if (Aj[jj] == k)
+            x += Ax[jj];
+    }
+    Bx[i] = x;
+''', 'csr_sample_values_kern', no_return=True)
+
+
 class IndexMixin(object):
     """
     This class provides common dispatching and validation logic for indexing.
diff --git a/cupyx/scipy/sparse/compressed.py b/cupyx/scipy/sparse/compressed.py
--- a/cupyx/scipy/sparse/compressed.py
+++ b/cupyx/scipy/sparse/compressed.py
@@ -499,7 +499,20 @@ def _get_sliceXslice(self, row, col):
 
     def _get_arrayXarray(self, row, col):
         # inner indexing
-        raise NotImplementedError()
+        idx_dtype = self.indices.dtype
+        M, N = self._swap(*self.shape)
+        major, minor = self._swap(row, col)
+        major = cupy.asarray(major, dtype=idx_dtype)
+        minor = cupy.asarray(minor, dtype=idx_dtype)
+
+        val = _index._csr_sample_values(
+            M, N, self.indptr, self.indices, self.data,
+            major.ravel(), minor.ravel())
+
+        if major.ndim == 1:
+            # Scipy returns `matrix` here
+            return cupy.expand_dims(val, 0)
+        return self.__class__(val.reshape(major.shape))
 
     def _get_columnXarray(self, row, col):
         # outer indexing


Test patch:
diff --git a/tests/cupyx_tests/scipy_tests/sparse_tests/test_index.py b/tests/cupyx_tests/scipy_tests/sparse_tests/test_index.py
--- a/tests/cupyx_tests/scipy_tests/sparse_tests/test_index.py
+++ b/tests/cupyx_tests/scipy_tests/sparse_tests/test_index.py
@@ -50,11 +50,11 @@ def _run(self, maj, min=None, flip_for_csc=True,
 
         if min is not None:
 
-            expected = expected[maj_h, min_h]
             actual = a[maj, min]
+            expected = expected[maj_h, min_h]
         else:
-            expected = expected[maj_h]
             actual = a[maj]
+            expected = expected[maj_h]
 
         if compare_dense:
             actual = actual.todense()
@@ -71,7 +71,7 @@ def _run(self, maj, min=None, flip_for_csc=True,
                 actual.data, expected.data)
         else:
             testing.assert_array_equal(
-                actual, expected)
+                actual, numpy.asarray(expected))
 
     @staticmethod
     def _get_index_combos(idx):
@@ -102,13 +102,6 @@ def test_major_scalar(self):
 
     def test_major_slice_minor_slice(self):
         self._run(slice(1, 5), slice(1, 5))
-
-    def test_major_slice_minor_all(self):
-        self._run(slice(1, 5), slice(None))
-        self._run(slice(5, 1), slice(None))
-
-    def test_major_slice_with_step(self):
-
         self._run(slice(1, 20, 2), slice(1, 5, 1))
         self._run(slice(20, 1, 2), slice(1, 5, 1))
         self._run(slice(1, 15, 2), slice(1, 5, 1))
@@ -117,6 +110,15 @@ def test_major_slice_with_step(self):
         self._run(slice(20, 1, 5), slice(None))
         self._run(slice(1, 20, 5), slice(None))
 
+    def test_major_slice_minor_all(self):
+        self._run(slice(1, 5), slice(None))
+        self._run(slice(5, 1), slice(None))
+
+    def test_major_slice_minor_scalar(self):
+        self._run(slice(1, 5), 5)
+        self._run(slice(5, 1), 5)
+        self._run(slice(5, 1, -1), 5)
+
     def test_major_scalar_minor_slice(self):
         self._run(5, slice(1, 5))
         self._run(numpy.array(5), slice(1, 5))
@@ -140,6 +142,15 @@ def test_major_all_minor_slice(self):
     def test_major_all_minor_all(self):
         self._run(slice(None), slice(None))
 
+    def test_ellipsis(self):
+        self._run(Ellipsis, flip_for_csc=False)
+        self._run(Ellipsis, 1, flip_for_csc=False)
+        self._run(1, Ellipsis, flip_for_csc=False)
+        self._run(Ellipsis, slice(None), flip_for_csc=False)
+        self._run(slice(None), Ellipsis, flip_for_csc=False)
+        self._run(Ellipsis, slice(1, None), flip_for_csc=False)
+        self._run(slice(1, None), Ellipsis, flip_for_csc=False)
+
     # Major Indexing
 
     def test_major_bool_fancy(self):
@@ -204,6 +215,10 @@ def test_major_slice_minor_bool(self):
                   compare_dense=True)
 
     def test_major_all_minor_fancy(self):
+
+        self._run(slice(None), [1, 5, 2, 3, 4, 5, 4, 1, 5])
+        self._run(slice(None), [0, 3, 4, 1, 1, 5, 5, 2, 3, 4, 5, 4, 1, 5])
+
         self._run(slice(None), [1, 5, 4, 5, 2, 4, 1])
 
         for idx in self._get_index_combos([1, 5, 4, 5, 2, 4, 1]):
@@ -223,14 +238,32 @@ def test_major_scalar_minor_fancy(self):
         for idx in self._get_index_combos([1, 5, 4, 1, 2]):
             self._run(5, idx)
 
-    def test_ellipsis(self):
-        self._run(Ellipsis, flip_for_csc=False)
-        self._run(Ellipsis, 1, flip_for_csc=False)
-        self._run(1, Ellipsis, flip_for_csc=False)
-        self._run(Ellipsis, slice(None), flip_for_csc=False)
-        self._run(slice(None), Ellipsis, flip_for_csc=False)
-        self._run(Ellipsis, slice(1, None), flip_for_csc=False)
-        self._run(slice(1, None), Ellipsis, flip_for_csc=False)
+    # Inner Indexing
+
+    def test_major_fancy_minor_fancy(self):
+
+        for idx in self._get_index_combos([1, 5, 4]):
+            self._run(idx, idx)
+
+        self._run([1, 5, 4], [1, 5, 4])
+
+        maj = self._get_index_combos([2, 0, 10, 0, 2])
+        min = self._get_index_combos([9, 2, 1, 0, 2])
+
+        for (idx1, idx2) in zip(maj, min):
+            self._run(idx1, idx2)
+
+        self._run([2, 0, 10, 0], [9, 2, 1, 0])
+
+        maj = self._get_index_combos([2, 0, 2])
+        min = self._get_index_combos([2, 1, 1])
+
+        for (idx1, idx2) in zip(maj, min):
+            self._run(idx1, idx2)
+
+        self._run([2, 0, 2], [2, 1, 2])
+
+    # Bad Indexing
 
     def test_bad_indexing(self):
         with pytest.raises(IndexError):
@@ -241,3 +274,6 @@ def test_bad_indexing(self):
 
         with pytest.raises(ValueError):
             self._run([1, 2, 3], [1, 2, 3, 4])
+
+        with pytest.raises(IndexError):
+            self._run([[0, 0], [1, 1]])