graphium.nn.pyg_layers¶

Implementations of GNN layers based on PyG in the library

Contents

Gated GCN Layer
GPS Layer
GIN and GINE Layers
MPNN Layer
PNA Layer
Pooling Layer

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 `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:

Name	Type	Description
`bool`	`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:

Name	Type	Description
`int`	`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', **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:

Name	Type	Description	Default
`in_dim`	`int`	Input feature dimensions of the layer	required
`out_dim`	`int`	Output feature dimensions of the layer, and for the edges	required
`in_dim_edges`	`int`	Input edge-feature dimensions of the layer	required
`activation`	`Union[Callable, str]`	activation function to use in the layer	`'relu'`
`dropout`	`float`	The ratio of units to dropout. Must be between 0 and 1	`0.0`
`normalization`	`Union[str, Callable]`	Normalization to use. Choices: "none" or `None`: No normalization "batch_norm": Batch normalization "layer_norm": Layer normalization `Callable`: Any callable function	`'none'`

`aggregate(sigma_ij, index, Bx_j, Bx)` ¶

aggregation function of the layer

Parameters:

Name	Type	Description	Default
`sigma_ij`	`torch.Tensor`	the output from message() function with dim [n_edges, out_dim]	required
`index`	`torch.Tensor`	dim [n_edges]	required
`Bx_j`	`torch.Tensor`	dim [n_edges, out_dim]	required
`Bx`	`torch.Tensor`	dim [n_nodes, out_dim]	required

Returns:

Name	Type	Description
`out`	`torch.Tensor`	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:

Name	Type	Description
`batch`	`Union[Data, 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 `True` for the current class

`message(Dx_i, Ex_j, Ce)` ¶

message function

Parameters:

Name	Type	Description	Default
`Dx_i`	`torch.Tensor`	tensor with dimension [n_edges, out_dim]	required
`Ex_j`	`torch.Tensor`	tensor with dimension [n_edges, out_dim]	required
`Ce`	`torch.Tensor`	tensor with dimension [n_edges, out_dim]	required

Returns:

Name	Type	Description
`sigma_ij`	`torch.Tensor`	tensor with dimension [n_edges, out_dim]

`update(aggr_out, Ax)` ¶

update function of the layer

Parameters:

Name	Type	Description	Default
`aggr_out`	`torch.Tensor`	the output from aggregate() function with dim [n_nodes, out_dim]	required
`Ax`	`torch.Tensor`	tensor with dim [n_nodes, out_dim]	required

Returns:

Name	Type	Description
`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:

Name	Type	Description
`bool`	`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:

Name	Type	Description
`bool`	`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:

Name	Type	Description
`int`	`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='full-attention', 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 Maddrell-Mander, Adam Sanders, Hatem Helal, Deniz Beker, Ladislav Rampášek, Dominique Beaini https://arxiv.org/abs/2212.02229

Parameters:

Name	Type	Description	Default
`in_dim`	`int`	Input node feature dimensions of the layer	required
`out_dim`	`int`	Output node feature dimensions of the layer	required
`in_dim`	`int`	Input edge feature dimensions of the layer	required
`out_dim`	`int`	Output edge feature dimensions of the layer	required
`in_dim_edges`	`Optional[int]`	input edge-feature dimensions of the layer	`None`
`out_dim_edges`	`Optional[int]`	output edge-feature dimensions of the layer	`None`
`activation`	`Union[Callable, str]`	activation function to use in the layer	`'relu'`
`dropout`	`float`	The ratio of units to dropout. Must be between 0 and 1	`0.0`
`node_residual`	`Optional[bool]`	If node residual is used after on the gnn layer output	`True`
`normalization`	`Union[str, Callable]`	Normalization to use. Choices: "none" or `None`: No normalization "batch_norm": Batch normalization "layer_norm": Layer normalization `Callable`: Any callable function	`'none'`
`mpnn_type`	`str`	type of mpnn used, choose from "pyg:gin", "pyg:gine", "pyg:gated-gcn", "pyg:pna-msgpass" and "pyg:mpnnplus"	`'pyg:gine'`
`mpnn_kwargs`		kwargs for mpnn layer	`None`
`attn_type`	`str`	type of attention used, choose from "full-attention" and "none"	`'full-attention'`
`attn_kwargs`		kwargs for attention layer	`None`
`droppath_rate_attn`	`float`	stochastic depth drop rate for attention layer https://arxiv.org/abs/1603.09382	`0.0`
`droppath_rate_ffn`	`float`	stochastic depth drop rate for ffn layer https://arxiv.org/abs/1603.09382	`0.0`
`mpnn_type`	`str`	Type of MPNN layer to use. Choices specified in PYG_LAYERS_DICT	`'pyg:gine'`
`mpnn_kwargs`		Keyword arguments to pass to the MPNN layer	`None`
`attn_type`	`str`	Type of attention layer to use. Choices specified in ATTENTION_LAYERS_DICT	`'full-attention'`
`biased_attention_key`	`Optional[str]`	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	`None`
`attn_kwargs`		Keyword arguments to pass to the attention layer	`None`

`forward(batch)` ¶

forward function of the layer

Parameters:

Name	Type	Description	Default
`batch`	`Batch`	pyg Batch graphs to pass through the layer	required

Returns:

Name	Type	Description
`batch`	`Batch`	pyg Batch graphs

`layer_supports_edges()` ¶

Return a boolean specifying if the layer type supports edges or not.

Returns:

Name	Type	Description
`supports_edges`	`bool`	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:

Name	Type	Description
`bool`	`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:

Name	Type	Description
`bool`	`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:

Name	Type	Description
`int`	`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 pytorch-geometric implementation of GINConv

Parameters:

Name	Type	Description	Default
`in_dim`	`int`	Input feature dimensions of the layer	required
`out_dim`	`int`	Output feature dimensions of the layer	required
`activation`	`Union[Callable, str]`	activation function to use in the layer	`'relu'`
`dropout`	`float`	The ratio of units to dropout. Must be between 0 and 1	`0.0`
`normalization`	`Union[str, Callable]`	Normalization to use. Choices: "none" or `None`: No normalization "batch_norm": Batch normalization "layer_norm": Layer normalization `Callable`: Any callable function	`'none'`
`init_eps`		Initial :math:`\epsilon` value, default: `0`.	required
`learn_eps`		If True, :math:`\epsilon` will be a learnable parameter.	required

`forward(batch)` ¶

forward function of the layer

Parameters:

Name	Type	Description	Default
`batch`	`Union[Data, Batch]`	pyg Batch graphs to pass through the layer	required

Returns:

Name	Type	Description
`batch`	`Union[Data, Batch]`	pyg Batch graphs

`layer_supports_edges()` ¶

Return a boolean specifying if the layer type supports edges or not.

Returns:

Name	Type	Description
`supports_edges`	`bool`	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:

Name	Type	Description
`bool`	`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:

Name	Type	Description
`bool`	`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:

Name	Type	Description
`int`	`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 Pre-training 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 pytorch-geometric implementation of GINEConv

Parameters:

Name	Type	Description	Default
`in_dim`	`int`	Input feature dimensions of the layer	required
`out_dim`	`int`	Output feature dimensions of the layer	required
`activation`	`Union[Callable, str]`	activation function to use in the layer	`'relu'`
`dropout`	`float`	The ratio of units to dropout. Must be between 0 and 1	`0.0`
`normalization`	`Union[str, Callable]`	Normalization to use. Choices: "none" or `None`: No normalization "batch_norm": Batch normalization "layer_norm": Layer normalization `Callable`: Any callable function	`'none'`
`init_eps`		Initial :math:`\epsilon` value, default: `0`.	required
`learn_eps`		If True, :math:`\epsilon` will be a learnable parameter.	required

`forward(batch)` ¶

forward function of the layer

Parameters:

Name	Type	Description	Default
`batch`	`Union[Data, Batch]`	pyg Batch graphs to pass through the layer	required

Returns:

Name	Type	Description
`batch`	`Union[Data, Batch]`	pyg Batch graphs

`layer_supports_edges()` ¶

Return a boolean specifying if the layer type supports edges or not.

Returns:

Name	Type	Description
`supports_edges`	`bool`	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:

Name	Type	Description
`bool`	`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:

Name	Type	Description
`bool`	`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:

Name	Type	Description
`int`	`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 Maddrell-Mander, Adam Sanders,
Hatem Helal, Deniz Beker, Ladislav Rampášek, Dominique Beaini
https://arxiv.org/abs/2212.02229

Parameters:

Name	Type	Description	Default
`in_dim`	`int`	Input feature dimensions of the nodes	`64`
`out_dim`	`int`	Output feature dimensions of the nodes	`64`
`activation`	`Union[str, Callable]`	Activation function to use in the edge and node model	`'gelu'`
`dropout`	`float`	The ratio of units to dropout at the end within apply_norm_activation_dropout.	`0.3`
`normalization`	`Union[str, Callable]`	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	`'layer_norm'`
`gather_from`	`str`	The method to gather features from. Could choose from: "senders", "receivers" and "both".	`'both'`
`scatter_to`	`str`	The method to scatter features to. Could choose from: "senders", "receivers" and "both".	`'both'`
`node_combine_method`	`str`	The method to combine the node features, Could choose from: "sum" and "concat".	`'concat'`
`aggregation_method`	`Optional[List[Union[str, Aggregation]]]`	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"	`['sum']`
`num_node_mlp`	`int`	Number of mlp layer used for node model	`2`
`mlp_expansion_ratio`	`int`	Expansion ratio for node and edge mlp	`4`
`use_edges`	`bool`	If edge features are used	`True`
`in_dim_edges`	`Optional[int]`	Input feature dimensions of the edges	`32`
`out_dim_edges`	`Optional[int]`	Output feature dimensions of the edges	`32`
`num_edge_mlp`	`Optional[int]`	Number of mlp layer used for edge model	`2`
`edge_dropout_rate`	`Optional[float]`	dropout rate for the edges	`0.0035`

`aggregate_features(input_features, senders, receivers, sender_features, receiver_features, size)` ¶

Function to aggregate (scatter) messages built from node and edge features.

Parameters:

Name	Type	Description	Default
`input_features`	`Tensor`	Edge features of the batch	required
`senders`	`Union[IntTensor, LongTensor]`	Senders of the edge_index of the batch	required
`receivers`	`Union[IntTensor, LongTensor]`	Receivers of the edge_index of the batch	required
`sender_features`	`Tensor`	Senders features gathered from the gather_features function	required
`receiver_features`	`Tensor`	Receiver features gathered from the gather_features function	required
`size`	`int`	size of the aggregation, equals to the total number of nodes	required

Returns:

Name	Type	Description
`Tensor`	`Tensor`	Aggregated node features

`forward(batch)` ¶

Forward function of the MPNN Plus layer

Parameters:

Name	Type	Description	Default
`batch`	`Batch`	pyg Batch graph to pass through the layer	required

Returns:

Name	Type	Description
`batch`	`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:

Name	Type	Description	Default
`input_features`	`Tensor`	Node features of the batch	required
`senders`	`Union[IntTensor, LongTensor]`	Senders of the edge_index of the batch	required
`receivers`	`Union[IntTensor, LongTensor]`	Receivers of the edge_index of the batch	required

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:

Name	Type	Description
`supports_edges`	`bool`	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:

Name	Type	Description
`bool`	`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:

Name	Type	Description
`bool`	`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:

Name	Type	Description
`int`	`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 pre-transformation, 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 pytorch-geometric implementation of PNAConv

Parameters:

Name	Type	Description	Default
`in_dim`	`int`	Input feature dimensions of the layer	required
`out_dim`	`int`	Output feature dimensions of the layer	required
`aggregators`	`List[str]`	Set of aggregation function identifiers, e.g. "mean", "max", "min", "std", "sum", "var", "moment3". The results from all aggregators will be concatenated.	required
`scalers`	`List[str]`	Set of scaling functions identifiers e.g. "identidy", "amplification", "attenuation" The results from all scalers will be concatenated	required
`activation`	`Union[Callable, str]`	activation function to use in the layer	`'relu'`
`dropout`	`float`	The ratio of units to dropout. Must be between 0 and 1	`0.0`
`normalization`	`Union[str, Callable]`	Normalization to use. Choices: "none" or `None`: No normalization "batch_norm": Batch normalization "layer_norm": Layer normalization `Callable`: Any callable function	`'none'`
`avg_d`	`Dict[str, float]`	Average degree of nodes in the training set, used by scalers to normalize	`{'log': 1.0, 'lin': 1.0}`
`last_activation`	`Union[Callable, str]`	activation function to use in the last layer of the internal MLP	`'none'`
`posttrans_layers`	`int`	number of layers in the MLP transformation after the aggregation	`1`
`pretrans_layers`	`int`	number of layers in the transformation before the aggregation	`1`
`in_dim_edges`	`int`	size of the edge features. If 0, edges are ignored	`0`

`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:

Name	Type	Description
`out`	`Tensor`	aggregated features

`forward(batch)` ¶

forward function of the layer

Parameters:

Name	Type	Description	Default
`batch`	`Union[Data, Batch]`	pyg Batch graphs	required

Returns:

Name	Type	Description
`batch`	`Union[Data, Batch]`	pyg Batch graphs

`layer_supports_edges()` ¶

Return a boolean specifying if the layer type supports edges or not.

Returns:

Name	Type	Description
`bool`	`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:

Name	Type	Description
`feat`	`Tensor`	the message

Pooling Layers¶

`graphium.nn.pyg_layers.pooling_pyg` ¶

`PoolingWrapperPyg` ¶

Bases: ModuleWrap

`forward(g, feature, *args, **kwargs)` ¶

forward function

Parameters:

Name	Type	Description	Default
`g`	`Batch`	the pyg batch graph	required
`feature`	`Tensor`	the node features	required

Returns:

Type	Description
	the pooled features

`VirtualNodePyg` ¶

Bases: nn.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:

Name	Type	Description	Default
`in_dim`	`int`	Input feature dimensions of the virtual node layer.	required
`out_dim`	`int`	Output feature dimensions of the virtual node layer.	required
`in_dim_edges`	`Optional[int]`	Input feature dimensions of the virtual node layer for the edges.	required
`out_dim_edges`	`Optional[int]`	Output feature dimensions of the virtual node layer for the edges.	required
`vn_type`	`Union[type(None), str]`	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	`'sum'`
`activation`	`Union[str, Callable]`	activation function to use in the neural network layer.	`'relu'`
`dropout`	`float`	The ratio of units to dropout. Must be between 0 and 1	`0.0`
`normalization`	`Union[str, Callable]`	Normalization to use. Choices: "none" or `None`: No normalization "batch_norm": Batch normalization "layer_norm": Layer normalization `Callable`: Any callable function	`'none'`
`bias`	`bool`	Whether to add a bias to the neural network	`True`
`residual`	`bool`	Whether all virtual nodes should be connected together via a residual connection	`True`
`use_edges`	`bool`	Boolean flag to select if edges are used in the global node aggregation and update of features	`False`

`forward(g, feat, vn_feat, edge_feat)` ¶

Apply the virtual node layer.

Parameters:

Name	Type	Description	Default
`g`	`Union[Data, Batch]`	PyG Graphs or Batched graphs.	required
`feat`	`torch.Tensor[..., N, Din]`	Node feature tensor, before convolution. `N` is the number of nodes, `Din` is the input features	required
`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	required
`edge_feat`	`torch.Tensor[..., E, Din]`	Edge feature tensor, before convolution. `E` is the number of edges, `Din` is the input features	required

Returns:

Type	Description
`Tensor`	`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
`Tensor`	`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
`Tensor`	`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:

Name	Type	Description	Default
`in_dim`	`int`	The dimension at the input layer of the pooling	required
`pooling`	`Union[str, List[str]]`	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	required

`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.

\[r^{(i)} = \frac{\log N_i}{N_i}\sum_{k=1}^{N_i} x^{(i)}_k\]

Parameters:

Name	Type	Description	Default
`x`	`Tensor`	Node feature matrix :math:`\mathbf{X} \in \mathbb{R}^{(N_1 + \ldots + N_B) \times F}`.	required
`batch`	`LongTensor`	Batch vector :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns each node to a specific example.	required
`size`	`int`	Batch-size :math:`B`. Automatically calculated if not given. (default: :obj:`None`)	required

Returns:

Type	Description
`Tensor`	the pooled features tensor

`scatter_std_pool(x, batch, dim=0, dim_size=None)` ¶

Returns batch-wise graph-level-outputs by taking the channel-wise 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:`\mathbf{X} \in \mathbb{R}^{(N_1 + \ldots + N_B) \times F}`.	required
`batch`	`LongTensor`	Batch vector :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns each node to a specific example.	required
`size`	`int`	Batch-size :math:`B`. Automatically calculated if not given. (default: :obj:`None`)	required

Returns:

Type	Description
	the pooled features tensor

Utils¶

`graphium.nn.pyg_layers.utils` ¶

`GaussianLayer` ¶

Bases: nn.Module

`init(num_kernels=128, in_dim=3)` ¶

Gaussian kernel function that applied on the all-to-all 3D distances.

Parameters:

Name	Type	Description	Default
`num_kernels`		Number of gaussian kernel used.	`128`

`PreprocessPositions` ¶

Bases: nn.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	Description	Default
`num_heads`	Number of attention heads used in self-attention.	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.

graphium.nn.pyg_layers¶

Gated GCN Layer¶

graphium.nn.pyg_layers.gated_gcn_pyg ¶

GatedGCNPyg ¶

layer_inputs_edges: bool property ¶

layer_outputs_edges: bool property ¶

out_dim_factor: int property ¶

__init__(in_dim, out_dim, in_dim_edges, out_dim_edges=None, activation='relu', dropout=0.0, normalization='none', **kwargs) ¶

aggregate(sigma_ij, index, Bx_j, Bx) ¶

forward(batch) ¶

layer_supports_edges() ¶

message(Dx_i, Ex_j, Ce) ¶

update(aggr_out, Ax) ¶

GPS Layer¶

graphium.nn.pyg_layers.gps_pyg ¶

GPSLayerPyg ¶

layer_inputs_edges: bool property ¶

layer_outputs_edges: bool property ¶

out_dim_factor: int property ¶

forward(batch) ¶

layer_supports_edges() ¶

GIN and GINE Layers¶

graphium.nn.pyg_layers.gin_pyg ¶

GINConvPyg ¶

layer_inputs_edges: bool property ¶

layer_outputs_edges: bool property ¶

out_dim_factor: int property ¶

__init__(in_dim, out_dim, activation='relu', dropout=0.0, normalization='none', **kwargs) ¶

forward(batch) ¶

layer_supports_edges() ¶

GINEConvPyg ¶

layer_inputs_edges: bool property ¶

layer_outputs_edges: bool property ¶

out_dim_factor: int property ¶

__init__(in_dim, out_dim, in_dim_edges=None, activation='relu', dropout=0.0, normalization='none', **kwargs) ¶

forward(batch) ¶

layer_supports_edges() ¶

MPNN Layer¶

graphium.nn.pyg_layers.mpnn_pyg ¶

MPNNPlusPyg ¶

layer_inputs_edges: bool property ¶

layer_outputs_edges: bool property ¶

out_dim_factor: int property ¶

aggregate_features(input_features, senders, receivers, sender_features, receiver_features, size) ¶

forward(batch) ¶

gather_features(input_features, senders, receivers) ¶

layer_supports_edges() ¶

PNA Layer¶

graphium.nn.pyg_layers.pna_pyg ¶

PNAMessagePassingPyg ¶

layer_inputs_edges: bool property ¶

layer_outputs_edges: bool property ¶

out_dim_factor: int property ¶

__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) ¶

