Skip to content

Models API Reference

The k3_node.models package contains complete, production-ready Graph Neural Network architectures implemented natively in Keras 3.

Standard GNN Architectures

BasicGNN

k3_node.models.BasicGNN

Bases: Layer

An abstract base class for implementing basic GNN models.

Parameters:

Name Type Description Default
in_channels int or tuple

Size of each input sample.

 required
hidden_channels int

Size of each hidden sample.

 required
num_layers int

Number of message passing layers.

 required
out_channels int

If not set to :obj:None, will apply a final linear transformation to convert hidden node embeddings to output size :obj:out_channels. (default: :obj:None)

 None
dropout float

Dropout probability. (default: :obj:0.)

 0.0
act str or Callable

The non-linear activation function to use. (default: :obj:"relu")

 'relu'
act_first bool

If set to :obj:True, activation is applied before normalization. (default: :obj:False)

 False
act_kwargs Dict[str, Any]

Arguments passed to the respective activation function defined by :obj:act. (default: :obj:None)

 None
norm str or Callable

The normalization function to use. (default: :obj:None)

 None
norm_kwargs Dict[str, Any]

Arguments passed to the respective normalization function defined by :obj:norm. (default: :obj:None)

 None
jk str

The Jumping Knowledge mode. If specified, the model will additionally apply a final linear transformation to transform node embeddings to the expected output feature dimensionality. (:obj:None, :obj:"last", :obj:"cat", :obj:"max", :obj:"lstm"). (default: :obj:None)

 None
**kwargs optional

Additional arguments of the underlying :class:torch_geometric.nn.conv.MessagePassing layers.

 {}

reset_parameters()

Resets all learnable parameters of the module.

GCN

k3_node.models.GCN

Bases: BasicGNN

The Graph Neural Network from the "Semi-supervised Classification with Graph Convolutional Networks" <https://arxiv.org/abs/1609.02907>_ paper, using the :class:~k3_node.layers.conv.GCNConv operator for message passing.

GraphSAGE

k3_node.models.GraphSAGE

Bases: BasicGNN

The Graph Neural Network from the "Inductive Representation Learning on Large Graphs" <https://arxiv.org/abs/1706.02216>_ paper, using the :class:~k3_node.layers.conv.SAGEConv operator for message passing.

GIN

k3_node.models.GIN

Bases: BasicGNN

The Graph Neural Network from the "How Powerful are Graph Neural Networks?" <https://arxiv.org/abs/1810.00826>_ paper, using the :class:~k3_node.layers.conv.GINConv operator for message passing.

GAT

k3_node.models.GAT

Bases: BasicGNN

The Graph Neural Network from "Graph Attention Networks" <https://arxiv.org/abs/1710.10903> or "How Attentive are Graph Attention Networks?" <https://arxiv.org/abs/2105.14491> papers, using the :class:~k3_node.layers.conv.GATConv or :class:~k3_node.layers.conv.GATv2Conv operator for message passing.

PNA

k3_node.models.PNA

Bases: BasicGNN

The Graph Neural Network from the "Principal Neighbourhood Aggregation for Graph Nets" <https://arxiv.org/abs/2004.05718>_ paper, using the :class:~k3_node.layers.conv.PNAConv operator for message passing.

EdgeCNN

k3_node.models.EdgeCNN

Bases: BasicGNN

The Graph Neural Network from the "Dynamic Graph CNN for Learning on Point Clouds" <https://arxiv.org/abs/1801.07829>_ paper, using the :class:~k3_node.layers.conv.EdgeConv operator for message passing.

MLP

k3_node.models.MLP

Bases: Layer

A Multi-Layer Perceptron (MLP) model.

There exists two ways to instantiate an MLP:

  1. By specifying explicit channel sizes, e.g., MLP([16, 32, 64, 128]) creates a three-layer MLP with differently sized hidden layers.

  2. By specifying fixed hidden channel sizes over a number of layers, e.g., MLP(in_channels=16, hidden_channels=32, out_channels=128, num_layers=3) creates a three-layer MLP with equally sized hidden layers.

Parameters:

Name Type Description Default
channel_list List[int] or int

List of input, intermediate and output channels such that len(channel_list) - 1 denotes the number of layers of the MLP. (default: None)

 None
in_channels int

Size of each input sample. Will override channel_list. (default: None)

 None
hidden_channels int

Size of each hidden sample. Will override channel_list. (default: None)

 None
out_channels int

Size of each output sample. Will override channel_list. (default: None)

 None
num_layers int

The number of layers. Will override channel_list. (default: None)

 None
