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

Copyright (c) 2023 Valence Labs, Recursion Pharmaceuticals. Use of this software is subject to the terms and conditions outlined in the LICENSE file. Unauthorized modification, distribution, or use is prohibited. Provided 'as is' without warranties of any kind.

Valence Labs, Recursion Pharmaceuticals are not liable for any damages arising from its use. Refer to the LICENSE file for the full terms and conditions.

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

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=1e-05, **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 edge-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

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 `True` for the current class

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

Use of this software is subject to the terms and conditions outlined in the LICENSE file. Unauthorized modification, distribution, or use is prohibited. Provided 'as is' without warranties of any kind.

Valence Labs, Recursion Pharmaceuticals and Graphcore Limited are not liable for any damages arising from its use. Refer to the LICENSE file for the full terms and conditions.

`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='full-attention', precision='32', biased_attention_key=None, attn_kwargs=None, force_consistent_in_dim=True, droppath_rate_attn=0.0, droppath_rate_ffn=0.0, hidden_dim_scaling=4.0, output_scale=1.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:

in_dim:
    Input node feature dimensions of the layer

out_dim:
    Output node feature dimensions of the layer

in_dim_edges:
    input edge-feature dimensions of the layer

out_dim_edges:
    output edge-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

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:gated-gcn", "pyg:pna-msgpass" and "pyg:mpnnplus"

mpnn_kwargs:
    kwargs for mpnn layer

attn_type:
    type of attention used, choose from "full-attention" and "none"

attn_kwargs:
    kwargs for attention layer

force_consistent_in_dim:
    whether to force the `embed_dim` to be the same as the `in_dim` for the attention and mpnn.
    The argument is only valid if `attn_type` is not None. If `embed_dim` is not provided,
    it will be set to `in_dim` by default, so this parameter won't have an effect.

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

output_scale:
    Float value that will be used to scale the activations, helps reduce growth of activations

    as the model gets deeper. Default value of 1.0 leaves the layer unchanged.

`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

`residual_add(feature, input_feature)` ¶

Residual additition layer. Allows information to propagate through the model by skipping the computational layers. Parameters: feature: The feature (typically nodes or edges) after message passing input_feature: The same feature from before message passing Returns: The addition of the two tensors.

`scale_activations(feature, scale_factor)` ¶

Scale Activations by a constant factor to stop growth of activation scale and reduce numerical stability issues at low precision

Parameters:

Name	Type	Description	Default
`feature`	`Tensor`	The feature to scale	required
`scale_factor`	`float`	The floating point scale factor	required

Returns:

Name	Type	Description
`Tensor`	`Tensor`	The scaled features

GIN and GINE Layers¶

`graphium.nn.pyg_layers.gin_pyg` ¶

Copyright (c) 2023 Valence Labs, Recursion Pharmaceuticals. Use of this software is subject to the terms and conditions outlined in the LICENSE file. Unauthorized modification, distribution, or use is prohibited. Provided 'as is' without warranties of any kind.

Valence Labs, Recursion Pharmaceuticals are not liable for any damages arising from its use. Refer to the LICENSE file for the full terms and conditions.

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

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

Use of this software is subject to the terms and conditions outlined in the LICENSE file. Unauthorized modification, distribution, or use is prohibited. Provided 'as is' without warranties of any kind.

Valence Labs, Recursion Pharmaceuticals and Graphcore Limited are not liable for any damages arising from its use. Refer to the LICENSE file for the full terms and conditions.

`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 Maddrell-Mander, 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` ¶

Copyright (c) 2023 Valence Labs, Recursion Pharmaceuticals. Use of this software is subject to the terms and conditions outlined in the LICENSE file. Unauthorized modification, distribution, or use is prohibited. Provided 'as is' without warranties of any kind.

Valence Labs, Recursion Pharmaceuticals are not liable for any damages arising from its use. Refer to the LICENSE file for the full terms and conditions.

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

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

Use of this software is subject to the terms and conditions outlined in the LICENSE file. Unauthorized modification, distribution, or use is prohibited. Provided 'as is' without warranties of any kind.

Valence Labs, Recursion Pharmaceuticals and Graphcore Limited are not liable for any damages arising from its use. Refer to the LICENSE file for the full terms and conditions.

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

\[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: 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` ¶

Use of this software is subject to the terms and conditions outlined in the LICENSE file. Unauthorized modification, distribution, or use is prohibited. Provided 'as is' without warranties of any kind.

Valence Labs, Recursion Pharmaceuticals and Graphcore Limited are not liable for any damages arising from its use. Refer to the LICENSE file for the full terms and conditions.

`GaussianLayer` ¶

Bases: Module

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

Gaussian kernel function that applied on the all-to-all 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	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', eps=1e-05, **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() ¶

residual_add(feature, input_feature) ¶

scale_activations(feature, scale_factor) ¶

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', eps=1e-05, **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()` ¶

`residual_add(feature, input_feature)` ¶

`scale_activations(feature, scale_factor)` ¶

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