graphium.nn.pyg_layers¶
Implementations of GNN layers based on PyG in the library
Gated GCN Layer¶
graphium.nn.pyg_layers.gated_gcn_pyg
¶
Unit tests for the different layers of graphium/nn/pyg_layers/...
The layers are not thoroughly tested due to the difficulty of testing them
adapated from https://github.com/rampasek/GraphGPS/blob/main/graphgps/layer/gps_layer.py
GatedGCNPyg
¶
Bases: MessagePassing
, BaseGraphStructure
layer_inputs_edges: bool
property
¶
Return a boolean specifying if the layer type
uses edges as input or not.
It is different from layer_supports_edges
since a layer that
supports edges can decide to not use them.
Returns:
Name  Type  Description 

bool 
bool

Always 
layer_outputs_edges: bool
property
¶
Abstract method. Return a boolean specifying if the layer type
uses edges as input or not.
It is different from layer_supports_edges
since a layer that
supports edges can decide to not use them.
Returns:
Name  Type  Description 

bool 
bool

Always 
out_dim_factor: int
property
¶
Get the factor by which the output dimension is multiplied for the next layer.
For standard layers, this will return 1
.
But for others, such as GatLayer
, the output is the concatenation
of the outputs from each head, so the out_dim gets multiplied by
the number of heads, and this function should return the number
of heads.
Returns:
int:
Always ``1`` for the current class
__init__(in_dim, out_dim, in_dim_edges, out_dim_edges=None, activation='relu', dropout=0.0, normalization='none', eps=1e05, **kwargs)
¶
ResGatedGCN: Residual Gated Graph ConvNets An Experimental Study of Neural Networks for Variable Graphs (Xavier Bresson and Thomas Laurent, ICLR 2018) https://arxiv.org/pdf/1711.07553v2.pdf
Parameters:
in_dim:
Input feature dimensions of the layer
out_dim:
Output feature dimensions of the layer, and for the edges
in_dim_edges:
Input edgefeature dimensions of the layer
activation:
activation function to use in the layer
dropout:
The ratio of units to dropout. Must be between 0 and 1
normalization:
Normalization to use. Choices:
 "none" or `None`: No normalization
 "batch_norm": Batch normalization
 "layer_norm": Layer normalization
 `Callable`: Any callable function
eps:
Epsilon value for the normalization `sum(gate_weights * messages) / (sum(gate_weights) + eps)`,
where `gate_weights` are the weights of the gates and follow a sigmoid function.
aggregate(sigma_ij, index, Bx_j, Bx)
¶
aggregation function of the layer Parameters: sigma_ij: the output from message() function with dim [n_edges, out_dim] index: dim [n_edges] Bx_j: dim [n_edges, out_dim] Bx: dim [n_nodes, out_dim] Returns: out: dim [n_nodes, out_dim]
forward(batch)
¶
Forward pass the Gated GCN layer extract the following from the batch: x, node features with dim [n_nodes, in_dim] e, edge features with dim [n_edges, in_dim] edge_index with dim [2, n_edges]
Parameters:
Name  Type  Description  Default 

batch 
Union[Data, Batch]

pyg Batch graph to pass through the layer 
required 
Returns: batch: pyg Batch graph
layer_supports_edges()
¶
Return a boolean specifying if the layer type supports edges or not.
Returns:
Name  Type  Description 

bool 
bool

