k3-node 1.0.0__py3-none-any.whl

k3_node/models/jumping_knowledge.py

@@ -0,0 +1,200 @@
+from typing import Dict, List, Optional
+
+import keras
+from keras import ops
+
+
+class JumpingKnowledge(keras.layers.Layer):
+    r"""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"`).
+
+    Args:
+        mode (str): The aggregation scheme to use
+            (:obj:`"cat"`, :obj:`"max"` or :obj:`"lstm"`).
+        channels (int, optional): The number of channels per representation.
+            Needs to be only set for LSTM-style aggregation.
+            (default: :obj:`None`)
+        num_layers (int, optional): The number of layers to aggregate. Needs to
+            be only set for LSTM-style aggregation. (default: :obj:`None`)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.models import JumpingKnowledge
+
+        # Node representations from 3 GNN layers
+        xs = [np.random.rand(10, 16).astype("float32") for _ in range(3)]
+        print(tuple(JumpingKnowledge("cat")(xs).shape))  # (10, 48): concatenate all layers
+        print(tuple(JumpingKnowledge("max")(xs).shape))  # (10, 16): element-wise max
+        print(tuple(JumpingKnowledge("lstm", channels=16, num_layers=3)(xs).shape))  # (10, 16): attention over layers
+        ```
+    """
+
+    def __init__(
+        self,
+        mode: str,
+        channels: Optional[int] = None,
+        num_layers: Optional[int] = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.mode = mode.lower()
+        assert self.mode in ['cat', 'max', 'lstm'], \
+            f"mode must be 'cat', 'max', or 'lstm', got '{mode}'"
+
+        self.channels = channels
+        self.num_layers = num_layers
+
+        if self.mode == 'lstm':
+            assert channels is not None, 'channels cannot be None for lstm'
+            assert num_layers is not None, 'num_layers cannot be None for lstm'
+            lstm_units = (num_layers * channels) // 2
+            self.lstm = keras.layers.Bidirectional(
+                keras.layers.LSTM(lstm_units, return_sequences=True),
+                merge_mode='concat',
+            )
+            self.att = keras.layers.Dense(1, use_bias=False)
+        else:
+            self.lstm = None
+            self.att = None
+
+    def reset_parameters(self) -> None:
+        r"""Resets all learnable parameters of the module."""
+        # Keras layers reinitialize on next call; explicit rebuild if needed
+        if self.lstm is not None and self.lstm.built:
+            for layer in self.lstm.layers:
+                if hasattr(layer, 'kernel') and layer.kernel is not None:
+                    layer.kernel.assign(
+                        keras.initializers.GlorotUniform()(layer.kernel.shape)
+                    )
+                if hasattr(layer, 'recurrent_kernel') and layer.recurrent_kernel is not None:
+                    layer.recurrent_kernel.assign(
+                        keras.initializers.Orthogonal()(layer.recurrent_kernel.shape)
+                    )
+                if hasattr(layer, 'bias') and layer.bias is not None:
+                    layer.bias.assign(ops.zeros(layer.bias.shape))
+        if self.att is not None and self.att.built:
+            if self.att.kernel is not None:
+                self.att.kernel.assign(
+                    keras.initializers.GlorotUniform()(self.att.kernel.shape)
+                )
+
+    def call(self, xs: List) -> object:
+        r"""Forward pass.
+
+        Args:
+            xs (List[Tensor]): List containing the layer-wise representations.
+        """
+        if self.mode == 'cat':
+            return ops.concatenate(xs, axis=-1)
+        elif self.mode == 'max':
+            return ops.max(ops.stack(xs, axis=-1), axis=-1)
+        else:  # lstm
+            assert self.lstm is not None and self.att is not None
+            x = ops.stack(xs, axis=1)  # [num_nodes, num_layers, num_channels]
+            alpha = self.lstm(x)        # [num_nodes, num_layers, 2*lstm_units]
+            alpha = self.att(alpha)     # [num_nodes, num_layers, 1]
+            alpha = ops.squeeze(alpha, axis=-1)  # [num_nodes, num_layers]
+            alpha = ops.softmax(alpha, axis=-1)
+            return ops.sum(x * ops.expand_dims(alpha, axis=-1), axis=1)
+
+    def __repr__(self) -> str:
+        if self.mode == 'lstm':
+            return (f'{self.__class__.__name__}({self.mode}, '
+                    f'channels={self.channels}, layers={self.num_layers})')
+        return f'{self.__class__.__name__}({self.mode})'
+
+
+class HeteroJumpingKnowledge(keras.layers.Layer):
+    r"""A heterogeneous version of the :class:`JumpingKnowledge` module.
+
+    Args:
+        types (List[str]): The keys of the input dictionary.
+        mode (str): The aggregation scheme to use
+            (:obj:`"cat"`, :obj:`"max"` or :obj:`"lstm"`).
+        channels (int, optional): The number of channels per representation.
+            Needs to be only set for LSTM-style aggregation.
+            (default: :obj:`None`)
+        num_layers (int, optional): The number of layers to aggregate. Needs to
+            be only set for LSTM-style aggregation. (default: :obj:`None`)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.models import HeteroJumpingKnowledge
+
+        # Per node type: representations from 3 GNN layers
+        xs_dict = {
+            "author": [np.random.rand(3, 16).astype("float32") for _ in range(3)],
+            "paper": [np.random.rand(4, 16).astype("float32") for _ in range(3)],
+        }
+        model = HeteroJumpingKnowledge(["author", "paper"], mode="cat")
+        out_dict = model(xs_dict)
+        print(tuple(out_dict["author"].shape), tuple(out_dict["paper"].shape))  # (3, 48) (4, 48)
+        ```
+    """
+
+    def __init__(
+        self,
+        types: List[str],
+        mode: str,
+        channels: Optional[int] = None,
+        num_layers: Optional[int] = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.mode = mode.lower()
+        self.types = list(types)
+
+        self.jk_dict = {
+            key: JumpingKnowledge(mode, channels, num_layers)
+            for key in types
+        }
+
+    def reset_parameters(self) -> None:
+        r"""Resets all learnable parameters of the module."""
+        for jk in self.jk_dict.values():
+            jk.reset_parameters()
+
+    def call(self, xs_dict: Dict[str, List]) -> Dict[str, object]:
+        r"""Forward pass.
+
+        Args:
+            xs_dict (Dict[str, List[Tensor]]): A dictionary holding a
+                list of layer-wise representation for each type.
+        """
+        return {key: self.jk_dict[key](xs_dict[key]) for key in self.types}
+
+    def __repr__(self) -> str:
+        if self.mode == 'lstm':
+            jk = next(iter(self.jk_dict.values()))
+            return (f'{self.__class__.__name__}('
+                    f'num_types={len(self.jk_dict)}, '
+                    f'mode={self.mode}, channels={jk.channels}, '
+                    f'layers={jk.num_layers})')
+        return (f'{self.__class__.__name__}(num_types={len(self.jk_dict)}, '
+                f'mode={self.mode})')
+

k3_node/models/label_prop.py

@@ -0,0 +1,110 @@
+from typing import Callable, Optional
+
+import keras
+from keras import ops
+
+from k3_node.layers.conv.message_passing import MessagePassing
+from k3_node.layers.conv.utils import gcn_norm
+
+
+class LabelPropagation(MessagePassing):
+    r"""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.
+
+    Args:
+        num_layers (int): The number of propagations.
+        alpha (float): The :math:`\alpha` coefficient.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.models import LabelPropagation
+
+        y = np.array([0, 1, 2, 0, 1, 2])  # node labels
+        train_mask = np.array([True, True, True, False, False, False])  # labels known for 3 nodes
+        edge_index = np.array([[0, 1, 2, 3, 4, 5], [3, 4, 5, 0, 1, 2]])
+        model = LabelPropagation(num_layers=3, alpha=0.9)
+        out = model(y, edge_index, train_mask)  # soft labels for every node
+        print(tuple(out.shape))  # (6, 3)
+        ```
+    """
+    def __init__(self, num_layers: int, alpha: float, **kwargs):
+        super().__init__(aggr='sum', **kwargs)
+        self.num_layers = num_layers
+        self.alpha = alpha
+
+    def build(self, input_shape=None):
+        self.built = True
+
+    # `mask` selects the labeled nodes; it is not a Keras sequence mask, so the output gets none
+    # (this also stops Keras from warning that the layer drops the mask).
+    supports_masking = True
+
+    def compute_mask(self, *args, **kwargs):
+        return None
+
+    def call(
+        self,
+        y,
+        edge_index,
+        mask=None,
+        edge_weight=None,
+        post_step: Optional[Callable] = None,
+    ):
+        shape = ops.shape(y)
+        if len(shape) == 1:
+            num_classes = int(ops.max(y)) + 1
+            y = ops.one_hot(y, num_classes)
+
+        y = ops.cast(y, dtype="float32")
+        out = y
+        if mask is not None:
+            mask_shape = ops.shape(mask)
+            # standardize_dtype: torch.bool does not compare equal to the string "bool"
+            if len(mask_shape) == 1 and keras.backend.standardize_dtype(mask.dtype) == "bool":
+                mask_expanded = ops.expand_dims(mask, axis=-1)
+                out = ops.where(mask_expanded, y, ops.zeros_like(y))
+            else:
+                out_zeros = ops.zeros_like(y)
+                out = ops.scatter_update(out_zeros, ops.expand_dims(mask, -1), ops.take(y, mask, axis=0))
+
+        if edge_weight is None:
+            num_nodes = ops.shape(y)[0]
+            edge_index, edge_weight = gcn_norm(
+                edge_index,
+                edge_weight=None,
+                num_nodes=num_nodes,
+                add_self_loops=False,
+                dtype=y.dtype,
+            )
+
+        res = (1.0 - self.alpha) * out
+        for _ in range(self.num_layers):
+            out = self.propagate(edge_index, x=out, edge_weight=edge_weight)
+            out = self.alpha * out + res
+            if post_step is not None:
+                out = post_step(out)
+            else:
+                out = ops.clip(out, 0.0, 1.0)
+
+        return out
+
+    def message(self, x_j, edge_weight=None):
+        if edge_weight is None:
+            return x_j
+        return ops.expand_dims(edge_weight, axis=-1) * x_j
+
+    def __repr__(self) -> str:
+        return (f'{self.__class__.__name__}(num_layers={self.num_layers}, '
+                f'alpha={self.alpha})')
+

k3_node/models/lightgcn.py

@@ -0,0 +1,171 @@
+from typing import Optional, Union
+
+import keras
+from keras import ops
+import numpy as np
+
+from k3_node.layers.conv import LGConv
+
+
+class BPRLoss:
+    r"""The Bayesian Personalized Ranking (BPR) loss."""
+    def __init__(self, lambda_reg: float = 0.0):
+        self.lambda_reg = lambda_reg
+
+    def __call__(self, positives, negatives, parameters=None):
+        diff = positives - negatives
+        # log(sigmoid(x)) = -log(1 + exp(-x)) or ops.log_sigmoid if available
+        log_prob = ops.mean(-ops.softplus(-diff))
+
+        regularization = 0.0
+        if self.lambda_reg != 0.0 and parameters is not None:
+            regularization = self.lambda_reg * ops.sum(ops.square(parameters))
+            regularization = regularization / ops.cast(ops.shape(positives)[0], "float32")
+
+        return -log_prob + regularization
+
+
+class LightGCN(keras.Model):
+    r"""The LightGCN model from the `"LightGCN: Simplifying and Powering
+    Graph Convolution Network for Recommendation"
+    <https://arxiv.org/abs/2002.02126>`_ paper.
+
+    Args:
+        num_nodes (int): The number of nodes in the graph.
+        embedding_dim (int): The dimensionality of node embeddings.
+        num_layers (int): The number of :class:`LGConv` layers.
+        alpha (float or Tensor, optional): The scalar or vector specifying
+            the re-weighting coefficients for aggregating the final embedding.
+            (default: :obj:`None`)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.models import LightGCN
+
+        edge_index = np.array([[0, 1, 2, 3, 4, 5, 6, 7], [1, 2, 3, 4, 5, 6, 7, 0]])  # user-item interactions
+        edge_label_index = np.array([[0, 1, 2, 3], [4, 5, 6, 7]])  # pairs to score
+        model = LightGCN(num_nodes=50, embedding_dim=16, num_layers=2)
+        scores = model(edge_index, edge_label_index)
+        print(tuple(scores.shape))  # (4,)
+        print(tuple(model.recommend(edge_index, k=2).shape))  # (50, 2): top-2 recommendations per node
+        ```
+    """
+    def __init__(
+        self,
+        num_nodes: int,
+        embedding_dim: int,
+        num_layers: int,
+        alpha: Optional[Union[float, list]] = None,
+        **kwargs,
+    ):
+        super().__init__()
+
+        self.num_nodes = num_nodes
+        self.embedding_dim = embedding_dim
+        self.num_layers = num_layers
+
+        if alpha is None:
+            alpha = [1.0 / (num_layers + 1)] * (num_layers + 1)
+        elif isinstance(alpha, (int, float)):
+            alpha = [float(alpha)] * (num_layers + 1)
+        self.alpha_list = list(alpha)
+
+        self.embedding = keras.layers.Embedding(
+            num_nodes,
+            embedding_dim,
+            embeddings_initializer=keras.initializers.GlorotUniform(),
+        )
+        self.convs = [LGConv(**kwargs) for _ in range(num_layers)]
+
+    def build(self, input_shape=None):
+        self.embedding.build((None,))
+        self.built = True
+
+    def reset_parameters(self):
+        r"""Resets all learnable parameters of the module."""
+        if self.embedding.built:
+            self.embedding.embeddings.assign(
+                keras.initializers.GlorotUniform()(shape=(self.num_nodes, self.embedding_dim))
+            )
+        for conv in self.convs:
+            if hasattr(conv, "reset_parameters"):
+                conv.reset_parameters()
+
+    def get_embedding(self, edge_index, edge_weight=None):
+        r"""Returns the embedding of nodes in the graph."""
+        if not self.embedding.built:
+            self.embedding.build((None,))
+        # Embedding weights: shape [num_nodes, embedding_dim]
+        x = self.embedding.weights[0]
+        out = x * self.alpha_list[0]
+
+        for i in range(self.num_layers):
+            x = self.convs[i](x, edge_index, edge_weight=edge_weight)
+            out = out + x * self.alpha_list[i + 1]
+
+        return out
+
+    def call(self, edge_index, edge_label_index=None, edge_weight=None):
+        r"""Computes rankings for pairs of nodes."""
+        if edge_label_index is None:
+            edge_label_index = edge_index
+
+        out = self.get_embedding(edge_index, edge_weight)
+
+        out_src = ops.take(out, edge_label_index[0], axis=0)
+        out_dst = ops.take(out, edge_label_index[1], axis=0)
+
+        return ops.sum(out_src * out_dst, axis=-1)
+
+    def predict_link(
+        self,
+        edge_index,
+        edge_label_index=None,
+        edge_weight=None,
+        prob: bool = False,
+    ):
+        pred = ops.sigmoid(self(edge_index, edge_label_index, edge_weight))
+        return pred if prob else ops.round(pred)
+
+    def recommend(
+        self,
+        edge_index,
+        edge_weight=None,
+        src_index=None,
+        dst_index=None,
+        k: int = 1,
+        sorted: bool = True,
+    ):
+        out = self.get_embedding(edge_index, edge_weight)
+        out_src = ops.take(out, src_index, axis=0) if src_index is not None else out
+        out_dst = ops.take(out, dst_index, axis=0) if dst_index is not None else out
+
+        pred = out_src @ ops.transpose(out_dst)
+        top_indices = ops.top_k(pred, k=k, sorted=sorted)[1]
+
+        if dst_index is not None:
+            top_indices = ops.take(dst_index, top_indices, axis=0)
+
+        return top_indices
+
+    def link_pred_loss(self, pred, edge_label):
+        loss_fn = keras.losses.BinaryCrossentropy(from_logits=True)
+        return loss_fn(edge_label, pred)
+
+    def recommendation_loss(
+        self,
+        pos_edge_rank,
+        neg_edge_rank,
+        node_id=None,
+        lambda_reg: float = 1e-4,
+    ):
+        loss_fn = BPRLoss(lambda_reg)
+        emb = self.embedding.weights[0]
+        emb = emb if node_id is None else ops.take(emb, node_id, axis=0)
+        return loss_fn(pos_edge_rank, neg_edge_rank, emb)
+
+    def __repr__(self) -> str:
+        return (f'{self.__class__.__name__}({self.num_nodes}, '
+                f'{self.embedding_dim}, num_layers={self.num_layers})')
+

k3_node/models/linkx.py

@@ -0,0 +1,181 @@
+import math
+from typing import Optional
+
+import keras
+from keras import ops
+
+from k3_node.layers.conv.message_passing import MessagePassing
+from k3_node.layers.conv.utils import scatter
+from k3_node.layers.norm import BatchNorm
+from k3_node.models.mlp import MLP
+
+
+class SparseLinear(keras.layers.Layer):
+    r"""A sparse linear transformation operator computing :math:`\mathbf{A}\mathbf{W} + \mathbf{b}`.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.models import SparseLinear
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        edge_index = np.random.randint(0, 10, size=(2, 30))  # 30 random edges
+
+        layer = SparseLinear(in_channels=10, out_channels=16)  # in_channels = number of nodes
+        out = layer(edge_index)  # multiplies the adjacency matrix with a learned weight matrix
+        print(tuple(out.shape))  # (10, 16)
+        ```
+    """
+    def __init__(self, in_channels: int, out_channels: int, bias: bool = True, **kwargs):
+        super().__init__(**kwargs)
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.use_bias = bias
+
+        self.weight = self.add_weight(
+            shape=(in_channels, out_channels),
+            initializer=keras.initializers.GlorotUniform(),
+            trainable=True,
+            name="weight",
+        )
+        if bias:
+            self.bias = self.add_weight(
+                shape=(out_channels,),
+                initializer="zeros",
+                trainable=True,
+                name="bias",
+            )
+        else:
+            self.bias = None
+
+    def build(self, input_shape=None):
+        self.built = True
+
+    def reset_parameters(self):
+        self.weight.assign(
+            keras.initializers.GlorotUniform()(shape=(self.in_channels, self.out_channels))
+        )
+        if self.use_bias and self.bias is not None:
+            self.bias.assign(ops.zeros((self.out_channels,)))
+
+    def call(self, edge_index, edge_weight=None):
+        row, col = edge_index[0], edge_index[1]
+        weight_j = ops.take(self.weight, row, axis=0)
+        if edge_weight is not None:
+            weight_j = ops.expand_dims(edge_weight, -1) * weight_j
+
+        out = scatter(weight_j, col, dim=0, dim_size=self.in_channels, reduce="sum")
+        if self.use_bias and self.bias is not None:
+            out = out + self.bias
+        return out
+
+
+class LINKX(keras.Model):
+    r"""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.
+
+    Args:
+        num_nodes (int): The number of nodes in the graph.
+        in_channels (int): Size of each input sample.
+        hidden_channels (int): Size of each hidden sample.
+        out_channels (int): Size of each output sample.
+        num_layers (int): Number of layers of :math:`\textrm{MLP}_{f}`.
+        num_edge_layers (int, optional): Number of layers of
+            :math:`\textrm{MLP}_{\mathbf{A}}`. (default: :obj:`1`)
+        num_node_layers (int, optional): Number of layers of
+            :math:`\textrm{MLP}_{\mathbf{X}}`. (default: :obj:`1`)
+        dropout (float, optional): Dropout probability. (default: :obj:`0.0`)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.models import LINKX
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        edge_index = np.random.randint(0, 10, size=(2, 30))  # 30 random edges
+
+        # LINKX learns from node features and adjacency separately, which suits heterophilic graphs
+        model = LINKX(num_nodes=10, in_channels=8, hidden_channels=32, out_channels=4, num_layers=2)
+        out = model(x, edge_index)
+        print(tuple(out.shape))  # (10, 4)
+        ```
+    """
+    def __init__(
+        self,
+        num_nodes: int,
+        in_channels: int,
+        hidden_channels: int,
+        out_channels: int,
+        num_layers: int,
+        num_edge_layers: int = 1,
+        num_node_layers: int = 1,
+        dropout: float = 0.0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.num_nodes = num_nodes
+        self.in_channels = in_channels
+        self.hidden_channels = hidden_channels
+        self.out_channels = out_channels
+        self.num_edge_layers = num_edge_layers
+
+        self.edge_lin = SparseLinear(num_nodes, hidden_channels)
+
+        if num_edge_layers > 1:
+            self.edge_norm = BatchNorm(hidden_channels)
+            channels = [hidden_channels] * num_edge_layers
+            self.edge_mlp = MLP(channels, dropout=0.0, act_first=True)
+        else:
+            self.edge_norm = None
+            self.edge_mlp = None
+
+        channels = [in_channels] + [hidden_channels] * num_node_layers
+        self.node_mlp = MLP(channels, dropout=0.0, act_first=True)
+
+        self.cat_lin1 = keras.layers.Dense(hidden_channels)
+        self.cat_lin2 = keras.layers.Dense(hidden_channels)
+
+        channels = [hidden_channels] * num_layers + [out_channels]
+        self.final_mlp = MLP(channels, dropout=dropout, act_first=True)
+
+    def build(self, input_shape=None):
+        self.edge_lin.build()
+        self.built = True
+
+    def reset_parameters(self):
+        r"""Resets all learnable parameters of the module."""
+        self.edge_lin.reset_parameters()
+        if self.edge_norm is not None and hasattr(self.edge_norm, "reset_parameters"):
+            self.edge_norm.reset_parameters()
+        if self.edge_mlp is not None and hasattr(self.edge_mlp, "reset_parameters"):
+            self.edge_mlp.reset_parameters()
+        self.node_mlp.reset_parameters()
+        self.final_mlp.reset_parameters()
+
+    def call(self, x, edge_index=None, edge_weight=None, training=None):
+        if edge_index is None and isinstance(x, (tuple, list)):
+            if len(x) >= 2:
+                x, edge_index = x[0], x[1]
+        out = self.edge_lin(edge_index, edge_weight)
+
+        if self.edge_norm is not None and self.edge_mlp is not None:
+            out = ops.relu(out)
+            out = self.edge_norm(out, training=training)
+            out = self.edge_mlp(out, training=training)
+
+        out = out + self.cat_lin1(out)
+
+        if x is not None:
+            x = self.node_mlp(x, training=training)
+            out = out + x
+            out = out + self.cat_lin2(x)
+
+        return self.final_mlp(ops.relu(out), training=training)
+
+    def __repr__(self) -> str:
+        return (f'{self.__class__.__name__}(num_nodes={self.num_nodes}, '
+                f'in_channels={self.in_channels}, '
+                f'out_channels={self.out_channels})')
+