dropout float or List[float]

Dropout probability of each hidden embedding. (default: 0.)

 0.0
act str or Callable

The non-linear activation function to use. (default: "relu")

 'relu'
act_first bool

If set to True, activation is applied before normalization. (default: False)

 False
norm str or Callable

The normalization function to use. (default: "batch_norm")

 'batch_norm'
norm_kwargs dict

Arguments passed to the respective normalization function. (default: None)

 None
plain_last bool

If set to False, will apply non-linearity, normalization and dropout to the last layer as well. (default: True)

 True
bias bool or List[bool]

If set to False, the module will not learn additive biases. (default: True)

 True

num_layers property

The number of layers.

reset_parameters()

Resets all learnable parameters of the module.

Graph Autoencoders & Self-Supervised Models

GAE

k3_node.models.GAE

The Graph Auto-Encoder model from the "Variational Graph Auto-Encoders" <https://arxiv.org/abs/1611.07308>_ paper based on user-defined encoder and decoder models.

Parameters:

Name Type Description Default
encoder

The encoder module.

 required
decoder optional

The decoder module. If set to None, will default to InnerProductDecoder. (default: None)

 None

__call__(*args, **kwargs)

Alias for encode.

decode(*args, **kwargs)

Runs the decoder and computes edge probabilities.

encode(*args, **kwargs)

Runs the encoder and computes node-wise latent variables.

recon_loss(z, pos_edge_index, neg_edge_index=None)

Given latent variables z, computes the binary cross entropy loss for positive edges pos_edge_index and negative sampled edges.

reset_parameters()

Resets all learnable parameters of the module.

test(z, pos_edge_index, neg_edge_index)

Given latent variables z, positive edges pos_edge_index and negative edges neg_edge_index, computes area under the ROC curve (AUC) and average precision (AP) scores.

VGAE

k3_node.models.VGAE

Bases: GAE

The Variational Graph Auto-Encoder model from the "Variational Graph Auto-Encoders" <https://arxiv.org/abs/1611.07308>_ paper.

Parameters:

Name Type Description Default
encoder

The encoder module to compute :math:\mu and :math:\log\sigma^2.

 required
decoder optional

The decoder module. If set to None, will default to InnerProductDecoder. (default: None)

 None

kl_loss(mu=None, logstd=None)

Computes the KL loss, either for the passed arguments mu and logstd, or based on latent variables from last encoding.

ARGA

k3_node.models.ARGA

Bases: GAE

The Adversarially Regularized Graph Auto-Encoder model from the "Adversarially Regularized Graph Autoencoder for Graph Embedding" <https://arxiv.org/abs/1802.04407>_ paper.

Parameters:

Name Type Description Default
encoder

The encoder module.

 required
discriminator

The discriminator module.

 required
decoder optional

The decoder module. If set to None, will default to InnerProductDecoder. (default: None)

 None

discriminator_loss(z)

Computes the loss of the discriminator.

reg_loss(z)

Computes the regularization loss of the encoder.

ARGVA

k3_node.models.ARGVA

Bases: ARGA

The Adversarially Regularized Variational Graph Auto-Encoder model from the "Adversarially Regularized Graph Autoencoder for Graph Embedding" <https://arxiv.org/abs/1802.04407>_ paper.

Parameters:

Name Type Description Default
encoder

The encoder module to compute :math:\mu and :math:\log\sigma^2.

 required
discriminator

The discriminator module.

 required
decoder optional

The decoder module. If set to None, will default to InnerProductDecoder. (default: None)

 None

InnerProductDecoder

k3_node.models.InnerProductDecoder

The inner product decoder from the "Variational Graph Auto-Encoders" <https://arxiv.org/abs/1611.07308>_ paper.

.. math:: \sigma(\mathbf{Z}\mathbf{Z}^{\top})

where :math:\mathbf{Z} \in \mathbb{R}^{N \times d} denotes the latent space produced by the encoder.

__call__(z, edge_index, sigmoid=True)

Decodes the latent variables z into edge probabilities for the given node-pairs edge_index.

forward_all(z, sigmoid=True)

Decodes the latent variables z into a probabilistic dense adjacency matrix.

DeepGraphInfomax

k3_node.models.DeepGraphInfomax

Bases: Layer

The Deep Graph Infomax model from the "Deep Graph Infomax" <https://arxiv.org/abs/1809.10341>_ paper based on user-defined encoder and summary model :math:\mathcal{E} and :math:\mathcal{R} respectively, and a corruption function :math:\mathcal{C}.

