test_graph_convolution.py 文件源码

def test_csr_matvec(self):
        import scipy.sparse
        x = np.random.randn(5)
        A = scipy.sparse.csr_matrix(
            np.random.randn(4, 5).reshape((4, 5)), dtype=x.dtype)
        y = A.dot(x)
        y_gpu = cupy.empty((4,), dtype=x.dtype)
        A_data = cuda.to_gpu(A.data)
        A_indices = cuda.to_gpu(A.indices)
        A_indptr = cuda.to_gpu(A.indptr)
        x_gpu = cuda.to_gpu(x)
        graph_convolution.csr_matvec(y.shape[0], A_data, A_indices, A_indptr,
                                     x_gpu, y_gpu)
        testing.assert_allclose(
            y, cuda.to_cpu(y_gpu), **self.check_forward_options)