Always 
message(Dx_i, Ex_j, Ce)
¶
message function Parameters: Dx_i: tensor with dimension [n_edges, out_dim] Ex_j: tensor with dimension [n_edges, out_dim] Ce: tensor with dimension [n_edges, out_dim] Returns: sigma_ij: tensor with dimension [n_edges, out_dim]
update(aggr_out, Ax)
¶
update function of the layer Parameters: aggr_out: the output from aggregate() function with dim [n_nodes, out_dim] Ax: tensor with dim [n_nodes, out_dim] Returns: x: dim [n_nodes, out_dim] e_out: dim [n_edges, out_dim_edges]
GPS Layer¶
graphium.nn.pyg_layers.gps_pyg
¶
GPSLayerPyg
¶
Bases: BaseGraphModule
layer_inputs_edges: bool
property
¶
Return a boolean specifying if the layer type
uses edges as input or not.
It is different from layer_supports_edges
since a layer that
supports edges can decide to not use them.
Returns:
bool:
Always ``True`` for the current class
layer_outputs_edges: bool
property
¶
Abstract method. Return a boolean specifying if the layer type
uses edges as input or not.
It is different from layer_supports_edges
since a layer that
supports edges can decide to not use them.
Returns:
bool:
Always ``False`` for the current class
out_dim_factor: int
property
¶
Get the factor by which the output dimension is multiplied for the next layer.
For standard layers, this will return 1
.
But for others, such as GatLayer
, the output is the concatenation
of the outputs from each head, so the out_dim gets multiplied by
the number of heads, and this function should return the number
of heads.
Returns:
int:
Always ``1`` for the current class
__init__(in_dim, out_dim, in_dim_edges=None, out_dim_edges=None, activation='relu', dropout=0.0, node_residual=True, normalization='none', mpnn_type='pyg:gine', mpnn_kwargs=None, attn_type='fullattention', precision='32', biased_attention_key=None, attn_kwargs=None, droppath_rate_attn=0.0, droppath_rate_ffn=0.0, hidden_dim_scaling=4.0, **kwargs)
¶
GPS layer implementation in pyg adapated from https://github.com/rampasek/GraphGPS/blob/main/graphgps/layer/gps_layer.py GPS: Recipe for a General, Powerful, Scalable Graph Transformer Ladislav Rampášek, Mikhail Galkin, Vijay Prakash Dwivedi, Anh Tuan Luu, Guy Wolf, Dominique Beaini https://arxiv.org/abs/2205.12454
GPS++: An Optimised Hybrid MPNN/Transformer for Molecular Property Prediction Dominic Masters, Josef Dean, Kerstin Klaser, Zhiyi Li, Sam MaddrellMander, Adam Sanders, Hatem Helal, Deniz Beker, Ladislav Rampášek, Dominique Beaini https://arxiv.org/abs/2212.02229
Parameters:
in_dim:
Input node feature dimensions of the layer
out_dim:
Output node feature dimensions of the layer
in_dim:
Input edge feature dimensions of the layer
out_dim:
Output edge feature dimensions of the layer
in_dim_edges:
input edgefeature dimensions of the layer
out_dim_edges:
output edgefeature dimensions of the layer
activation:
activation function to use in the layer
dropout:
The ratio of units to dropout. Must be between 0 and 1
node_residual:
If node residual is used after on the gnn layer output
normalization:
Normalization to use. Choices:
 "none" or `None`: No normalization
 "batch_norm": Batch normalization
 "layer_norm": Layer normalization
 `Callable`: Any callable function