Parameters:

Name Type Description Default
hidden_channels int

The latent space dimensionality.

 required
encoder

The encoder module :math:\mathcal{E}.

 required
summary callable

The readout function :math:\mathcal{R}.

 required
corruption callable

The corruption function :math:\mathcal{C}.

 required

call(*args, **kwargs)

Returns the latent space for the input arguments, their corruptions and their summary representation.

discriminate(z, summary, sigmoid=True)

Given the patch-summary pair z and summary, computes the probability scores assigned to this patch-summary pair.

loss(pos_z, neg_z, summary)

Computes the mutual information maximization objective.

reset_parameters()

Resets all learnable parameters of the module.

test(train_z, train_y, test_z, test_y, solver='lbfgs', *args, **kwargs)

Evaluates latent space quality via a logistic regression downstream task.

Scalable Graph Transformers

Polynormer

k3_node.models.Polynormer

Bases: Layer

The Polynormer module from the "Polynormer: polynomial-expressive graph transformer in linear time" <https://arxiv.org/abs/2403.01232>_ paper.

Parameters:

Name Type Description Default
in_channels int

Input channels.

 required
hidden_channels int

Hidden channels.

 required
out_channels int

Output channels.

 required
local_layers int

The number of local attention layers. (default: :obj:7)

 7
global_layers int

The number of global attention layers. (default: :obj:2)

 2
in_dropout float

Input dropout rate. (default: :obj:0.15)

 0.15
dropout float

Dropout rate. (default: :obj:0.5)

 0.5
global_dropout float

Global dropout rate. (default: :obj:0.5)

 0.5
heads int

The number of heads. (default: :obj:1)

 1
beta float

Aggregate type. (default: :obj:0.9)

 0.9
qk_shared bool

Whether weight of query and key are shared. (default: :obj:True)

 False
pre_ln bool

Pre layer normalization. (default: :obj:False)

 False
post_bn bool

Post batch normalization. (default: :obj:True)

 True
local_attn bool

Whether use local attention (GATConv vs GCNConv). (default: :obj:False)

 False

call(x, edge_index, batch=None, training=None)

Forward pass.

Parameters:

Name Type Description Default
x Tensor

The input node features.

 required
edge_index Tensor

The edge indices.

 required
batch Tensor

The batch vector assigning each node to a graph. (default: :obj:None)

 None
training bool

Whether in training mode. (default: :obj:None)

 None

reset_parameters()

Resets all learnable parameters of the module.

SGFormer

k3_node.models.SGFormer

Bases: Model

The sgformer module from the "SGFormer: Simplifying and Empowering Transformers for Large-Graph Representations" <https://arxiv.org/abs/2306.10759>_ paper.

Parameters:

Name Type Description Default
in_channels int

Input channels.

 required
hidden_channels int

Hidden channels.

 required
out_channels int

Output channels.

 required
trans_num_layers int

The number of layers for all-pair attention. (default: :obj:2)

 2
trans_num_heads int

The number of heads for attention. (default: :obj:1)

 1
trans_dropout float

Global dropout rate. (default: :obj:0.5)

 0.5
gnn_num_layers int

The number of layers for GNN. (default: :obj:3)

 3
gnn_dropout float

GNN dropout rate. (default: :obj:0.5)

 0.5
graph_weight float

The weight balance global and gnn module. (default: :obj:0.5)

 0.5
aggregate str

Aggregate type (:obj:'add' or :obj:'cat'). (default: :obj:'add')

 'add'

LPFormer

k3_node.models.LPFormer

Bases: Layer

The LPFormer model from the "LPFormer: An Adaptive Graph Transformer for Link Prediction" <https://arxiv.org/abs/2310.11009>_ paper.

Parameters:

Name Type Description Default
in_channels int

Input feature dimension.

 required
hidden_channels int

Hidden dimension.

 required
num_gnn_layers int

Number of GCN layers. (default: 2)

 2
gnn_dropout float

GNN dropout rate. (default: 0.1)

 0.1
num_transformer_layers int

Number of Transformer layers. (default: 1)

 1
num_heads int

Number of attention heads. (default: 1)

 1
transformer_dropout float

Transformer dropout rate. (default: 0.1)

 0.1
ppr_thresholds list

Thresholds for PPR node categorization. (default: [0, 1e-4, 1e-2])

 None

call(batch, x, edge_index, ppr_matrix=None, training=False)

Forward pass of LPFormer.

Parameters:

Name Type Description Default
batch

Tensor of shape (2, B) with link pairs (u, v) to predict.

 required