aggregate(inputs, index, edge_index, dim_size=None) ¶

forward(batch) ¶

layer_supports_edges() ¶

message(x_i, x_j, edge_feat) ¶

Pooling Layers¶

graphium.nn.pyg_layers.pooling_pyg ¶

PoolingWrapperPyg ¶

forward(g, feature, *args, **kwargs) ¶

VirtualNodePyg ¶

__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) ¶

forward(g, feat, vn_feat, edge_feat) ¶

parse_pooling_layer_pyg(in_dim, pooling, feat_type='node', **kwargs) ¶

scatter_logsum_pool(x, batch, dim=0, dim_size=None) ¶

scatter_std_pool(x, batch, dim=0, dim_size=None) ¶

Utils¶

graphium.nn.pyg_layers.utils ¶

GaussianLayer ¶

__init__(num_kernels=128, in_dim=3) ¶

PreprocessPositions ¶

__init__(num_heads, embed_dim, num_kernel, in_dim=3, num_layers=2, activation='gelu', first_normalization='none') ¶

forward(batch, max_num_nodes_per_graph, on_ipu, positions_3d_key) ¶

triplets(edge_index, num_nodes) ¶

`graphium.nn.pyg_layers.gated_gcn_pyg` ¶

`GatedGCNPyg` ¶

