opengt.encoder

class opengt.encoder.ast_encoder.ASTEdgeEncoder(emb_dim)

Bases: Module

The Abstract Syntax Tree (AST) Edge Encoder used for ogbg-code2 dataset.

Edge attributes are generated by augment_edge function dynamically and are expected to be: edge_attr[:,0]: whether it is AST edge (0) for next-token edge (1) edge_attr[:,1]: whether it is original direction (0) or inverse direction (1)

Parameters:

emb_dim (int) – Output edge embedding dimension

Input:

batch.edge_attr (torch.Tensor): Default edge feature.

Output:

batch.edge_attr (torch.Tensor): emb_dim-dimensional vector

class opengt.encoder.ast_encoder.ASTNodeEncoder(emb_dim)

Bases: Module

The Abstract Syntax Tree (AST) Node Encoder used for ogbg-code2 dataset.

Parameters:

emb_dim (int) – Output node embedding dimension

Input:

batch.x (torch.Tensor): Default node feature. The first and second column represents node type and node attributes. batch.node_depth (torch.Tensor): The depth of the node in the AST.

Output:

batch.x (torch.Tensor): emb_dim-dimensional vector

opengt.encoder.composed_encoders.concat_node_encoders(encoder_classes, pe_enc_names)

A factory that creates a new Encoder class that concatenates functionality of the given list of two or three Encoder classes. First Encoder is expected to be a dataset-specific encoder, and the rest PE Encoders.

Parameters:

• encoder_classes – List of node encoder classes

• pe_enc_names – List of PE embedding Encoder names, used to query a dict with their desired PE embedding dims. That dict can only be created during the runtime, once the config is loaded.

Returns:

new node encoder class

class opengt.encoder.equivstable_laplace_pos_encoder.EquivStableLapPENodeEncoder(dim_emb)

Bases: Module

Equivariant and Stable Laplace Positional Embedding node encoder.

This encoder simply transforms the k-dim node LapPE to d-dim to be later used at the local GNN module as edge weights. Based on the approach proposed in paper https://openreview.net/pdf?id=e95i1IHcWj

Parameters:

dim_emb – Size of final node embedding

Input:

batch.EigVals (torch.Tensor): Eigenvalues of the graph Laplacian. batch.EigVecs (torch.Tensor): Eigenvectors of the graph Laplacian.

Output:

batch.pe_EquivStableLapPE (torch.Tensor): Positional encoding of the

graph Laplacian eigenvalues and eigenvectors.

class opengt.encoder.ER_edge_encoder.EREdgeEncoder(emb_dim, use_edge_attr=False, expand_edge_attr=False)

Bases: Module

Effective Resistance Edge Encoder

This encoder calculates the effective resistance edge embedding mentioned in https://arxiv.org/pdf/2206.11941

Parameters:

• emb_dim (int) – Size of final edge embedding

• use_edge_attr (bool) – If True, the edge attributes are prepended to the effective resistance embedding.

• expand_edge_attr (bool) – If True, the edge attributes are linearly projected to (dim_in - emb_dim) size.

class opengt.encoder.ER_node_encoder.ERNodeEncoder(dim_emb, expand_x=True)

Bases: Module

Effective Resistance Node Encoder

ER of size dim_pe will get appended to each node feature vector. If expand_x set True, original node features will be first linearly projected to (dim_emb - dim_pe) size and the concatenated with ER.

Parameters:

• dim_emb (int) – Size of final node embedding

• expand_x (bool) – Expand node features x from dim_in to (dim_emb - dim_pe)

class opengt.encoder.exp_edge_fixer.ExpanderEdgeFixer(add_edge_index=False, num_virt_node=0)

Bases: Module

Gets the batch and sets new edge indices + global nodes. Creates expander_edge_index and expander_edge_attr for the batch instead of modifying the original edge_index and edge_attr. Also adds virtual nodes and edges for the batch and stores them in batch.virt_h, batch.virt_edge_index and batch.virt_edge_attr.

Used in Exphormer model.

Parameters:

• add_edge_index (bool) – If True, adds the edge index to the batch.

• num_virt_node (int) – Number of virtual nodes to add to the batch.

class opengt.encoder.feature_encoder.FeatureEncoder(dim_in)

Bases: Module

Encodes node and edge features. Receives the encoder type from the config file.

Parameters:

dim_in (int) – Input feature dimension

opengt.encoder.graphormer_encoder.add_graph_token(data, token)

Helper function to augment a batch of PyG graphs with a graph token each. Note that the token is automatically replicated to fit the batch.

Parameters:

• data (torch_geometric.data.Data) – A PyG data object holding a single graph

• token (torch.Tensor) – A tensor containing the graph token values

Returns:

The augmented data object.

Return type:

data (torch_geometric.data.Data)

class opengt.encoder.kernel_pos_encoder.ElstaticSENodeEncoder(dim_emb, expand_x=True)

Bases: KernelPENodeEncoder

Electrostatic interactions Structural Encoding node encoder.

