1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
|
import torch
from torch_geometric.data import Data
from torch_geometric.nn import GPSE
from torch_geometric.nn.models.gpse import IdentityHead
from torch_geometric.transforms import AddGPSE
num_nodes = 6
gpse_inner_dim = 512
def test_gpse():
x = torch.randn(num_nodes, 4)
edge_index = torch.tensor([[0, 1, 0, 4, 1, 4, 2, 3, 3, 5],
[1, 0, 4, 0, 4, 1, 3, 2, 5, 3]])
data = Data(x=x, edge_index=edge_index)
model = GPSE()
model.post_mp = IdentityHead()
transform = AddGPSE(model)
assert str(transform) == 'AddGPSE()'
out = transform(data)
assert out.pestat_GPSE.size() == (num_nodes, gpse_inner_dim)
|
