Skip to content

GitLab

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

test metis; ignore pycharm .idea/

parent 363c29ad
Show whitespace changes
Inline Side-by-side
Showing with 8 additions and 3 deletions
test/test_metis.py
View file @ 7c891506
...@@ -7,11 +7,16 @@ from .utils import devices ...@@ -7,11 +7,16 @@ from .utils import devices
@pytest.mark.parametrize('device', devices) @pytest.mark.parametrize('device', devices)
def test_metis(device): def test_metis(device):
mat = SparseTensor.from_dense(torch.randn((6, 6), device=device)) weighted_mat = SparseTensor.from_dense(torch.randn((6, 6), device=device))
mat, partptr, perm = mat.partition(num_parts=2, recursive=False) mat, partptr, perm = weighted_mat.partition(num_parts=2, recursive=False)
assert partptr.numel() == 3 assert partptr.numel() == 3
assert perm.numel() == 6 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 partptr.numel() == 3
assert perm.numel() == 6 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