x

Node features of shape (N, in_channels).

 required
edge_index

Graph edge index (2, E).

 required
ppr_matrix

Optional precomputed PPR matrix of shape (N, N).

 None

GPSE

k3_node.models.GPSE

Bases: Layer

The Graph Positional and Structural Encoder (GPSE) model from the "Graph Positional and Structural Encoder" <https://arxiv.org/abs/2307.07107>_ paper.

Molecular & Physical Models

AttentiveFP

k3_node.models.AttentiveFP

Bases: Layer

The Attentive FP model for molecular representation learning from the "Pushing the Boundaries of Molecular Representation for Drug Discovery with the Graph Attention Mechanism" <https://pubs.acs.org/doi/10.1021/acs.jmedchem.9b00959>_ paper, based on graph attention mechanisms.

Parameters:

Name Type Description Default
in_channels int

Size of each input sample.

 required
hidden_channels int

Hidden node feature dimensionality.

 required
out_channels int

Size of each output sample.

 required
edge_dim int

Edge feature dimensionality.

 required
num_layers int

Number of GNN layers.

 required
num_timesteps int

Number of iterative refinement steps for global readout.

 required
dropout float

Dropout probability. (default: 0.0)

 0.0

NeuralFingerprint

k3_node.models.NeuralFingerprint

Bases: Layer

The Neural Fingerprint model from the "Convolutional Networks on Graphs for Learning Molecular Fingerprints" <https://arxiv.org/abs/1509.09292>__ paper to generate fingerprints of molecules.

Parameters:

Name Type Description Default
in_channels int

Size of each input sample.

 required
hidden_channels int

Size of each hidden sample.

 required
out_channels int

Size of each output fingerprint.

 required
num_layers int

Number of layers.

 required
**kwargs optional

Additional arguments of :class:~k3_node.layers.conv.MFConv.

 {}

reset_parameters()

Resets all learnable parameters of the module.

SchNet

k3_node.models.SchNet

Bases: Layer

The continuous-filter convolutional neural network SchNet from the "SchNet: A Continuous-filter Convolutional Neural Network for Modeling Quantum Interactions" <https://arxiv.org/abs/1706.08566>_ paper.

Parameters:

Name Type Description Default
hidden_channels int

Hidden embedding size. (default: 128)

 128
num_filters int

The number of filters to use. (default: 128)

 128
num_interactions int

The number of interaction blocks. (default: 6)

 6
num_gaussians int

The number of gaussians. (default: 50)

 50
cutoff float

Cutoff distance. (default: 10.0)

 10.0
interaction_graph callable

Interaction graph builder. (default: None)

 None
max_num_neighbors int

Maximum neighbors per atom. (default: 32)

 32
readout str

Readout pooling (add, sum, mean). (default: "add")

 'add'
dipole bool

Predict dipole moment magnitude. (default: False)

 False
mean float

Mean of target property. (default: None)

 None
std float

Standard deviation of target property. (default: None)

 None
atomref tensor

Reference atomic values. (default: None)

 None

DimeNet

k3_node.models.DimeNet

Bases: Layer

The directional message passing neural network (DimeNet) from the "Directional Message Passing for Molecular Graphs" <https://arxiv.org/abs/2003.03123>_ paper.

DimeNetPlusPlus

k3_node.models.DimeNetPlusPlus

Bases: DimeNet

The DimeNet++ from the "Fast and Uncertainty-Aware Directional Message Passing for Non-Equilibrium Molecules" <https://arxiv.org/abs/2011.14115>_ paper.

GNNFF

k3_node.models.GNNFF

Bases: Layer

The Graph Neural Network Force Field (GNNFF) from the "Accurate and scalable graph neural network force field and molecular dynamics with direct force architecture" <https://www.nature.com/articles/s41524-021-00543-3>_ paper. :class:GNNFF directly predicts atomic forces from automatically extracted features of the local atomic environment.

Parameters:

Name Type Description Default
hidden_node_channels int

Hidden node embedding size.

 required
hidden_edge_channels int

Hidden edge embedding size.

 required
num_layers int

Number of message passing blocks.

 required
cutoff float

Cutoff distance. (default: 5.0)

 5.0
max_num_neighbors int

Maximum neighbors per node. (default: 32)

 32

ViSNet

k3_node.models.ViSNet

Bases: Layer

The equivariant vector-scalar interactive graph neural network (ViSNet) from the "Enhancing Geometric Representations for Molecules with Equivariant Vector-Scalar Interactive Message Passing" <https://arxiv.org/abs/2210.16518>_ paper.

