| 24 | |
| 25 | |
| 26 | def rand_coo(shape, nnz, dev, nz_dim=None): |
| 27 | # Create a sparse matrix without duplicate entries. |
| 28 | nnzid = np.random.choice(shape[0] * shape[1], nnz, replace=False) |
| 29 | nnzid = torch.tensor(nnzid, device=dev).long() |
| 30 | row = torch.div(nnzid, shape[1], rounding_mode="floor") |
| 31 | col = nnzid % shape[1] |
| 32 | if nz_dim is None: |
| 33 | val = torch.randn(nnz, device=dev, requires_grad=True) |
| 34 | else: |
| 35 | val = torch.randn(nnz, nz_dim, device=dev, requires_grad=True) |
| 36 | indices = torch.stack([row, col]) |
| 37 | indices = rand_stride(indices) |
| 38 | val = rand_stride(val) |
| 39 | return spmatrix(indices, val, shape) |
| 40 | |
| 41 | |
| 42 | def rand_csr(shape, nnz, dev, nz_dim=None): |