`layer_inputs_edges: bool` `property` ¶

`layer_outputs_edges: bool` `property` ¶

`out_dim_factor: int` `property` ¶

`init(in_dim, out_dim, in_dim_edges, out_dim_edges=None, activation='relu', dropout=0.0, normalization='none', **kwargs)` ¶

`aggregate(sigma_ij, index, Bx_j, Bx)` ¶

`forward(batch)` ¶

`layer_supports_edges()` ¶

`message(Dx_i, Ex_j, Ce)` ¶

`update(aggr_out, Ax)` ¶

`graphium.nn.pyg_layers.gps_pyg` ¶

`GPSLayerPyg` ¶

`layer_inputs_edges: bool` `property` ¶

`layer_outputs_edges: bool` `property` ¶

`out_dim_factor: int` `property` ¶

`forward(batch)` ¶

`layer_supports_edges()` ¶

`graphium.nn.pyg_layers.gin_pyg` ¶

`GINConvPyg` ¶

`layer_inputs_edges: bool` `property` ¶

`layer_outputs_edges: bool` `property` ¶

`out_dim_factor: int` `property` ¶

`init(in_dim, out_dim, activation='relu', dropout=0.0, normalization='none', **kwargs)` ¶

`forward(batch)` ¶

`layer_supports_edges()` ¶