Semi-Supervised & Non-Homophilous Models

LabelPropagation

k3_node.models.LabelPropagation

Bases: MessagePassing

The label propagation operator, firstly introduced in the "Learning from Labeled and Unlabeled Data with Label Propagation" <http://mlg.eng.cam.ac.uk/zoubin/papers/CMU-CALD-02-107.pdf>_ paper.

.. math:: \mathbf{Y}^{\prime} = \alpha \cdot \mathbf{D}^{-1/2} \mathbf{A} \mathbf{D}^{-1/2} \mathbf{Y} + (1 - \alpha) \mathbf{Y},

where unlabeled data is inferred by labeled data via propagation. This concrete implementation here is derived from the "Combining Label Propagation And Simple Models Out-performs Graph Neural Networks" <https://arxiv.org/abs/2010.13993>_ paper.

Parameters:

Name Type Description Default
num_layers int

The number of propagations.

 required
alpha float

The :math:\alpha coefficient.

 required

CorrectAndSmooth

k3_node.models.CorrectAndSmooth

Bases: Layer

The correct and smooth (C&S) post-processing model from the "Combining Label Propagation And Simple Models Out-performs Graph Neural Networks" <https://arxiv.org/abs/2010.13993>_ paper.

Parameters:

Name Type Description Default
num_correction_layers int

The number of propagations :math:L_1.

 required
correction_alpha float

The :math:\alpha_1 coefficient.

 required
num_smoothing_layers int

The number of propagations :math:L_2.

 required
smoothing_alpha float

The :math:\alpha_2 coefficient.

 required
autoscale bool

If set to :obj:True, will automatically determine the scaling factor :math:\gamma. (default: :obj:True)

 True
scale float

The scaling factor :math:\gamma, in case :obj:autoscale = False. (default: :obj:1.0)

 1.0

RECT_L

k3_node.models.RECT_L

Bases: Layer

The RECT model, i.e. its supervised RECT-L part, from the "Network Embedding with Completely-imbalanced Labels" <https://arxiv.org/abs/2007.03545>_ paper.

Parameters:

Name Type Description Default
in_channels int

Size of each input sample.

 required
hidden_channels int

Intermediate size of each sample.

 required
normalize bool

Whether to add self-loops and compute symmetric normalization coefficients on-the-fly. (default: :obj:True)

 True
dropout float

The dropout probability. (default: :obj:0.0)

 0.0

get_semantic_labels(x, y, mask)

Replaces the original labels by their class-centers.

reset_parameters()

Resets all learnable parameters of the module.

PMLP

k3_node.models.PMLP

Bases: Layer

The P(ropagational)MLP model from the "Graph Neural Networks are Inherently Good Generalizers: Insights by Bridging GNNs and MLPs" <https://arxiv.org/abs/2212.09034>_ paper.

:class:PMLP is identical to a standard MLP during training, but then adopts a GNN architecture during testing.

Parameters:

Name Type Description Default
in_channels int

Size of each input sample.

 required
hidden_channels int

Size of each hidden sample.

 required
out_channels int

Size of each output sample.

 required
num_layers int

The number of layers.

 required
dropout float

Dropout probability of each hidden embedding. (default: :obj:0.)

 0.0
norm bool

If set to :obj:False, will not apply batch normalization. (default: :obj:True)

 True
bias bool

If set to :obj:False, the module will not learn additive biases. (default: :obj:True)

 True

call(x, edge_index=None, training=None)

Forward pass.

Parameters:

Name Type Description Default
x Tensor

The node features of shape [N, in_channels].

 required
edge_index Tensor

The edge indices. Required during inference. (default: :obj:None)

 None
training bool

Override the instance-level self.training flag. (default: :obj:None)

 None

reset_parameters()

Resets all learnable parameters of the module.

LINKX

k3_node.models.LINKX

Bases: Layer

The LINKX model from the "Large Scale Learning on Non-Homophilous Graphs: New Benchmarks and Strong Simple Methods" <https://arxiv.org/abs/2110.14446>_ paper.

Parameters:

Name Type Description Default
num_nodes int

The number of nodes in the graph.

 required
in_channels int

Size of each input sample.

 required
hidden_channels int

Size of each hidden sample.

 required
out_channels int

Size of each output sample.

 required
num_layers int

Number of layers of :math:\textrm{MLP}_{f}.

 required
num_edge_layers int

Number of layers of :math:\textrm{MLP}_{\mathbf{A}}. (default: :obj:1)

 1
num_node_layers int

