1import pytest 2 3pytestmark = pytest.mark.gpu 4 5import dask.array as da 6from dask.array.utils import assert_eq 7 8cupy = pytest.importorskip("cupy") 9cupyx = pytest.importorskip("cupyx") 10 11 12def test_sparse_hstack_vstack_csr(): 13 pytest.importorskip("cupyx") 14 x = cupy.arange(24, dtype=cupy.float32).reshape(4, 6) 15 16 sp = da.from_array(x, chunks=(2, 3), asarray=False, fancy=False) 17 sp = sp.map_blocks(cupyx.scipy.sparse.csr_matrix, dtype=cupy.float32) 18 19 y = sp.compute() 20 21 assert cupyx.scipy.sparse.isspmatrix(y) 22 assert_eq(x, y.todense()) 23 24 25@pytest.mark.parametrize("axis", [0, 1]) 26def test_sparse_concatenate(axis): 27 pytest.importorskip("cupyx") 28 29 rs = da.random.RandomState(RandomState=cupy.random.RandomState) 30 meta = cupyx.scipy.sparse.csr_matrix((0, 0)) 31 32 xs = [] 33 ys = [] 34 for i in range(2): 35 x = rs.random((1000, 10), chunks=(100, 10)) 36 x[x < 0.9] = 0 37 xs.append(x) 38 ys.append(x.map_blocks(cupyx.scipy.sparse.csr_matrix, meta=meta)) 39 40 z = da.concatenate(ys, axis=axis) 41 z = z.compute() 42 43 if axis == 0: 44 sp_concatenate = cupyx.scipy.sparse.vstack 45 elif axis == 1: 46 sp_concatenate = cupyx.scipy.sparse.hstack 47 z_expected = sp_concatenate( 48 [cupyx.scipy.sparse.csr_matrix(e.compute()) for e in xs] 49 ) 50 51 assert (z.toarray() == z_expected.toarray()).all() 52 53 54@pytest.mark.parametrize("sp_format", ["csr", "csc"]) 55def test_sparse_dot(sp_format): 56 pytest.importorskip("cupyx") 57 58 if sp_format == "csr": 59 sp_matrix = cupyx.scipy.sparse.csr_matrix 60 elif sp_format == "csc": 61 sp_matrix = cupyx.scipy.sparse.csc_matrix 62 dtype = "f" 63 density = 0.3 64 x_shape, x_chunks = (4, 8), (2, 4) 65 y_shape, y_chunks = (8, 6), (4, 3) 66 x = cupy.random.random(x_shape, dtype=dtype) 67 y = cupy.random.random(y_shape, dtype=dtype) 68 x[x < 1 - density] = 0 69 y[y < 1 - density] = 0 70 z = x.dot(y) 71 72 da_x = da.from_array(x, chunks=x_chunks, asarray=False, fancy=False) 73 da_y = da.from_array(y, chunks=y_chunks, asarray=False, fancy=False) 74 da_x = da_x.map_blocks(sp_matrix, dtype=dtype) 75 da_y = da_y.map_blocks(sp_matrix, dtype=dtype) 76 da_z = da.dot(da_x, da_y).compute() 77 78 assert cupyx.scipy.sparse.isspmatrix(da_z) 79 assert_eq(z, da_z.todense()) 80