mpnn_type:
type of mpnn used, choose from "pyg:gin", "pyg:gine", "pyg:gatedgcn", "pyg:pnamsgpass" and "pyg:mpnnplus"
mpnn_kwargs:
kwargs for mpnn layer
attn_type:
type of attention used, choose from "fullattention" and "none"
attn_kwargs:
kwargs for attention layer
droppath_rate_attn:
stochastic depth drop rate for attention layer https://arxiv.org/abs/1603.09382
droppath_rate_ffn:
stochastic depth drop rate for ffn layer https://arxiv.org/abs/1603.09382
mpnn_type:
Type of MPNN layer to use. Choices specified in PYG_LAYERS_DICT
mpnn_kwargs:
Keyword arguments to pass to the MPNN layer
attn_type:
Type of attention layer to use. Choices specified in ATTENTION_LAYERS_DICT
biased_attention_key:
indicates if biased attention is used by specifying a key corresponding to the pyg attribute in the batch (processed by the gaussian kernel encoder)
default: None means biased attention is not used
attn_kwargs:
Keyword arguments to pass to the attention layer
forward(batch)
¶
forward function of the layer Parameters: batch: pyg Batch graphs to pass through the layer Returns: batch: pyg Batch graphs
layer_supports_edges()
¶
Return a boolean specifying if the layer type supports edges or not.
Returns:
supports_edges: bool
Always ``True`` for the current class
GIN and GINE Layers¶
graphium.nn.pyg_layers.gin_pyg
¶
GINConvPyg
¶
Bases: BaseGraphModule
layer_inputs_edges: bool
property
¶
Return a boolean specifying if the layer type
uses edges as input or not.
It is different from layer_supports_edges
since a layer that
supports edges can decide to not use them.
Returns:
bool:
Always ``False`` for the current class
layer_outputs_edges: bool
property
¶
Abstract method. Return a boolean specifying if the layer type
uses edges as input or not.
It is different from layer_supports_edges
since a layer that
supports edges can decide to not use them.
Returns:
bool:
Always ``False`` for the current class
out_dim_factor: int
property
¶
Get the factor by which the output dimension is multiplied for the next layer.
For standard layers, this will return 1
.
But for others, such as GatLayer
, the output is the concatenation
of the outputs from each head, so the out_dim gets multiplied by
the number of heads, and this function should return the number
of heads.
Returns:
int:
Always ``1`` for the current class
__init__(in_dim, out_dim, activation='relu', dropout=0.0, normalization='none', **kwargs)
¶
GIN: Graph Isomorphism Networks HOW POWERFUL ARE GRAPH NEURAL NETWORKS? (Keyulu Xu, Weihua Hu, Jure Leskovec and Stefanie Jegelka, ICLR 2019) https://arxiv.org/pdf/1810.00826.pdf
[!] code uses the pytorchgeometric implementation of GINConv
Parameters:
in_dim:
Input feature dimensions of the layer
out_dim:
Output feature dimensions of the layer
activation:
activation function to use in the layer
dropout:
The ratio of units to dropout. Must be between 0 and 1
normalization:
Normalization to use. Choices:
 "none" or `None`: No normalization
 "batch_norm": Batch normalization
 "layer_norm": Layer normalization
 `Callable`: Any callable function
init_eps :
Initial :math:`\epsilon` value, default: ``0``.
learn_eps :
If True, :math:`\epsilon` will be a learnable parameter.
forward(batch)
¶
forward function of the layer Parameters: batch: pyg Batch graphs to pass through the layer Returns: batch: pyg Batch graphs
layer_supports_edges()
¶
Return a boolean specifying if the layer type supports edges or not.
Returns:
supports_edges: bool
Always ``False`` for the current class
GINEConvPyg
¶
Bases: BaseGraphModule
layer_inputs_edges: bool
property
¶
Return a boolean specifying if the layer type
uses edges as input or not.
It is different from layer_supports_edges
since a layer that
supports edges can decide to not use them.
Returns:
bool:
Always ``True`` for the current class
layer_outputs_edges: bool
property
¶
Abstract method. Return a boolean specifying if the layer type
uses edges as input or not.
It is different from layer_supports_edges
since a layer that
supports edges can decide to not use them.
Returns:
bool:
Always ``False`` for the current class
out_dim_factor: int
property
¶
Get the factor by which the output dimension is multiplied for the next layer.
For standard layers, this will return 1
.
But for others, such as GatLayer
, the output is the concatenation
of the outputs from each head, so the out_dim gets multiplied by
the number of heads, and this function should return the number
of heads.
Returns:
int:
Always ``1`` for the current class
__init__(in_dim, out_dim, in_dim_edges=None, activation='relu', dropout=0.0, normalization='none', **kwargs)
¶
GINE: Graph Isomorphism Networks with Edges Strategies for Pretraining Graph Neural Networks Weihua Hu, Bowen Liu, Joseph Gomes, Marinka Zitnik, Percy Liang, Vijay Pande, Jure Leskovec https://arxiv.org/abs/1905.12265
[!] code uses the pytorchgeometric implementation of GINEConv
Parameters:
in_dim:
Input feature dimensions of the layer
out_dim:
Output feature dimensions of the layer
activation:
activation function to use in the layer
dropout:
The ratio of units to dropout. Must be between 0 and 1
normalization:
Normalization to use. Choices:
 "none" or `None`: No normalization
 "batch_norm": Batch normalization
 "layer_norm": Layer normalization
 `Callable`: Any callable function
