Fast CSR matvec with dense vector skipping output allocation. The result is added to the specificed output array, so the output should be manually zeroed prior to calling this routine, if necessary.
(A_csr, x_vec, out_vec)
| 254 | |
| 255 | |
| 256 | def csr_matvecs(A_csr, x_vec, out_vec): |
| 257 | """ |
| 258 | Fast CSR matvec with dense vector skipping output allocation. The result is |
| 259 | added to the specificed output array, so the output should be manually |
| 260 | zeroed prior to calling this routine, if necessary. |
| 261 | """ |
| 262 | # Check format but don't convert |
| 263 | if A_csr.format != "csr": |
| 264 | raise ValueError("Matrix must be in CSR format.") |
| 265 | # Check shapes |
| 266 | M, N = A_csr.shape |
| 267 | if x_vec.ndim != 2 or out_vec.ndim != 2: |
| 268 | raise ValueError("Only matrices allowed for input and output.") |
| 269 | n, kx = x_vec.shape |
| 270 | m, ko = out_vec.shape |
| 271 | if M != m or N != n: |
| 272 | raise ValueError(f"Matrix shape {(M,N)} does not match input {(n,)} and output {(m,)} shapes.") |
| 273 | if kx != ko: |
| 274 | raise ValueError("Output size does not match input size.") |
| 275 | # Apply matvecs |
| 276 | if SPLIT_CSR_MATVECS: |
| 277 | for k in range(kx): |
| 278 | _sparsetools.csr_matvec(M, N, A_csr.indptr, A_csr.indices, A_csr.data, x_vec[:,k], out_vec[:,k]) |
| 279 | else: |
| 280 | _sparsetools.csr_matvecs(M, N, kx, A_csr.indptr, A_csr.indices, A_csr.data, x_vec, out_vec) |
| 281 | return out_vec |
| 282 | |
| 283 | |
| 284 | def add_sparse(A, B): |