(val_shape)
| 119 | |
| 120 | @pytest.mark.parametrize("val_shape", [(3), (3, 2)]) |
| 121 | def test_dense(val_shape): |
| 122 | ctx = F.ctx() |
| 123 | |
| 124 | row = torch.tensor([1, 1, 2]).to(ctx) |
| 125 | col = torch.tensor([2, 4, 3]).to(ctx) |
| 126 | val = torch.randn(val_shape).to(ctx) |
| 127 | A = from_coo(row, col, val) |
| 128 | A_dense = A.to_dense() |
| 129 | |
| 130 | shape = A.shape + val.shape[1:] |
| 131 | mat = torch.zeros(shape, device=ctx) |
| 132 | mat[row, col] = val |
| 133 | assert torch.allclose(A_dense, mat) |
| 134 | |
| 135 | |
| 136 | @pytest.mark.parametrize("dense_dim", [None, 4]) |