Number of layers of :math:\textrm{MLP}_{\mathbf{X}}. (default: :obj:1)

 1
dropout float

Dropout probability. (default: :obj:0.0)

 0.0

reset_parameters()

Resets all learnable parameters of the module.

ARLinkPredictor

k3_node.models.ARLinkPredictor

Bases: Layer

Link predictor using Attract-Repel embeddings from the paper "Pseudo-Euclidean Attract-Repel Embeddings for Undirected Graphs" <https://arxiv.org/abs/2106.09671>_.

This model splits node embeddings into: attract and repel. The edge prediction score is computed as the dot product of attract components minus the dot product of repel components.

Parameters:

Name Type Description Default
in_channels int

Size of each input sample.

 required
hidden_channels int

Size of hidden embeddings.

 required
out_channels int

Size of output embeddings. If set to None, will default to hidden_channels. (default: None)

 None
num_layers int

Number of message passing layers. (default: 2)

 2
dropout float

Dropout probability. (default: 0.0)

 0.0
attract_ratio float

Ratio to use for attract component. Must be between 0 and 1. (default: 0.5)

 0.5

calculate_r_fraction(attract_z, repel_z)

Calculate the R-fraction (proportion of energy in repel space).

decode(attract_z, repel_z, edge_index)

Decode edge scores from attract-repel embeddings.

encode(x, *args, **kwargs)

Encode node features into attract-repel embeddings.

Recommendation & Temporal Models

LightGCN

k3_node.models.LightGCN

Bases: Layer

The LightGCN model from the "LightGCN: Simplifying and Powering Graph Convolution Network for Recommendation" <https://arxiv.org/abs/2002.02126>_ paper.

Parameters:

Name Type Description Default
num_nodes int

The number of nodes in the graph.

 required
embedding_dim int

The dimensionality of node embeddings.

 required
num_layers int

The number of :class:LGConv layers.

 required
alpha float or Tensor

The scalar or vector specifying the re-weighting coefficients for aggregating the final embedding. (default: :obj:None)

 None

call(edge_index, edge_label_index=None, edge_weight=None)

Computes rankings for pairs of nodes.

get_embedding(edge_index, edge_weight=None)

Returns the embedding of nodes in the graph.

reset_parameters()

Resets all learnable parameters of the module.

BPRLoss

k3_node.models.BPRLoss

The Bayesian Personalized Ranking (BPR) loss.

SignedGCN

k3_node.models.SignedGCN

Bases: Layer

The signed graph convolutional network model from the "Signed Graph Convolutional Network" <https://arxiv.org/abs/1808.06354>_ paper.

Parameters:

Name Type Description Default
in_channels int

Size of each input sample.

 required
hidden_channels int

Size of each hidden sample.

 required
num_layers int

Number of layers.

 required
lamb float

Balances the contributions of the overall objective. (default: :obj:5)

 5.0
bias bool

If set to :obj:False, all layers will not learn an additive bias. (default: :obj:True)

 True

TGNMemory

k3_node.models.TGNMemory

Bases: Layer

The Temporal Graph Network (TGN) memory model from the "Temporal Graph Networks for Deep Learning on Dynamic Graphs" <https://arxiv.org/abs/2006.10637>_ paper.

Parameters:

Name Type Description Default
num_nodes int

The number of nodes to save memories for.

 required
raw_msg_dim int

The raw message dimensionality.

 required
memory_dim int

The hidden memory dimensionality.

 required
time_dim int

The time encoding dimensionality.

 required
message_module Callable

Function combining source and destination node memory, raw message, and time encoding.

 required
aggregator_module Callable

Function aggregating messages to the same destination into a single representation.

 required

RENet

k3_node.models.RENet

Bases: Model

The Recurrent Event Network model from the "Recurrent Event Network for Reasoning over Temporal Knowledge Graphs" <https://arxiv.org/abs/1904.05530>_ paper.

Parameters:

Name Type Description Default
num_nodes int

The number of nodes in the knowledge graph.

 required
num_rels int

The number of relations in the knowledge graph.

 required
hidden_channels int

Hidden size of node and relation embeddings.

 required
seq_len int

The sequence length of past events.

 required
num_layers int

The number of recurrent layers. (default: :obj:1)

 1
dropout float

Dropout rate before final prediction. (default: :obj:0.0)

 0.0
bias bool

If set to :obj:False, all layers will not learn an additive bias. (default: :obj:True)

 True

test(logits, y)

Given ground-truth :obj:y, computes Mean Reciprocal Rank (MRR) and Hits at 1/3/10.

