Commit 7c891506 authored by bwdeng20's avatar bwdeng20
Browse files

test metis; ignore pycharm .idea/

parent 363c29ad
......@@ -7,11 +7,16 @@ from .utils import devices
@pytest.mark.parametrize('device', devices)
def test_metis(device):
mat = SparseTensor.from_dense(torch.randn((6, 6), device=device))
mat, partptr, perm = mat.partition(num_parts=2, recursive=False)
weighted_mat = SparseTensor.from_dense(torch.randn((6, 6), device=device))
mat, partptr, perm = weighted_mat.partition(num_parts=2, recursive=False)
assert partptr.numel() == 3
assert perm.numel() == 6
mat, partptr, perm = mat.partition(num_parts=2, recursive=True)
mat, partptr, perm = weighted_mat.partition(num_parts=2, recursive=True)
assert partptr.numel() == 3
assert perm.numel() == 6
unweighted_mat = SparseTensor.from_dense(torch.ones((6, 6), device=device))
mat, partptr, perm = unweighted_mat.partition(num_parts=2, recursive=True)
assert partptr.numel() == 3
assert perm.numel() == 6
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment