opengt.head

class opengt.head.graphormer_graph.GraphormerHead(dim_in, dim_out)

Bases: Module

Graphormer prediction head for graph prediction tasks.

Parameters:

• dim_in (int) – Input dimension.

• dim_out (int) – Output dimension. For binary prediction, dim_out=1.

Input:

batch.x (torch.Tensor): Node features. batch.y (torch.Tensor): Graph labels. batch.batch (torch.Tensor): Batch indices.

Output:

pred (torch.Tensor): Predicted graph labels. true (torch.Tensor): True graph labels.

class opengt.head.inductive_edge.GNNInductiveEdgeHead(dim_in, dim_out)

Bases: Module

GNN prediction head for inductive edge/link prediction tasks.

Implementation adapted from the transductive GraphGym’s GNNEdgeHead.

Parameters:

• dim_in (int) – Input dimension

• dim_out (int) – Output dimension. For binary prediction, dim_out=1.

Input:

batch.x (torch.Tensor): Node features. batch.edge_label (torch.Tensor): Edge labels.

Output:

pred (torch.Tensor): Predicted edge labels. true (torch.Tensor): True edge labels.

class opengt.head.inductive_node.GNNInductiveNodeHead(dim_in, dim_out)

Bases: Module

GNN prediction head for inductive node prediction tasks.

Parameters:

• dim_in (int) – Input dimension

• dim_out (int) – Output dimension. For binary prediction, dim_out=1.

Input:

batch.x (torch.Tensor): Node features. batch.y (torch.Tensor): Node labels.

Output:

pred (torch.Tensor): Predicted node labels. true (torch.Tensor): True node labels.

class opengt.head.infer_links.InferLinksHead(dim_in, dim_out)

Bases: Module

InferLinks prediction head for graph prediction tasks.

Parameters:

• dim_in (int) – Input dimension.

• dim_out (int) – Output dimension. For binary prediction, dim_out=1.

Input:

batch.x (torch.Tensor): Node features. batch.y (torch.Tensor): Edge labels. batch.complete_edge_index (torch.Tensor): Edge indices for the complete graph.

Output:

pred (torch.Tensor): Predicted edge labels. true (torch.Tensor): True edge labels.

class opengt.head.mlp_mixer_graph.MLPMixerGraphHead(dim_in, dim_out, L=2)

Bases: Module

Graph MLP Mixer prediction head for graph prediction tasks.

Note that this head is specially designed for Graph MLP Mixer (without pooling layer). Cannot work on other models.

Parameters:

• dim_in (int) – Input dimension.

• dim_out (int) – Output dimension. For binary prediction, dim_out=1.

• L (int) – Number of hidden layers.

Input:

batch.x (torch.Tensor): Graph embedding. batch.y (torch.Tensor): Graph labels.

Output:

pred (torch.Tensor): Predicted graph labels. true (torch.Tensor): True graph labels.

class opengt.head.ogb_code_graph.OGBCodeGraphHead(dim_in, dim_out, L=1)

Bases: Module

Sequence prediction head for ogbg-code2 graph-level prediction tasks.

Parameters:

• dim_in (int) – Input dimension.

• dim_out (int) – IGNORED, kept for GraphGym framework compatibility

• L (int) – Number of hidden layers.

class opengt.head.san_graph.SANGraphHead(dim_in, dim_out, L=2)

Bases: Module

SAN prediction head for graph prediction tasks.

Parameters:

• dim_in (int) – Input dimension.

• dim_out (int) – Output dimension. For binary prediction, dim_out=1.

• L (int) – Number of hidden layers.

Input:

batch.x (torch.Tensor): Node features. batch.y (torch.Tensor): Graph labels. batch.batch (torch.Tensor): Batch indices.

Output:

pred (torch.Tensor): Predicted graph labels. true (torch.Tensor): True graph labels.