(A)
| 24 | """ |
| 25 | |
| 26 | def _mat_gather(A): |
| 27 | nr = A.index_map(0).size_local |
| 28 | bs0 = A.block_size[0] |
| 29 | bs1 = A.block_size[1] |
| 30 | nbs = bs0 * bs1 |
| 31 | gatheredvals = np.concatenate( |
| 32 | MPI.COMM_WORLD.allgather(A.data[: nbs * A.indptr[nr]]) |
| 33 | ).reshape(-1, bs0, bs1) |
| 34 | gatheredptrs = MPI.COMM_WORLD.allgather(A.indptr[: nr + 1]) |
| 35 | cols = A.index_map(1).local_to_global(A.indices[: A.indptr[nr]]) |
| 36 | gatheredcols = np.concatenate(MPI.COMM_WORLD.allgather(cols)) |
| 37 | indptr = gatheredptrs[0] |
| 38 | for i in range(1, len(gatheredptrs)): |
| 39 | indptr = np.concatenate((indptr, (gatheredptrs[i][1:] + indptr[-1]))) |
| 40 | return bsr_matrix((gatheredvals, gatheredcols, indptr)) |
| 41 | |
| 42 | return _mat_gather |
| 43 |
nothing calls this directly
no test coverage detected