1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
|
import torch
from torch_geometric.testing import is_full_test
from torch_geometric.utils import normalized_cut
def test_normalized_cut():
row = torch.tensor([0, 1, 1, 1, 2, 2, 3, 3, 4, 4])
col = torch.tensor([1, 0, 2, 3, 1, 4, 1, 4, 2, 3])
edge_attr = torch.tensor(
[3.0, 3.0, 6.0, 3.0, 6.0, 1.0, 3.0, 2.0, 1.0, 2.0])
expected = torch.tensor([4.0, 4.0, 5.0, 2.5, 5.0, 1.0, 2.5, 2.0, 1.0, 2.0])
out = normalized_cut(torch.stack([row, col], dim=0), edge_attr)
assert torch.allclose(out, expected)
if is_full_test():
jit = torch.jit.script(normalized_cut)
out = jit(torch.stack([row, col], dim=0), edge_attr)
assert torch.allclose(out, expected)
|