init_eps :
Initial :math:`\epsilon` value, default: ``0``.
learn_eps :
If True, :math:`\epsilon` will be a learnable parameter.
forward(batch)
¶
forward function of the layer Parameters: batch: pyg Batch graphs to pass through the layer Returns: batch: pyg Batch graphs
layer_supports_edges()
¶
Return a boolean specifying if the layer type supports edges or not.
Returns:
supports_edges: bool
Always ``True`` for the current class
MPNN Layer¶
graphium.nn.pyg_layers.mpnn_pyg
¶
MPNNPlusPyg
¶
Bases: BaseGraphModule
layer_inputs_edges: bool
property
¶
Return a boolean specifying if the layer type
uses edges as input or not.
It is different from layer_supports_edges
since a layer that
supports edges can decide to not use them.
Returns:
bool:
Always ``True`` for the current class
layer_outputs_edges: bool
property
¶
Abstract method. Return a boolean specifying if the layer type
uses edges as input or not.
It is different from layer_supports_edges
since a layer that
supports edges can decide to not use them.
Returns:
bool:
Always ``True`` for the current class
out_dim_factor: int
property
¶
Get the factor by which the output dimension is multiplied for the next layer.
For standard layers, this will return 1
.
But for others, such as GatLayer
, the output is the concatenation
of the outputs from each head, so the out_dim gets multiplied by
the number of heads, and this function should return the number
of heads.
Returns:
int:
Always ``1`` for the current class
__init__(in_dim=64, out_dim=64, activation='gelu', dropout=0.3, normalization='layer_norm', gather_from='both', scatter_to='both', node_combine_method='concat', num_node_mlp=2, mlp_expansion_ratio=4, use_edges=True, in_dim_edges=32, out_dim_edges=32, aggregation_method=['sum'], num_edge_mlp=2, use_globals=True, edge_dropout_rate=0.0035, **kwargs)
¶
MPNNPlusPyg: InteractionNetwork layer witg edges and global feature, GPS++ type of GNN layer
GPS++: An Optimised Hybrid MPNN/Transformer for Molecular Property Prediction
Dominic Masters, Josef Dean, Kerstin Klaser, Zhiyi Li, Sam MaddrellMander, Adam Sanders,
Hatem Helal, Deniz Beker, Ladislav Rampášek, Dominique Beaini
https://arxiv.org/abs/2212.02229
Parameters:
in_dim:
Input feature dimensions of the nodes
out_dim:
Output feature dimensions of the nodes
activation:
Activation function to use in the edge and node model
dropout:
The ratio of units to dropout at the end within apply_norm_activation_dropout.
normalization:
Normalization to use with the edge, node models and at the end
within apply_norm_activation_dropout. Choices:
 "none" or `None`: No normalization
 "batch_norm": Batch normalization
 "layer_norm": Layer normalization
 `Callable`: Any callable function
gather_from:
The method to gather features from. Could choose from:
"senders", "receivers" and "both".
scatter_to:
The method to scatter features to. Could choose from:
"senders", "receivers" and "both".
node_combine_method:
The method to combine the node features, Could choose from:
"sum" and "concat".
aggregation_method:
Methods for aggregating (scatter) messages built from node and edge features.
Provide a list of `Aggregation` or strings.
supported strings are:
 "sum" / "add" (Default)
 "mean"
 "max"
 "min"
 "softmax"
 "median"
 "std"
 "var"
num_node_mlp:
Number of mlp layer used for node model
mlp_expansion_ratio:
Expansion ratio for node and edge mlp
use_edges:
If edge features are used
in_dim_edges:
Input feature dimensions of the edges
out_dim_edges:
Output feature dimensions of the edges
num_edge_mlp:
Number of mlp layer used for edge model
edge_dropout_rate:
dropout rate for the edges
aggregate_features(input_features, senders, receivers, sender_features, receiver_features, size)
¶
Function to aggregate (scatter) messages built from node and edge features.
Parameters:
input_features:
Edge features of the batch
senders:
Senders of the edge_index of the batch
receivers:
Receivers of the edge_index of the batch
sender_features:
Senders features gathered from the gather_features function
receiver_features:
Receiver features gathered from the gather_features function
size:
size of the aggregation, equals to the total number of nodes
Returns:
Name  Type  Description 

