Add tests for cupy.cusparse.csrsm2
@@ -687,3 +687,71 @@ class TestErrorSpmm(unittest.TestCase):
c = cupy.array(self.b, order='f')
with self.assertRaises(ValueError):
cupy.cusparse.spmm(a, b, c=c)
+
+
+@testing.parameterize(*testing.product({
+ 'lower': [True, False],
+ 'unit_diag': [True, False],
+ 'transa': ['N', 'T', 'H'],
+ 'blocking': [True, False],
+ 'level_info': [True, False],
+ 'format': ['csr', 'csc'],
+ 'nrhs': [None, 1, 4],
+ 'order': ['C', 'F']
+}))
+@testing.with_requires('scipy')
+class TestCsrsm2(unittest.TestCase):
+
+ n = 6
+ alpha = 1.0
+ density = 0.75
+ _tol = {'f': 1e-5, 'd': 1e-12}
+
+ def _setup(self, dtype):
+ dtype = numpy.dtype(dtype)
+ self.tol = self._tol[dtype.char.lower()]
+
+ a_shape = (self.n, self.n)
+ a = testing.shaped_random(a_shape, numpy, dtype=dtype, scale=1)
+ a_mask = testing.shaped_random(a_shape, numpy, dtype='f', scale=1)
+ a[a_mask > self.density] = 0
+ a_diag = numpy.diag(numpy.ones((self.n,), dtype=dtype))
+ if self.unit_diag:
+ a[a_diag > 0] = 0
+ a = a + a_diag
+ cp_a = cupy.array(a)
+ if self.unit_diag:
+ cp_a[a_diag > 0] = 0.1 # any number except 0
+ if self.format == 'csr':
+ self.a = sparse.csr_matrix(cp_a)
+ elif self.format == 'csc':
+ self.a = sparse.csc_matrix(cp_a)
+
+ b_shape = (self.n,) if self.nrhs is None else (self.n, self.nrhs)
+ b = numpy.arange(1, numpy.prod(b_shape) + 1,
+ dtype=dtype).reshape(b_shape)
+ b = b.copy(order=self.order)
+ self.b = cupy.array(b, order=self.order)
+
+ if self.lower:
+ a = numpy.tril(a)
+ else:
+ a = numpy.triu(a)
+ if self.transa == 'T':
+ a = a.T
+ elif self.transa == 'H':
+ a = a.conj().T
+ self.ref_x = numpy.linalg.solve(a, self.alpha * b)
+
+ @testing.for_dtypes('fdFD')
+ def test_csrsm2(self, dtype):
+ if (self.format == 'csc' and numpy.dtype(dtype).char in 'FD' and
+ self.transa == 'H'):
+ raise unittest.SkipTest()
+ self._setup(dtype)
+ x = self.b.copy(order=self.order)
+ cusparse.csrsm2(self.a, x, alpha=self.alpha,
+ lower=self.lower, unit_diag=self.unit_diag,
+ transa=self.transa, blocking=self.blocking,
+ level_info=self.level_info)
+ testing.assert_allclose(x, self.ref_x, atol=self.tol, rtol=self.tol)