`GINEConvPyg` ¶

`layer_inputs_edges: bool` `property` ¶

`layer_outputs_edges: bool` `property` ¶

`out_dim_factor: int` `property` ¶

`init(in_dim, out_dim, in_dim_edges=None, activation='relu', dropout=0.0, normalization='none', **kwargs)` ¶

`forward(batch)` ¶

`layer_supports_edges()` ¶

`graphium.nn.pyg_layers.mpnn_pyg` ¶

`MPNNPlusPyg` ¶

`layer_inputs_edges: bool` `property` ¶

`layer_outputs_edges: bool` `property` ¶

`out_dim_factor: int` `property` ¶

`aggregate_features(input_features, senders, receivers, sender_features, receiver_features, size)` ¶

`forward(batch)` ¶

`gather_features(input_features, senders, receivers)` ¶

`layer_supports_edges()` ¶

`graphium.nn.pyg_layers.pna_pyg` ¶

`PNAMessagePassingPyg` ¶

`layer_inputs_edges: bool` `property` ¶

`layer_outputs_edges: bool` `property` ¶

`out_dim_factor: int` `property` ¶

`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)` ¶

`aggregate(inputs, index, edge_index, dim_size=None)` ¶

`forward(batch)` ¶

`layer_supports_edges()` ¶

`message(x_i, x_j, edge_feat)` ¶

`graphium.nn.pyg_layers.pooling_pyg` ¶

`PoolingWrapperPyg` ¶

`forward(g, feature, *args, **kwargs)` ¶

`VirtualNodePyg` ¶

`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)` ¶

`forward(g, feat, vn_feat, edge_feat)` ¶

`parse_pooling_layer_pyg(in_dim, pooling, feat_type='node', **kwargs)` ¶

`scatter_logsum_pool(x, batch, dim=0, dim_size=None)` ¶

`scatter_std_pool(x, batch, dim=0, dim_size=None)` ¶

`graphium.nn.pyg_layers.utils` ¶

`GaussianLayer` ¶

`init(num_kernels=128, in_dim=3)` ¶

`PreprocessPositions` ¶

`init(num_heads, embed_dim, num_kernel, in_dim=3, num_layers=2, activation='gelu', first_normalization='none')` ¶

`forward(batch, max_num_nodes_per_graph, on_ipu, positions_3d_key)` ¶

`triplets(edge_index, num_nodes)` ¶