Test CUB spmv

leofang · web-flow · commit 17f2f4698b56 · 2020-06-11T03:17:06.000-04:00
style in line with the existing tests and #2598
diff --git a/tests/cupyx_tests/scipy_tests/sparse_tests/test_csr.py b/tests/cupyx_tests/scipy_tests/sparse_tests/test_csr.py
@@ -1524,3 +1524,26 @@ def test_getitem_slice_stop_too_small(self, xp, sp):
     @testing.numpy_cupy_allclose(sp_name='sp')
     def test_getitem_slice_stop_too_large(self, xp, sp):
         return _make(xp, sp, self.dtype)[None:4]
+
+
+@testing.parameterize(*testing.product({
+    'make_method': [
+        '_make', '_make_unordered', '_make_empty', '_make_duplicate',
+        '_make_shape'],
+    'dtype': [numpy.float32, numpy.float64, cupy.complex64, cupy.complex128],
+}))
+@testing.with_requires('scipy')
+@testing.gpu
+@unittest.skipIf(cupy.cuda.cub_enabled is False, 'The CUB module is not built')
+class TestCUBspmv(unittest.TestCase):
+    @property
+    def make(self):
+        return globals()[self.make_method]
+
+    @testing.numpy_cupy_allclose(sp_name='sp')
+    def test_mul_dense_vector(self, xp, sp):
+        assert cupy.cuda.cub_enabled
+
+        m = self.make(xp, sp, self.dtype)
+        x = xp.arange(4).astype(self.dtype)
+        return m * x
