(shape, nnz, dev, nz_dim=None)
| 40 | |
| 41 | |
| 42 | def rand_csr(shape, nnz, dev, nz_dim=None): |
| 43 | # Create a sparse matrix without duplicate entries. |
| 44 | nnzid = np.random.choice(shape[0] * shape[1], nnz, replace=False) |
| 45 | nnzid = torch.tensor(nnzid, device=dev).long() |
| 46 | row = torch.div(nnzid, shape[1], rounding_mode="floor") |
| 47 | col = nnzid % shape[1] |
| 48 | if nz_dim is None: |
| 49 | val = torch.randn(nnz, device=dev, requires_grad=True) |
| 50 | else: |
| 51 | val = torch.randn(nnz, nz_dim, device=dev, requires_grad=True) |
| 52 | indptr = torch.zeros(shape[0] + 1, device=dev, dtype=torch.int64) |
| 53 | for r in row.tolist(): |
| 54 | indptr[r + 1] += 1 |
| 55 | indptr = torch.cumsum(indptr, 0) |
| 56 | row_sorted, row_sorted_idx = torch.sort(row) |
| 57 | indices = col[row_sorted_idx] |
| 58 | indptr = rand_stride(indptr) |
| 59 | indices = rand_stride(indices) |
| 60 | val = rand_stride(val) |
| 61 | return from_csr(indptr, indices, val, shape=shape) |
| 62 | |
| 63 | |
| 64 | def rand_csc(shape, nnz, dev, nz_dim=None): |
nothing calls this directly
no test coverage detected