From 353b7a49e14ffd4558020584cb76a20e6fd58b19 Mon Sep 17 00:00:00 2001 From: Severin Magel Date: Fri, 2 Oct 2026 18:34:32 -0400 Subject: [PATCH] Move transformer models into models/transformer DeepIce and GRIT are transformer architectures but lived under `models/gnn`. Move them beside `ISeeCube` in `models/transformer` (DeepIce's module renamed to `deepice.py` to match the class), and leave deprecation shims at the old `gnn.icemix` and `gnn.grit` paths that re-export the classes and warn on import. The classes themselves are unchanged, so pretrained configs -- which resolve `class_name` by walking `graphnet.models` -- still load. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_016pqUKMMW7ZPkhc9EChzCux --- examples/04_training/06_train_icemix_model.py | 2 +- examples/04_training/08_train_grit_model.py | 2 +- src/graphnet/models/gnn/__init__.py | 2 - src/graphnet/models/gnn/grit.py | 169 ++--------------- src/graphnet/models/gnn/icemix.py | 174 ++---------------- src/graphnet/models/transformer/__init__.py | 2 + src/graphnet/models/transformer/deepice.py | 167 +++++++++++++++++ src/graphnet/models/transformer/grit.py | 164 +++++++++++++++++ 8 files changed, 360 insertions(+), 322 deletions(-) create mode 100644 src/graphnet/models/transformer/deepice.py create mode 100644 src/graphnet/models/transformer/grit.py diff --git a/examples/04_training/06_train_icemix_model.py b/examples/04_training/06_train_icemix_model.py index 7395f7c8f..45ffd5a95 100644 --- a/examples/04_training/06_train_icemix_model.py +++ b/examples/04_training/06_train_icemix_model.py @@ -16,7 +16,7 @@ from graphnet.data.constants import FEATURES, TRUTH from graphnet.models import StandardModel from graphnet.models.detector.prometheus import Prometheus -from graphnet.models.gnn import DeepIce +from graphnet.models.transformer import DeepIce from graphnet.models.graphs import KNNGraph from graphnet.models.graphs.nodes import IceMixNodes from graphnet.models.task.reconstruction import ( diff --git a/examples/04_training/08_train_grit_model.py b/examples/04_training/08_train_grit_model.py index da18d99d8..56d58e09e 100644 --- a/examples/04_training/08_train_grit_model.py +++ b/examples/04_training/08_train_grit_model.py @@ -11,7 +11,7 @@ from graphnet.data.constants import FEATURES, TRUTH from graphnet.models import StandardModel from graphnet.models.detector.prometheus import Prometheus -from graphnet.models.gnn import GRIT +from graphnet.models.transformer import GRIT from graphnet.models.graphs import KNNGraphRRWP from graphnet.models.task.reconstruction import EnergyReconstruction from graphnet.training.callbacks import PiecewiseLinearLR diff --git a/src/graphnet/models/gnn/__init__.py b/src/graphnet/models/gnn/__init__.py index 9bbc1fb20..1b5ba7d0a 100644 --- a/src/graphnet/models/gnn/__init__.py +++ b/src/graphnet/models/gnn/__init__.py @@ -5,6 +5,4 @@ from .dynedge_jinst import DynEdgeJINST from .dynedge_kaggle_tito import DynEdgeTITO from .RNN_tito import RNN_TITO -from .icemix import DeepIce from .particlenet import ParticleNeT -from .grit import GRIT diff --git a/src/graphnet/models/gnn/grit.py b/src/graphnet/models/gnn/grit.py index baeda22b1..64c8aefbe 100644 --- a/src/graphnet/models/gnn/grit.py +++ b/src/graphnet/models/gnn/grit.py @@ -1,164 +1,19 @@ -"""Implementation of GRIT, a graph transformer model. +"""Deprecated module path for `GRIT`. -Original author: Liheng Ma -Original code: https://github.com/LiamMa/GRIT -Paper: "Graph Inductive Biases in Transformers without Message Passing", - https://arxiv.org/abs/2305.17589 - -Adapted by: Philip Weigel +`GRIT` is a graph-transformer model and now lives in +`graphnet.models.transformer`. This module re-exports it so existing imports +keep working, and warns on import that the path has moved. """ -import torch.nn as nn -from torch import Tensor -from torch_geometric.data import Data +from warnings import warn -from graphnet.models.gnn.gnn import GNN +from graphnet.models.transformer.grit import GRIT -from graphnet.models.components.grit_layers import ( - GritTransformerLayer, - SANGraphHead, -) -from graphnet.models.components.embedding import ( - RRWPLinearEdgeEncoder, - RRWPLinearNodeEncoder, - LinearNodeEncoder, - LinearEdgeEncoder, - RWSELinearNodeEncoder, +warn( + "`graphnet.models.gnn.grit` is deprecated; import `GRIT` from " + "`graphnet.models.transformer` instead.", + DeprecationWarning, + stacklevel=2, ) - -class GRIT(GNN): - """GRIT is a graph transformer model. - - Original code: - https://github.com/LiamMa/GRIT/blob/main/grit/network/grit_model.py - """ - - def __init__( - self, - nb_inputs: int, - hidden_dim: int, - nb_outputs: int = 1, - ksteps: int = 21, - n_layers: int = 10, - n_heads: int = 8, - pad_to_full_graph: bool = True, - add_node_attr_as_self_loop: bool = False, - dropout: float = 0.0, - fill_value: float = 0.0, - norm: nn.Module = nn.BatchNorm1d, - attn_dropout: float = 0.2, - edge_enhance: bool = True, - update_edges: bool = True, - attn_clamp: float = 5.0, - activation: nn.Module = nn.ReLU, - attn_activation: nn.Module = nn.ReLU, - norm_edges: bool = True, - enable_edge_transform: bool = True, - pred_head_layers: int = 2, - pred_head_activation: nn.Module = nn.ReLU, - pred_head_pooling: str = "mean", - position_encoding: str = "NoPE", - ): - """Construct `GRIT` model. - - Args: - nb_inputs: Number of inputs. - hidden_dim: Size of hidden dimension. - nb_outputs: Size of output dimension. - ksteps: Number of random walk steps. - n_layers: Number of GRIT layers. - n_heads: Number of heads in MHA. - pad_to_full_graph: Pad to form fully-connected graph. - add_node_attr_as_self_loop: Adds node attr as an self-edge. - dropout: Dropout probability. - fill_value: Padding value. - norm: Uninstantiated normalization layer. - Either `torch.nn.BatchNorm1d` or `torch.nn.LayerNorm`. - attn_dropout: Attention dropout probability. - edge_enhance: Applies learnable weight matrix with node-pair in - output node calculation for MHA. - update_edges: Update edge values after GRIT layer. - attn_clamp: Clamp absolute value of attention scores to a value. - activation: Uninstantiated activation function. - E.g. `torch.nn.ReLU` - attn_activation: Uninstantiated attention activation function. - E.g. `torch.nn.ReLU` - norm_edges: Apply normalization layer to edges. - enable_edge_transform: Apply transformation to edges. - pred_head_layers: Number of layers in the prediction head. - pred_head_activation: Uninstantiated prediction head activation - function. E.g. `torch.nn.ReLU` - pred_head_pooling: Pooling function to use for the prediction head, - either "mean" (default) or "add". - position_encoding: Method of position encoding. - """ - super().__init__(nb_inputs, nb_outputs) - self.position_encoding = position_encoding.lower() - if self.position_encoding == "nope": - encoders = [ - LinearNodeEncoder(nb_inputs, hidden_dim), - LinearEdgeEncoder(hidden_dim), - ] - elif self.position_encoding == "rrwp": - encoders = [ - LinearNodeEncoder(nb_inputs, hidden_dim), - LinearEdgeEncoder(hidden_dim), - RRWPLinearNodeEncoder(ksteps, hidden_dim), - RRWPLinearEdgeEncoder( - ksteps, - hidden_dim, - pad_to_full_graph=pad_to_full_graph, - add_node_attr_as_self_loop=add_node_attr_as_self_loop, - fill_value=fill_value, - ), - ] - elif self.position_encoding == "rwse": - encoders = [ - LinearNodeEncoder(nb_inputs, hidden_dim - (ksteps - 1)), - RWSELinearNodeEncoder(ksteps - 1, hidden_dim), - ] - self.encoders = nn.ModuleList(encoders) - - layers = [] - for _ in range(n_layers): - layers.append( - GritTransformerLayer( - in_dim=hidden_dim, - out_dim=hidden_dim, - num_heads=n_heads, - dropout=dropout, - activation=activation, - attn_dropout=attn_dropout, - norm=norm, - residual=True, - norm_edges=norm_edges, - enable_edge_transform=enable_edge_transform, - update_edges=update_edges, - attn_activation=attn_activation, - attn_clamp=attn_clamp, - attn_edge_enhance=edge_enhance, - ) - ) - self.layers = nn.ModuleList(layers) - self.head = SANGraphHead( - dim_in=hidden_dim, - dim_out=nb_outputs, - L=pred_head_layers, - activation=pred_head_activation, - pooling=pred_head_pooling, - ) - - def forward(self, x: Data) -> Tensor: - """Forward pass.""" - for encoder in self.encoders: - x = encoder(x) - - # Apply GRIT layers - for layer in self.layers: - x = layer(x) - - # Graph head - x = self.head(x) - - return x +__all__ = ["GRIT"] diff --git a/src/graphnet/models/gnn/icemix.py b/src/graphnet/models/gnn/icemix.py index cd3e3f6cc..1a51ec183 100644 --- a/src/graphnet/models/gnn/icemix.py +++ b/src/graphnet/models/gnn/icemix.py @@ -1,167 +1,19 @@ -"""Implementation of IceMix. +"""Deprecated module path for `DeepIce`. -This method was a solution submitted to the IceCube - Neutrinos in Deep Ice -Kaggle competition. - -Solution by DrHB: https://github.com/DrHB/icecube-2nd-place +`DeepIce` is a transformer model and now lives in +`graphnet.models.transformer`. This module re-exports it so existing imports +keep working, and warns on import that the path has moved. """ -import torch -import torch.nn as nn -from typing import Set, Dict, Any, Optional - -from graphnet.models.components.attention_blocks import ( - Block_rel, - Block, -) -from graphnet.models.components.embedding import ( - FourierEncoderEPJC, - SpacetimeEncoderEPJC, -) -from graphnet.models.gnn.dynedge import DynEdge -from graphnet.models.gnn.gnn import GNN -from graphnet.models.utils import array_to_sequence - -from torch_geometric.utils import to_dense_batch -from torch_geometric.data import Data -from torch import Tensor - - -class DeepIce(GNN): - """DeepIce model.""" - - def __init__( - self, - hidden_dim: int = 384, - mlp_ratio: int = 4, - seq_length: int = 192, - depth: int = 12, - head_size: int = 32, - depth_rel: int = 4, - n_rel: int = 1, - scaled_emb: bool = False, - include_dynedge: bool = False, - dynedge_args: Optional[Dict[str, Any]] = None, - n_features: int = 6, - ): - """Construct `DeepIce`. +from warnings import warn - Args: - hidden_dim: The latent feature dimension. - mlp_ratio: Mlp expansion ratio of FourierEncoderEPJC and - Transformer. - seq_length: The base feature dimension. - depth: The depth of the transformer. - head_size: The size of the attention heads. - depth_rel: The depth of the relative transformer. - n_rel: The number of relative transformer layers to use. - scaled_emb: Whether to scale the sinusoidal positional embeddings. - include_dynedge: If True, pulse-level predictions from `DynEdge` - will be added as features to the model. - dynedge_args: Initialization arguments for DynEdge. If not - provided, DynEdge will be initialized with the original Kaggle - Competition settings. If `include_dynedge` is False, this - argument have no impact. - n_features: The number of features in the input data. - """ - super().__init__(seq_length, hidden_dim) - fourier_out_dim = hidden_dim // 2 if include_dynedge else hidden_dim - self.fourier_ext = FourierEncoderEPJC( - seq_length=seq_length, - mlp_dim=None, - output_dim=fourier_out_dim, - scaled=scaled_emb, - n_features=n_features, - ) - self.rel_pos = SpacetimeEncoderEPJC(head_size) - self.sandwich = nn.ModuleList( - [ - Block_rel( - input_dim=hidden_dim, num_heads=hidden_dim // head_size - ) - for _ in range(depth_rel) - ] - ) - self.cls_token = nn.Linear(hidden_dim, 1, bias=False) - self.blocks = nn.ModuleList( - [ - Block( - input_dim=hidden_dim, - num_heads=hidden_dim // head_size, - mlp_ratio=mlp_ratio, - drop_path=0.0 * (i / (depth - 1)), - init_values=1, - ) - for i in range(depth) - ] - ) - self.n_rel = n_rel +from graphnet.models.transformer.deepice import DeepIce - if include_dynedge and dynedge_args is None: - self.warning_once("Running with default DynEdge settings") - self.dyn_edge = DynEdge( - nb_inputs=9, - nb_neighbours=9, - post_processing_layer_sizes=[336, hidden_dim // 2], - dynedge_layer_sizes=[ - (128, 256), - (336, 256), - (336, 256), - (336, 256), - ], - global_pooling_schemes=None, - activation_layer="gelu", - add_norm_layer=True, - skip_readout=True, - ) - elif include_dynedge and not (dynedge_args is None): - self.dyn_edge = DynEdge(**dynedge_args) - - self.include_dynedge = include_dynedge - - @torch.jit.ignore - def no_weight_decay(self) -> Set: - """cls_tocken should not be subject to weight decay during training.""" - return {"cls_token"} - - def forward(self, data: Data) -> Tensor: - """Apply learnable forward pass.""" - x0, mask, seq_length = array_to_sequence( - data.x, data.batch, padding_value=0 - ) - x = self.fourier_ext(x0, seq_length) - rel_pos_bias = self.rel_pos(x0) - batch_size = mask.shape[0] - if self.include_dynedge: - graph = self.dyn_edge(data) - graph, _ = to_dense_batch(graph, data.batch) - x = torch.cat([x, graph], 2) - - attn_mask = torch.zeros(mask.shape, device=mask.device) - attn_mask[~mask] = -torch.inf - - for i, blk in enumerate(self.sandwich): - x = blk(x, attn_mask, rel_pos_bias) - if i + 1 == self.n_rel: - rel_pos_bias = None - - mask = torch.cat( - [ - torch.ones( - batch_size, 1, dtype=mask.dtype, device=mask.device - ), - mask, - ], - 1, - ) - attn_mask = torch.zeros(mask.shape, device=mask.device) - attn_mask[~mask] = -torch.inf - cls_token = self.cls_token.weight.unsqueeze(0).expand( - batch_size, -1, -1 - ) - x = torch.cat([cls_token, x], 1) - - for blk in self.blocks: - x = blk(x, None, attn_mask) +warn( + "`graphnet.models.gnn.icemix` is deprecated; import `DeepIce` from " + "`graphnet.models.transformer` instead.", + DeprecationWarning, + stacklevel=2, +) - return x[:, 0] +__all__ = ["DeepIce"] diff --git a/src/graphnet/models/transformer/__init__.py b/src/graphnet/models/transformer/__init__.py index a0aac6f2b..ebdd3071f 100644 --- a/src/graphnet/models/transformer/__init__.py +++ b/src/graphnet/models/transformer/__init__.py @@ -1,3 +1,5 @@ """Transformer-specific modules.""" +from .deepice import DeepIce +from .grit import GRIT from .iseecube import ISeeCube diff --git a/src/graphnet/models/transformer/deepice.py b/src/graphnet/models/transformer/deepice.py new file mode 100644 index 000000000..cd3e3f6cc --- /dev/null +++ b/src/graphnet/models/transformer/deepice.py @@ -0,0 +1,167 @@ +"""Implementation of IceMix. + +This method was a solution submitted to the IceCube - Neutrinos in Deep Ice +Kaggle competition. + +Solution by DrHB: https://github.com/DrHB/icecube-2nd-place +""" + +import torch +import torch.nn as nn +from typing import Set, Dict, Any, Optional + +from graphnet.models.components.attention_blocks import ( + Block_rel, + Block, +) +from graphnet.models.components.embedding import ( + FourierEncoderEPJC, + SpacetimeEncoderEPJC, +) +from graphnet.models.gnn.dynedge import DynEdge +from graphnet.models.gnn.gnn import GNN +from graphnet.models.utils import array_to_sequence + +from torch_geometric.utils import to_dense_batch +from torch_geometric.data import Data +from torch import Tensor + + +class DeepIce(GNN): + """DeepIce model.""" + + def __init__( + self, + hidden_dim: int = 384, + mlp_ratio: int = 4, + seq_length: int = 192, + depth: int = 12, + head_size: int = 32, + depth_rel: int = 4, + n_rel: int = 1, + scaled_emb: bool = False, + include_dynedge: bool = False, + dynedge_args: Optional[Dict[str, Any]] = None, + n_features: int = 6, + ): + """Construct `DeepIce`. + + Args: + hidden_dim: The latent feature dimension. + mlp_ratio: Mlp expansion ratio of FourierEncoderEPJC and + Transformer. + seq_length: The base feature dimension. + depth: The depth of the transformer. + head_size: The size of the attention heads. + depth_rel: The depth of the relative transformer. + n_rel: The number of relative transformer layers to use. + scaled_emb: Whether to scale the sinusoidal positional embeddings. + include_dynedge: If True, pulse-level predictions from `DynEdge` + will be added as features to the model. + dynedge_args: Initialization arguments for DynEdge. If not + provided, DynEdge will be initialized with the original Kaggle + Competition settings. If `include_dynedge` is False, this + argument have no impact. + n_features: The number of features in the input data. + """ + super().__init__(seq_length, hidden_dim) + fourier_out_dim = hidden_dim // 2 if include_dynedge else hidden_dim + self.fourier_ext = FourierEncoderEPJC( + seq_length=seq_length, + mlp_dim=None, + output_dim=fourier_out_dim, + scaled=scaled_emb, + n_features=n_features, + ) + self.rel_pos = SpacetimeEncoderEPJC(head_size) + self.sandwich = nn.ModuleList( + [ + Block_rel( + input_dim=hidden_dim, num_heads=hidden_dim // head_size + ) + for _ in range(depth_rel) + ] + ) + self.cls_token = nn.Linear(hidden_dim, 1, bias=False) + self.blocks = nn.ModuleList( + [ + Block( + input_dim=hidden_dim, + num_heads=hidden_dim // head_size, + mlp_ratio=mlp_ratio, + drop_path=0.0 * (i / (depth - 1)), + init_values=1, + ) + for i in range(depth) + ] + ) + self.n_rel = n_rel + + if include_dynedge and dynedge_args is None: + self.warning_once("Running with default DynEdge settings") + self.dyn_edge = DynEdge( + nb_inputs=9, + nb_neighbours=9, + post_processing_layer_sizes=[336, hidden_dim // 2], + dynedge_layer_sizes=[ + (128, 256), + (336, 256), + (336, 256), + (336, 256), + ], + global_pooling_schemes=None, + activation_layer="gelu", + add_norm_layer=True, + skip_readout=True, + ) + elif include_dynedge and not (dynedge_args is None): + self.dyn_edge = DynEdge(**dynedge_args) + + self.include_dynedge = include_dynedge + + @torch.jit.ignore + def no_weight_decay(self) -> Set: + """cls_tocken should not be subject to weight decay during training.""" + return {"cls_token"} + + def forward(self, data: Data) -> Tensor: + """Apply learnable forward pass.""" + x0, mask, seq_length = array_to_sequence( + data.x, data.batch, padding_value=0 + ) + x = self.fourier_ext(x0, seq_length) + rel_pos_bias = self.rel_pos(x0) + batch_size = mask.shape[0] + if self.include_dynedge: + graph = self.dyn_edge(data) + graph, _ = to_dense_batch(graph, data.batch) + x = torch.cat([x, graph], 2) + + attn_mask = torch.zeros(mask.shape, device=mask.device) + attn_mask[~mask] = -torch.inf + + for i, blk in enumerate(self.sandwich): + x = blk(x, attn_mask, rel_pos_bias) + if i + 1 == self.n_rel: + rel_pos_bias = None + + mask = torch.cat( + [ + torch.ones( + batch_size, 1, dtype=mask.dtype, device=mask.device + ), + mask, + ], + 1, + ) + attn_mask = torch.zeros(mask.shape, device=mask.device) + attn_mask[~mask] = -torch.inf + cls_token = self.cls_token.weight.unsqueeze(0).expand( + batch_size, -1, -1 + ) + x = torch.cat([cls_token, x], 1) + + for blk in self.blocks: + x = blk(x, None, attn_mask) + + return x[:, 0] diff --git a/src/graphnet/models/transformer/grit.py b/src/graphnet/models/transformer/grit.py new file mode 100644 index 000000000..baeda22b1 --- /dev/null +++ b/src/graphnet/models/transformer/grit.py @@ -0,0 +1,164 @@ +"""Implementation of GRIT, a graph transformer model. + +Original author: Liheng Ma +Original code: https://github.com/LiamMa/GRIT +Paper: "Graph Inductive Biases in Transformers without Message Passing", + https://arxiv.org/abs/2305.17589 + +Adapted by: Philip Weigel +""" + +import torch.nn as nn +from torch import Tensor +from torch_geometric.data import Data + +from graphnet.models.gnn.gnn import GNN + +from graphnet.models.components.grit_layers import ( + GritTransformerLayer, + SANGraphHead, +) +from graphnet.models.components.embedding import ( + RRWPLinearEdgeEncoder, + RRWPLinearNodeEncoder, + LinearNodeEncoder, + LinearEdgeEncoder, + RWSELinearNodeEncoder, +) + + +class GRIT(GNN): + """GRIT is a graph transformer model. + + Original code: + https://github.com/LiamMa/GRIT/blob/main/grit/network/grit_model.py + """ + + def __init__( + self, + nb_inputs: int, + hidden_dim: int, + nb_outputs: int = 1, + ksteps: int = 21, + n_layers: int = 10, + n_heads: int = 8, + pad_to_full_graph: bool = True, + add_node_attr_as_self_loop: bool = False, + dropout: float = 0.0, + fill_value: float = 0.0, + norm: nn.Module = nn.BatchNorm1d, + attn_dropout: float = 0.2, + edge_enhance: bool = True, + update_edges: bool = True, + attn_clamp: float = 5.0, + activation: nn.Module = nn.ReLU, + attn_activation: nn.Module = nn.ReLU, + norm_edges: bool = True, + enable_edge_transform: bool = True, + pred_head_layers: int = 2, + pred_head_activation: nn.Module = nn.ReLU, + pred_head_pooling: str = "mean", + position_encoding: str = "NoPE", + ): + """Construct `GRIT` model. + + Args: + nb_inputs: Number of inputs. + hidden_dim: Size of hidden dimension. + nb_outputs: Size of output dimension. + ksteps: Number of random walk steps. + n_layers: Number of GRIT layers. + n_heads: Number of heads in MHA. + pad_to_full_graph: Pad to form fully-connected graph. + add_node_attr_as_self_loop: Adds node attr as an self-edge. + dropout: Dropout probability. + fill_value: Padding value. + norm: Uninstantiated normalization layer. + Either `torch.nn.BatchNorm1d` or `torch.nn.LayerNorm`. + attn_dropout: Attention dropout probability. + edge_enhance: Applies learnable weight matrix with node-pair in + output node calculation for MHA. + update_edges: Update edge values after GRIT layer. + attn_clamp: Clamp absolute value of attention scores to a value. + activation: Uninstantiated activation function. + E.g. `torch.nn.ReLU` + attn_activation: Uninstantiated attention activation function. + E.g. `torch.nn.ReLU` + norm_edges: Apply normalization layer to edges. + enable_edge_transform: Apply transformation to edges. + pred_head_layers: Number of layers in the prediction head. + pred_head_activation: Uninstantiated prediction head activation + function. E.g. `torch.nn.ReLU` + pred_head_pooling: Pooling function to use for the prediction head, + either "mean" (default) or "add". + position_encoding: Method of position encoding. + """ + super().__init__(nb_inputs, nb_outputs) + self.position_encoding = position_encoding.lower() + if self.position_encoding == "nope": + encoders = [ + LinearNodeEncoder(nb_inputs, hidden_dim), + LinearEdgeEncoder(hidden_dim), + ] + elif self.position_encoding == "rrwp": + encoders = [ + LinearNodeEncoder(nb_inputs, hidden_dim), + LinearEdgeEncoder(hidden_dim), + RRWPLinearNodeEncoder(ksteps, hidden_dim), + RRWPLinearEdgeEncoder( + ksteps, + hidden_dim, + pad_to_full_graph=pad_to_full_graph, + add_node_attr_as_self_loop=add_node_attr_as_self_loop, + fill_value=fill_value, + ), + ] + elif self.position_encoding == "rwse": + encoders = [ + LinearNodeEncoder(nb_inputs, hidden_dim - (ksteps - 1)), + RWSELinearNodeEncoder(ksteps - 1, hidden_dim), + ] + self.encoders = nn.ModuleList(encoders) + + layers = [] + for _ in range(n_layers): + layers.append( + GritTransformerLayer( + in_dim=hidden_dim, + out_dim=hidden_dim, + num_heads=n_heads, + dropout=dropout, + activation=activation, + attn_dropout=attn_dropout, + norm=norm, + residual=True, + norm_edges=norm_edges, + enable_edge_transform=enable_edge_transform, + update_edges=update_edges, + attn_activation=attn_activation, + attn_clamp=attn_clamp, + attn_edge_enhance=edge_enhance, + ) + ) + self.layers = nn.ModuleList(layers) + self.head = SANGraphHead( + dim_in=hidden_dim, + dim_out=nb_outputs, + L=pred_head_layers, + activation=pred_head_activation, + pooling=pred_head_pooling, + ) + + def forward(self, x: Data) -> Tensor: + """Forward pass.""" + for encoder in self.encoders: + x = encoder(x) + + # Apply GRIT layers + for layer in self.layers: + x = layer(x) + + # Graph head + x = self.head(x) + + return x