class opengt.encoder.kernel_pos_encoder.HKdiagSENodeEncoder(dim_emb, expand_x=True)

Bases: KernelPENodeEncoder

Heat kernel (diagonal) Structural Encoding node encoder.

class opengt.encoder.kernel_pos_encoder.KernelPENodeEncoder(dim_emb, expand_x=True)

Bases: Module

Configurable kernel-based Positional Encoding node encoder.

The choice of which kernel-based statistics to use is configurable through setting of kernel_type. Based on this, the appropriate config is selected, and also the appropriate variable with precomputed kernel stats is then selected from PyG Data graphs in forward function. E.g., supported are ‘RWSE’, ‘HKdiagSE’, ‘ElstaticSE’.

PE of size dim_pe will get appended to each node feature vector. If expand_x set True, original node features will be first linearly projected to (dim_emb - dim_pe) size and the concatenated with PE.

Parameters:

• dim_emb (int) – Size of final node embedding

• expand_x (bool) – Expand node features x from dim_in to (dim_emb - dim_pe)

class opengt.encoder.kernel_pos_encoder.RWSENodeEncoder(dim_emb, expand_x=True)

Bases: KernelPENodeEncoder

Random Walk Structural Encoding node encoder.

class opengt.encoder.laplace_pos_encoder.LapPENodeEncoder(dim_emb, expand_x=True)

Bases: Module

Laplace Positional Embedding node encoder.

LapPE of size dim_pe will get appended to each node feature vector. If expand_x set True, original node features will be first linearly projected to (dim_emb - dim_pe) size and the concatenated with LapPE.

Parameters:

• dim_emb (int) – Size of final node embedding

• expand_x (bool) – Expand node features x from dim_in to (dim_emb - dim_pe)

class opengt.encoder.linear_edge_encoder.LinearEdgeEncoder(emb_dim)

Bases: Module

Linear edge encoder that applies a linear transformation to the edge features.

Parameters:

emb_dim (int) – The dimension of the output edge features.

class opengt.encoder.linear_node_encoder.LinearNodeEncoder(emb_dim)

Bases: Module

Linear node encoder that applies a linear transformation to the input features.

Parameters:

emb_dim (int) – The dimension of the output node features.

class opengt.encoder.ppa_encoder.PPAEdgeEncoder(emb_dim)

Bases: Module

Edge encoder for PPA that applies a linear transformation to the edge features.

Parameters:

emb_dim (int) – The dimension of the output edge features.

class opengt.encoder.ppa_encoder.PPANodeEncoder(emb_dim)

Bases: Module

Uniform input node embedding for PPA that has no node features.

Parameters:

emb_dim (int) – The dimension of the output node features.

The RRWP encoder for GRIT Adapted from https://github.com/LiamMa/GRIT

class opengt.encoder.rrwp_encoder.PadToFullGraphEdgeEncoder(**kwargs)

Bases: Module

Padding to Full Attention

class opengt.encoder.rrwp_encoder.RRWPLinearEdgeEncoder(emb_dim, out_dim, batchnorm=False, layernorm=False, use_bias=False, pad_to_full_graph=True, fill_value=0.0, add_node_attr_as_self_loop=False, overwrite_old_attr=False)

Bases: Module

Merge RRWP with given edge-attr and Zero-padding to all pairs of node FC_1(RRWP) + FC_2(edge-attr) - FC_2 given by the TypedictEncoder in same cases - Zero-padding for non-existing edges in fully-connected graph - (optional) add node-attr as the E_{i,i}’s attr

note: assuming node-attr and edge-attr is with the same dimension after Encoders

class opengt.encoder.rrwp_encoder.RRWPLinearEdgeMaskedEncoder(emb_dim, out_dim, batchnorm=False, layernorm=False, use_bias=False, fill_value=0.0, add_node_attr_as_self_loop=False, overwrite_old_attr=False, mask_index_name='edge_index')

Bases: Module

RRWP Linear + Sparse-Attention Masking

class opengt.encoder.rrwp_encoder.RRWPLinearNodeEncoder(emb_dim, out_dim, use_bias=False, batchnorm=False, layernorm=False, pe_name='rrwp')

Bases: Module

FC_1(RRWP) + FC_2 (Node-attr) note: FC_2 is given by the Typedict encoder of node-attr in some cases Parameters: num_classes (int) - the number of classes for the embedding mapping to learn

opengt.encoder.rrwp_encoder.full_edge_index(edge_index, batch=None)

Returns the Full batched sparse adjacency matrices given by edge indices. Returns batched sparse adjacency matrices with exactly those edges that are not in the input edge_index while ignoring self-loops. Implementation inspired by torch_geometric.utils.to_dense_adj :param edge_index: The edge indices. :type edge_index: torch.Tensor :param batch: Batch vector, which assigns each node to a specific example.

Returns:

Complementary edge index.

SignNet https://arxiv.org/abs/2202.13013 based on https://github.com/cptq/SignNet-BasisNet

