| 38 | @pytest.mark.parametrize("col", [(0, 1, 2, 2), (1, 3, 3, 4)]) |
| 39 | @pytest.mark.parametrize("shape", [None, (5, 5), (5, 6)]) |
| 40 | def test_from_coo(dense_dim, row, col, shape): |
| 41 | val_shape = (len(row),) |
| 42 | if dense_dim is not None: |
| 43 | val_shape += (dense_dim,) |
| 44 | ctx = F.ctx() |
| 45 | val = torch.randn(val_shape).to(ctx) |
| 46 | row = torch.tensor(row).to(ctx) |
| 47 | col = torch.tensor(col).to(ctx) |
| 48 | mat = from_coo(row, col, val, shape) |
| 49 | |
| 50 | if shape is None: |
| 51 | shape = (torch.max(row).item() + 1, torch.max(col).item() + 1) |
| 52 | |
| 53 | mat_row, mat_col = mat.coo() |
| 54 | mat_val = mat.val |
| 55 | |
| 56 | assert mat.shape == shape |
| 57 | assert mat.nnz == row.numel() |
| 58 | assert mat.dtype == val.dtype |
| 59 | assert torch.allclose(mat_val, val) |
| 60 | assert torch.allclose(mat_row, row) |
| 61 | assert torch.allclose(mat_col, col) |
| 62 | |
| 63 | |
| 64 | @pytest.mark.parametrize("dense_dim", [None, 4]) |