Tensor 
Tensor

Aggregated node features 
forward(batch)
¶
Forward function of the MPNN Plus layer Parameters: batch: pyg Batch graph to pass through the layer Returns: batch: pyg Batch graph with updated node and edge features
gather_features(input_features, senders, receivers)
¶
Function to gather node features based on the senders and receivers of the edge indices.
Parameters:
input_features:
Node features of the batch
senders:
Senders of the edge_index of the batch
receivers:
Receivers of the edge_index of the batch
Output
Gathered node features (sender and receiver) summed up or concatenated Gathered sender features Gathered receiver features
layer_supports_edges()
¶
Return a boolean specifying if the layer type supports edges or not.
Returns:
supports_edges: bool
Always ``True`` for the current class
PNA Layer¶
graphium.nn.pyg_layers.pna_pyg
¶
PNAMessagePassingPyg
¶
Bases: MessagePassing
, BaseGraphStructure
layer_inputs_edges: bool
property
¶
Return a boolean specifying if the layer type
uses edges as input or not.
It is different from layer_supports_edges
since a layer that
supports edges can decide to not use them.
Returns:
bool:
Returns ``self.edge_features``
layer_outputs_edges: bool
property
¶
Abstract method. Return a boolean specifying if the layer type
uses edges as input or not.
It is different from layer_supports_edges
since a layer that
supports edges can decide to not use them.
Returns:
bool:
Always ``False`` for the current class
out_dim_factor: int
property
¶
Get the factor by which the output dimension is multiplied for the next layer.
For standard layers, this will return 1
.
But for others, such as GatLayer
, the output is the concatenation
of the outputs from each head, so the out_dim gets multiplied by
the number of heads, and this function should return the number
of heads.
Returns:
int:
Always ``1`` for the current class
__init__(in_dim, out_dim, aggregators, scalers, activation='relu', dropout=0.0, normalization='none', avg_d={'log': 1.0, 'lin': 1.0}, last_activation='none', posttrans_layers=1, pretrans_layers=1, in_dim_edges=0, **kwargs)
¶
Implementation of the message passing architecture of the PNA message passing layer,
previously known as PNALayerComplex
. This layer applies an MLP as
pretransformation to the concatenation of \([h_u, h_v, e_{uv}]\) to generate
the messages, with \(h_u\) the node feature, \(h_v\) the neighbour node features,
and \(e_{uv}\) the edge feature between the nodes \(u\) and \(v\).
After the pretransformation, it aggregates the messages multiple aggregators and scalers, concatenates their results, then applies an MLP on the concatenated features.
PNA: Principal Neighbourhood Aggregation Gabriele Corso, Luca Cavalleri, Dominique Beaini, Pietro Lio, Petar Velickovic https://arxiv.org/abs/2004.05718
[!] code adapted from pytorchgeometric implementation of PNAConv
Parameters:
in_dim:
Input feature dimensions of the layer
out_dim:
Output feature dimensions of the layer
aggregators:
Set of aggregation function identifiers,
e.g. "mean", "max", "min", "std", "sum", "var", "moment3".
The results from all aggregators will be concatenated.
scalers:
Set of scaling functions identifiers
e.g. "identidy", "amplification", "attenuation"
The results from all scalers will be concatenated
activation:
activation function to use in the layer
dropout:
The ratio of units to dropout. Must be between 0 and 1
normalization:
Normalization to use. Choices:
 "none" or `None`: No normalization
 "batch_norm": Batch normalization
 "layer_norm": Layer normalization
 `Callable`: Any callable function