class opengt.encoder.signnet_pos_encoder.GINDeepSigns(in_channels, hidden_channels, out_channels, num_layers, k, dim_pe, rho_num_layers, use_bn=False, use_ln=False, dropout=0.5, activation='relu')

Bases: Module

Sign invariant neural network with MLP aggregation. f(v1, …, vk) = rho(enc(v1) + enc(-v1), …, enc(vk) + enc(-vk))

class opengt.encoder.signnet_pos_encoder.MaskedGINDeepSigns(in_channels, hidden_channels, out_channels, num_layers, dim_pe, rho_num_layers, use_bn=False, use_ln=False, dropout=0.5, activation='relu')

Bases: Module

Sign invariant neural network with sum pooling and DeepSet. f(v1, …, vk) = rho(enc(v1) + enc(-v1), …, enc(vk) + enc(-vk))

class opengt.encoder.signnet_pos_encoder.SignNetNodeEncoder(dim_emb, expand_x=True)

Bases: Module

SignNet Positional Embedding node encoder. https://arxiv.org/abs/2202.13013 https://github.com/cptq/SignNet-BasisNet

Uses precomputated Laplacian eigen-decomposition, but instead of eigen-vector sign flipping + DeepSet/Transformer, computes the PE as: SignNetPE(v_1, … , v_k) = rho ( [phi(v_i) + phi(-v_i)]^k_i=1 ) where phi is GIN network applied to k first non-trivial eigenvectors, and rho is an MLP if k is a constant, but if all eigenvectors are used then rho is DeepSet with sum-pooling.

SignNetPE of size dim_pe will get appended to each node feature vector. If expand_x set True, original node features will be first linearly projected to (dim_emb - dim_pe) size and the concatenated with SignNetPE.

Parameters:

• dim_emb (int) – Size of final node embedding.

• expand_x (bool) – Expand node features x from dim_in to (dim_emb - dim_pe).

class opengt.encoder.sine_encoder.SineEncoder(hidden_dim=128)

Bases: Module

SineEncoder encodes the eigenvalues of the graph Laplacian into a higher dimensional space using sine and cosine functions. It is used in SpecFormer model.

Parameters:

hidden_dim (int) – The dimension of the hidden layer.

Input:

batch.x (torch.Tensor): Eigenvalues of the graph Laplacian.

Output:

batch.x (torch.Tensor): Encoded eigenvalues in a higher dimensional space.

The SPD encoder for GRIT Adapted from https://github.com/LiamMa/GRIT

class opengt.encoder.spd_encoder.SPDEdgeEncoder(in_dim, out_dim, batchnorm=False, layernorm=False, use_bias=False, pad_to_full_graph=True, pe_name='spd', pad_0th=False, overwrite_old_attr=False)

Bases: Module

Shortest-path-distance (SPD) Embedding Encoder

Parameters:

• in_dim (int) – The input dimension of the edge features.

• out_dim (int) – The output dimension of the edge features.

• batchnorm (bool) – Whether to apply batch normalization. Default: False.

• layernorm (bool) – Whether to apply layer normalization. Default: False.

• use_bias (bool) – Whether to use bias in the linear layer. Default: False.

• pad_to_full_graph (bool) – Whether to pad to a full graph. Default: True.

• pe_name (str) – The name of the positional encoding. Default: “spd”.

• pad_0th (bool) – Whether to pad the 0-th embedding with a zero vector. Default: False.

• overwrite_old_attr (bool) – Whether to overwrite old attributes. Default: False.

class opengt.encoder.type_dict_encoder.TypeDictEdgeEncoder(emb_dim)

Bases: Module

Generic edge encoder for datasets with edge features that consist of only one type dictionary thus require a single nn.Embedding layer.

Parameters:

emb_dim (int) – The dimension of the output edge features.

class opengt.encoder.type_dict_encoder.TypeDictNodeEncoder(emb_dim)

Bases: Module

Generic node encoder for datasets with node features that consist of only one type dictionary thus require a single nn.Embedding layer.

Parameters:

emb_dim (int) – The dimension of the output node features.

class opengt.encoder.voc_superpixels_encoder.VOCEdgeEncoder(emb_dim)

Bases: Module

Dataset specific edge encoder for VOCSuperpixels dataset. Applies a linear transformation to the input edge features.

Parameters:

emb_dim (int) – The dimension of the output edge features.

class opengt.encoder.voc_superpixels_encoder.VOCNodeEncoder(emb_dim)

Bases: Module

Dataset specific node encoder for VOCSuperpixels dataset. Applies a linear transformation to the input node features.

Parameters:

emb_dim (int) – The dimension of the output node features.

class opengt.encoder.wl_encoder.WLSENodeEncoder(dim_emb, expand_x=True)

Bases: Module

Encodes the Weisfeiler-Lehman subgraph type of each node in the graph.

Parameters:

• dim_emb (int) – The dimension of the embedding.

• expand_x (bool) – If True, expands the input node features by the embedding dimension. Default: True.