1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
|
import torch
from torch_geometric.data import Data, HeteroData
from torch_geometric.transforms import Constant
def test_constant():
assert str(Constant()) == 'Constant(value=1.0)'
x = torch.tensor([[-1, 0], [0, 0], [2, 0]], dtype=torch.float)
edge_index = torch.tensor([[0, 1], [1, 2]])
data = Data(edge_index=edge_index, num_nodes=3)
data = Constant()(data)
assert len(data) == 3
assert data.edge_index.tolist() == edge_index.tolist()
assert data.x.tolist() == [[1], [1], [1]]
assert data.num_nodes == 3
data = Data(edge_index=edge_index, x=x)
data = Constant()(data)
assert len(data) == 2
assert data.edge_index.tolist() == edge_index.tolist()
assert data.x.tolist() == [[-1, 0, 1], [0, 0, 1], [2, 0, 1]]
data = HeteroData()
data['v'].x = x
data = Constant()(data)
assert len(data) == 1
assert data['v'].x.tolist() == [[-1, 0, 1], [0, 0, 1], [2, 0, 1]]
data = HeteroData()
data['v'].x = x
data['w'].x = x
data = Constant(node_types='w')(data)
assert len(data) == 1
assert data['v'].x.tolist() == x.tolist()
assert data['w'].x.tolist() == [[-1, 0, 1], [0, 0, 1], [2, 0, 1]]
|
