(shape, nnz, dev, nz_dim=None)
| 84 | |
| 85 | |
| 86 | def rand_diag(shape, nnz, dev, nz_dim=None): |
| 87 | nnz = min(shape) |
| 88 | if nz_dim is None: |
| 89 | val = torch.randn(nnz, device=dev, requires_grad=True) |
| 90 | else: |
| 91 | val = torch.randn(nnz, nz_dim, device=dev, requires_grad=True) |
| 92 | return diag(val, shape) |
| 93 | |
| 94 | |
| 95 | def rand_coo_uncoalesced(shape, nnz, dev): |