Graph Embedding Models

Node2Vec

k3_node.models.Node2Vec

Bases: Layer

The Node2Vec model from the "node2vec: Scalable Feature Learning for Networks" <https://arxiv.org/abs/1607.00653>_ paper where random walks of length :obj:walk_length are sampled in a given graph, and node embeddings are learned via negative sampling optimization.

Parameters:

Name Type Description Default
edge_index

The edge indices.

 required
embedding_dim int

The size of each embedding vector.

 required
walk_length int

The walk length.

 required
context_size int

The actual context size which is considered for positive samples.

 required
walks_per_node int

The number of walks to sample for each node. (default: :obj:1)

 1
p float

Likelihood of immediately revisiting a node in the walk. (default: :obj:1.0)

 1.0
q float

Control parameter to interpolate between breadth-first strategy and depth-first strategy. (default: :obj:1.0)

 1.0
num_negative_samples int

The number of negative samples to use for each positive sample. (default: :obj:1)

 1
num_nodes int

The number of nodes. (default: :obj:None)

 None

call(batch=None)

Returns the embeddings for the nodes in :obj:batch.

loss(pos_rw, neg_rw)

Computes the loss given positive and negative random walks.

MetaPath2Vec

k3_node.models.MetaPath2Vec

Bases: Layer

The MetaPath2Vec model from the "metapath2vec: Scalable Representation Learning for Heterogeneous Networks" <https://ericdongyx.github.io/papers/ KDD17-dong-chawla-swami-metapath2vec.pdf>_ paper where random walks based on a given :obj:metapath are sampled in a heterogeneous graph, and node embeddings are learned via negative sampling optimization.

Parameters:

Name Type Description Default
edge_index_dict Dict[Tuple[str, str, str], Tensor]

Dictionary holding edge indices for each edge type.

 required
embedding_dim int

The size of each embedding vector.

 required
metapath List[Tuple[str, str, str]]

The sequence of edge types denoting the metapath.

 required
walk_length int

The walk length.

 required
context_size int

The context size considered for positive samples.

 required
walks_per_node int

The number of walks to sample for each node. (default: :obj:1)

 1
num_negative_samples int

The number of negative samples. (default: :obj:1)

 1
num_nodes_dict Dict[str, int]

The number of nodes for each node type. (default: :obj:None)

 None

call(node_type, batch=None)

Returns the embeddings for the nodes in :obj:batch of type :obj:node_type.

loss(pos_rw, neg_rw)

Computes the loss given positive and negative random walks.

Architectural Blocks

DeepGCNLayer

k3_node.models.DeepGCNLayer

Bases: Layer

The skip connection operations from the "DeepGCNs: Can GCNs Go as Deep as CNNs?" <https://arxiv.org/abs/1904.03751> and "All You Need to Train Deeper GCNs" <https://arxiv.org/abs/2006.07739> papers. The implemented skip connections includes the pre-activation residual connection ("res+"), the residual connection ("res"), the dense connection ("dense") and no connections ("plain").

  • Res+ ("res+"):

.. math:: \text{Normalization}\to\text{Activation}\to\text{Dropout}\to \text{GraphConv}\to\text{Res}

  • Res ("res") / Dense ("dense") / Plain ("plain"):

.. math:: \text{GraphConv}\to\text{Normalization}\to\text{Activation}\to \text{Res/Dense/Plain}\to\text{Dropout}

Parameters:

Name Type Description Default
conv optional

the GCN operator. (default: None)

 None
norm optional

the normalization layer. (default: None)

 None
act optional

the activation layer. (default: None)

 None
block str

The skip connection operation to use ("res+", "res", "dense" or "plain"). (default: "res+")

 'res+'
dropout float

Whether to apply dropout. (default: 0.)

 0.0
ckpt_grad bool

Kept for API compatibility; a no-op, since there is no gradient-checkpointing API shared uniformly across Keras backends. (default: False)

 False

reset_parameters()

Resets all learnable parameters of the module.

GroupAddRev

k3_node.models.GroupAddRev

Bases: Layer

The Grouped Reversible GNN module from the "Graph Neural Networks with 1000 Layers" <https://arxiv.org/abs/2106.07476>_ paper.

Parameters:

Name Type Description Default
conv Layer or List[Layer]

A seed GNN layer or list of GNN layers.

 required
split_dim int

The dimension across which to split groups. (default: :obj:-1)

 -1
num_groups int

The number of groups. (default: :obj:None)

 None

JumpingKnowledge

k3_node.models.JumpingKnowledge

Bases: Layer

