Skip to content

Commit

Permalink
extend test coverage.
Browse files Browse the repository at this point in the history
  • Loading branch information
mfbalin committed Oct 11, 2024
1 parent 20a35a7 commit c3e8e89
Showing 1 changed file with 4 additions and 3 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -37,13 +37,14 @@ def test_rank_sort_and_unique_and_compact(dtype, rank):
assert_equal(offsets1, offsets2)
assert offsets1.is_pinned() and offsets2.is_pinned()

res3 = torch.ops.graphbolt.rank_sort(nodes_list1, rank, WORLD_SIZE)
# Test with the reverse order of ntypes. See if results are equivalent.
res3 = torch.ops.graphbolt.rank_sort(nodes_list1[::-1], rank, WORLD_SIZE)

# This function is deterministic. Call with identical arguments and check.

Check warning on line 43 in tests/python/pytorch/graphbolt/impl/test_cooperative_minibatching_utils.py

View workflow job for this annotation

GitHub Actions / lintrunner

UFMT format

Run `lintrunner -a` to apply this patch.
for (nodes1, idx1, offsets1), (nodes3, idx3, offsets3) in zip(res1, res3):
for (nodes1, idx1, offsets1), (nodes3, idx3, offsets3) in zip(res1, reversed(res3)):
assert_equal(nodes1, nodes3)
assert_equal(idx1, idx3)
assert_equal(offsets1, offsets3)
assert_equal(offsets1.diff(), offsets3.diff())

# The dependency on the rank argument is simply a permutation.
res4 = torch.ops.graphbolt.rank_sort(nodes_list1, 0, WORLD_SIZE)
Expand Down

0 comments on commit c3e8e89

Please sign in to comment.