avg_d:
Average degree of nodes in the training set, used by scalers to normalize
last_activation:
activation function to use in the last layer of the internal MLP
posttrans_layers:
number of layers in the MLP transformation after the aggregation
pretrans_layers:
number of layers in the transformation before the aggregation
in_dim_edges:
size of the edge features. If 0, edges are ignored
aggregate(inputs, index, edge_index, dim_size=None)
¶
aggregate function
Parameters:
Name  Type  Description  Default 

inputs 
Tensor

input features 
required 
index 
Tensor

index of the nodes 
required 
edge_index 
Tensor

edge index 
required 
dim_size 
Optional[int]

dimension size 
None

Returns: out: aggregated features
forward(batch)
¶
forward function of the layer Parameters: batch: pyg Batch graphs Returns: batch: pyg Batch graphs
layer_supports_edges()
¶
Return a boolean specifying if the layer type supports edges or not.
Returns:
bool:
Always ``True`` for the current class
message(x_i, x_j, edge_feat)
¶
message function
Parameters:
Name  Type  Description  Default 

x_i 
Tensor

node features 
required 
x_j 
Tensor

neighbour node features 
required 
edge_feat 
OptTensor

edge features 
required 
Returns: feat: the message
Pooling Layers¶
graphium.nn.pyg_layers.pooling_pyg
¶
PoolingWrapperPyg
¶
Bases: ModuleWrap
forward(g, feature, *args, **kwargs)
¶
forward function Parameters: g: the pyg batch graph feature: the node features Returns: the pooled features
VirtualNodePyg
¶
Bases: Module
__init__(in_dim, out_dim, in_dim_edges, out_dim_edges, vn_type='sum', activation='relu', dropout=0.0, normalization='none', bias=True, residual=True, use_edges=False, **kwargs)
¶
The VirtualNode is a layer that pool the features of the graph, applies a neural network layer on the pooled features, then add the result back to the node features of every node.
Parameters:
in_dim:
Input feature dimensions of the virtual node layer.
out_dim:
Output feature dimensions of the virtual node layer.
in_dim_edges:
Input feature dimensions of the virtual node layer for the edges.
out_dim_edges:
Output feature dimensions of the virtual node layer for the edges.
vn_type:
The type of the virtual node. Choices are:
 "none": No pooling
 "sum": Sum all the nodes for each graph
 "mean": Mean all the nodes for each graph
 "logsum": Mean all the nodes then multiply by log(num_nodes) for each graph
 "max": Max all the nodes for each graph
 "min": Min all the nodes for each graph
 "std": Standard deviation of all the nodes for each graph
activation:
activation function to use in the neural network layer.
dropout:
The ratio of units to dropout. Must be between 0 and 1
normalization:
Normalization to use. Choices:
 "none" or `None`: No normalization
 "batch_norm": Batch normalization
 "layer_norm": Layer normalization
 `Callable`: Any callable function
bias:
Whether to add a bias to the neural network
residual:
Whether all virtual nodes should be connected together
via a residual connection
use_edges:
Boolean flag to select if edges are used in the global node
aggregation and update of features
forward(g, feat, vn_feat, edge_feat)
¶
Apply the virtual node layer.
Parameters:
g:
PyG Graphs or Batched graphs.
feat (torch.Tensor[..., N, Din]):
Node feature tensor, before convolution.
`N` is the number of nodes, `Din` is the input features
vn_feat (torch.Tensor[..., M, Din]):
Graph feature of the previous virtual node, or `None`
`M` is the number of graphs, `Din` is the input features.
It is added to the result after the MLP, as a residual connection
edge_feat (torch.Tensor[..., E, Din]):
Edge feature tensor, before convolution.
`E` is the number of edges, `Din` is the input features
Returns:
`feat = torch.Tensor[..., N, Dout]`:
Node feature tensor, after convolution and residual.
`N` is the number of nodes, `Dout` is the output features of the layer and residual
`vn_feat = torch.Tensor[..., M, Dout]`:
Graph feature tensor to be used at the next virtual node, or `None`
`M` is the number of graphs, `Dout` is the output features
`edge_feat = torch.Tensor[..., N, Dout]`:
Edge feature tensor, after convolution and residual  if edges are used, otherwise returned unchanged.
`N` is the number of edges, `Dout` is the output features of the layer and residual
parse_pooling_layer_pyg(in_dim, pooling, feat_type='node', **kwargs)
¶
Select the pooling layers from a list of strings, and put them in a Module that concatenates their outputs.
Parameters:
in_dim:
The dimension at the input layer of the pooling
pooling:
The list of pooling layers to use. The accepted strings are:
 "none": No pooling
 "sum": Sum all the nodes for each graph
 "mean": Mean all the nodes for each graph
 "logsum": Mean all the nodes then multiply by log(num_nodes) for each graph
 "max": Max all the nodes for each graph
 "min": Min all the nodes for each graph
 "std": Standard deviation of all the nodes for each graph
