Function test_gine_layer

tests/test_layers.py:15–21

()

def test_gine_layer():
    layer = GINELayer()
    x = torch.eye(4)
    edge_index = (torch.tensor([0, 0, 0, 1, 1, 2]), torch.tensor([1, 2, 3, 2, 3, 3]))
    graph = Graph(x=x, edge_index=edge_index, edge_attr=torch.randn(6, 4))
    x = layer(graph, x)
    assert tuple(x.shape) == (4, 4)


if __name__ == "__main__":