The Jumping Knowledge layer aggregation module from the "Representation Learning on Graphs with Jumping Knowledge Networks" <https://arxiv.org/abs/1806.03536>_ paper.

Jumping knowledge is performed based on either concatenation (:obj:"cat")

.. math::

\mathbf{x}_v^{(1)} \, \Vert \, \ldots \, \Vert \, \mathbf{x}_v^{(T)},

max pooling (:obj:"max")

.. math::

\max \left( \mathbf{x}_v^{(1)}, \ldots, \mathbf{x}_v^{(T)} \right),

or weighted summation

.. math::

\sum_{t=1}^T \alpha_v^{(t)} \mathbf{x}_v^{(t)}

with attention scores :math:\alpha_v^{(t)} obtained from a bi-directional LSTM (:obj:"lstm").

Parameters:

Name Type Description Default
mode str

The aggregation scheme to use (:obj:"cat", :obj:"max" or :obj:"lstm").

 required
channels int

The number of channels per representation. Needs to be only set for LSTM-style aggregation. (default: :obj:None)

 None
num_layers int

The number of layers to aggregate. Needs to be only set for LSTM-style aggregation. (default: :obj:None)

 None

call(xs)

Forward pass.

Parameters:

Name Type Description Default
xs List[Tensor]

List containing the layer-wise representations.

 required

reset_parameters()

Resets all learnable parameters of the module.

MetaLayer

k3_node.models.MetaLayer

Bases: Layer

A meta layer for building any kind of graph network, inspired by the "Relational Inductive Biases, Deep Learning, and Graph Networks" <https://arxiv.org/abs/1806.01261>_ paper.

A graph network takes a graph as input and returns an updated graph as output (with same connectivity). The input graph has node features :obj:x, edge features :obj:edge_attr as well as graph-level features :obj:u. The output graph has the same structure, but updated features.

Edge features, node features as well as global features are updated by calling the modules :obj:edge_model, :obj:node_model and :obj:global_model, respectively.

To allow for batch-wise graph processing, all callable functions take an additional argument :obj:batch, which determines the assignment of edges or nodes to their specific graphs.

Parameters:

Name Type Description Default
edge_model callable

A callable which updates a graph's edge features based on its source and target node features, its current edge features and its global features. (default: :obj:None)

 None
node_model callable

A callable which updates a graph's node features based on its current node features, its graph connectivity, its edge features and its global features. (default: :obj:None)

 None
global_model callable

A callable which updates a graph's global features based on its node features, its graph connectivity, its edge features and its current global features. (default: :obj:None)

 None

Example::

from keras import layers
from k3_node.models import MetaLayer
from k3_node.layers.conv.utils import scatter

class EdgeModel(layers.Layer):
    def __init__(self):
        super().__init__()
        self.mlp = layers.Dense(5)
    def call(self, src, dst, edge_attr, u, batch):
        out = ops.concatenate([src, dst, edge_attr, u[batch]], axis=1)
        return self.mlp(out)

class NodeModel(layers.Layer):
    def __init__(self):
        super().__init__()
        self.mlp1 = layers.Dense(10)
        self.mlp2 = layers.Dense(10)
    def call(self, x, edge_index, edge_attr, u, batch):
        row, col = edge_index[0], edge_index[1]
        out = ops.concatenate([x[row], edge_attr], axis=1)
        out = scatter(self.mlp1(out), col, dim_size=ops.shape(x)[0])
        out = ops.concatenate([x, out, u[batch]], axis=1)
        return self.mlp2(out)

class GlobalModel(layers.Layer):
    def __init__(self):
        super().__init__()
        self.mlp = layers.Dense(20)
    def call(self, x, edge_index, edge_attr, u, batch):
        out = ops.concatenate([u, scatter(x, batch)], axis=1)
        return self.mlp(out)

op = MetaLayer(EdgeModel(), NodeModel(), GlobalModel())
x, edge_attr, u = op(x, edge_index, edge_attr, u, batch)

call(x, edge_index, edge_attr=None, u=None, batch=None)

Forward pass.

Parameters:

Name Type Description Default
x Tensor

The node features of shape [N, F_x].

 required
edge_index Tensor

The edge indices of shape [2, E].

 required
edge_attr Tensor

The edge features of shape [E, F_e]. (default: :obj:None)

 None
u Tensor

The global graph features of shape [B, F_u]. (default: :obj:None)

 None
batch Tensor

The batch vector :math:\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N. (default: :obj:None)

 None

reset_parameters()

Resets all learnable parameters of the module.