scatter_logsum_pool(x, batch, dim=0, dim_size=None)
¶
Apply pooling over the nodes in the graph using a mean aggregation, but scaled by the log of the number of nodes. This gives the same expressive power as the sum, but helps deal with graphs that are significantly larger than others by using a logarithmic scale.
Parameters:
Name  Type  Description  Default 

x 
Tensor

Node feature matrix
:math: 
required 
batch 
LongTensor

Batch vector :math: 
required 
size 
int

Batchsize :math: 
required 
Returns: the pooled features tensor
scatter_std_pool(x, batch, dim=0, dim_size=None)
¶
Returns batchwise graphleveloutputs by taking the channelwise
minimum across the node dimension, so that for a single graph
:math:\mathcal{G}_i
its output is computed by
.. math:: \mathbf{r}i = \mathrm{max}{n=1}^{N_i} \, \mathbf{x}_n
Parameters:
Name  Type  Description  Default 

x 
Tensor

Node feature matrix
:math: 
required 
batch 
LongTensor

Batch vector :math: 
required 
size 
int

Batchsize :math: 
required 
Returns:
Type  Description 

the pooled features tensor 
Utils¶
graphium.nn.pyg_layers.utils
¶
GaussianLayer
¶
Bases: Module
__init__(num_kernels=128, in_dim=3)
¶
Gaussian kernel function that applied on the alltoall 3D distances.
Parameters: num_kernels: Number of gaussian kernel used.
PreprocessPositions
¶
Bases: Module
Compute 3D attention bias and 3D node features according to the 3D position information.
__init__(num_heads, embed_dim, num_kernel, in_dim=3, num_layers=2, activation='gelu', first_normalization='none')
¶
Parameters:
Name  Type  Description  Default 

num_heads 
Number of attention heads used in selfattention. 
required  
embed_dim 
Hidden dimension of node features. 
required  
num_kernel 
Number of gaussian kernels. 
required  
num_layers 
The number of layers in the MLP. 
2


activation 
The activation function used in the MLP. 
'gelu'


first_normalization 
The normalization function used before the gaussian kernel. 
'none'

forward(batch, max_num_nodes_per_graph, on_ipu, positions_3d_key)
¶
Inputs
batch: Batch object. max_num_nodes_per_graph: Maximum number of nodes per graph. on_ipu: If model rus on IPU. positions_3d_key: The key of the pyg graph object that contains the 3D positions.
triplets(edge_index, num_nodes)
¶
Generates triplets from the given edge indices. A triplet is defined as a path of length two, such that if node A is connected to node B, and node B is connected to node C, then there is a triplet (A, B, C).
Parameters:
Name  Type  Description  Default 

edge_index 
LongTensor

The edge indices. 
required 
num_nodes 
int

The number of nodes. 
required 
Returns:
Name  Type  Description 

col 
Tensor

The sink node indices of edges from the edge indices. 
row 
Tensor

The source node indices of edges from the edge indices. 
idx_i 
Tensor

The sink node indices of the triplets. 
idx_j 
Tensor

The middle node indices of the triplets. 
idx_k 
Tensor

The source node indices of the triplets. 
idx_kj 
Tensor

The indices of edges those from the source node to the middle node. 
idx_ji 
Tensor

The indices of edges those from the middle node to the sink node. 