Repository: datamol-io/graphium (language: python)

Issue:
Meta: goli inception
The below mostly focus on what is outside of the business logic of goli, meaning it's about coding good practices and guidelines to make the future integration to the platform easy and seamless. It also prepares the lib for a possible release to the public domain (open source).

---

What I recently did on the default branch:

- change the default branch from `main` to `master` (this is because one tool we use to manage the release process called `rever` does not support `main` yet).
- add CI for testing and linting the code (with black)
- add a dummy test (so we have a test folder)
- start a doc with `mkdocs` (similar but simpler to set up and fewer files than sphinx): private access to the doc at https://invivoai--goli.github.privpage.net/
- tweak the env.yml file
- use `datamol` dep: datamol is the numpy of molecule built on top of rdkit, it provides convenient functions such as `dm.to_mol`. We extensively use it on the platform in order to standardize the way we work with molecules.
- tweak readme to reflect changes
- move part of the readme doc to CONTRIBUTE.md
- add `__init__.py` files in various Python modules.

---

Infra tasks (todo by dom or/and hadrien):

- create a private S3 bucket to manage goli's dataset and experiment results
- 
---

What should be done + some guidelines (it also includes Dom's suggestions from the original README file):

- replace functions such as `to_mol` by datamol.
  - consider contributing to `datamol` when relevant.
- create a notebooks folder to highlight an interesting part of the code (take izanagi as an example: https://github.com/invivoai-platform/izanagi/tree/master/tutorials)
- cleanup existing GNN models and keep only the relevant ones
- convert `MolecularTransformer` by something function-based, simpler to use and more flexible.
- update `ModelWrapper` to the latest PL version and use their metrics API (https://pytorch-lightning.readthedocs.io/en/latest/metrics.html)
  - try to keep `ModelWrapper` as light as possible, relying most of the time on PL own API.
  - Don't hesitate to use their callback system: https://pytorch-lightning.readthedocs.io/en/latest/callbacks.html
  - `ModelWrapper` **must** serialize out-of-the-box: https://pytorch-lightning.readthedocs.io/en/latest/weights_loading.html
  - build code that is device agnostic (avoid explicit moves to cuda device) and let PL do the job for you.
- consider optuna as an hparams search engine: on this, I know this is very opinionated. I find optuna to be simpler and lighter than hydra. If you like hydra or already are comfortable with it, feel free to ignore optuna.
- for any file IO in goli, use [`fsspec`](https://filesystem-spec.readthedocs.io/en/latest/) so you can transparently read/write on a local, S3 or GS filesystem without the need for a different code for each.
- use a private S3 bucket to deal with the dataset and results of the experiments that are going to be performed.
  - make a python command using fsspec to download the dataset instead of using a GDrive (or use the S3 cli tool).
  - clearly separate the different files/folders in this repo: `dataset`, `pretrained-model`, `experiments`, etc
- use [omegaconf](https://github.com/omry/omegaconf) instead of hydra to read/load config files.
- Replace all the `print` statements by `logger.info` or equivalent using the `loguru` library and `from loguru import logger`.
- use [click](https://github.com/pallets/click) for all CLI commands.
- clearly separate experiments from re-usable source code.

---

Feel free to discuss high-level things in this ticket. To discuss/implement a specific point from above, please open a new ticket.

ping @DomInvivo 

Ipu fp16
Checklist:

- [x] Added a `news` entry: _copy `news/TEMPLATE.rst` to `news/my-feature-or-branch.rst`) and edit it._
- [x] data cast to FP16 before forward pass when executing in IPU
- [x] program outputs cast back to FP32 on the host
- [x] model weights in FP16 (this is done in PT Lightning, behind the scenes)
- [x] half partials in FP16
- [x] part of optimizer state in FP16, second order in FP32

---

Allow overriding training options for validation
Resolves #188 

Checklist:

- [ ] Added a `news` entry: _copy `news/TEMPLATE.rst` to `news/my-feature-or-branch.rst`) and edit it._

---



Discussion / hints:

Fix patch:
diff --git a/docs/api/goli.nn.md b/docs/api/goli.nn.md
index 9ce75de4b..ba78d3ba2 100644
--- a/docs/api/goli.nn.md
+++ b/docs/api/goli.nn.md
@@ -17,6 +17,7 @@
 ::: goli.nn.pyg_layers.gin_pyg
 ::: goli.nn.pyg_layers.gps_pyg
 ::: goli.nn.pyg_layers.pna_pyg
+::: goli.nn.pyg_layers.mpnn_pyg
 ::: goli.nn.pyg_layers.pooling_pyg
 ::: goli.nn.base_graph_layer
 ::: goli.nn.base_layers
diff --git a/expts/configs/config_pcqmv2_gpspp.yaml b/expts/configs/config_pcqmv2_gpspp.yaml
index 6112fbcad..3c7b21723 100644
--- a/expts/configs/config_pcqmv2_gpspp.yaml
+++ b/expts/configs/config_pcqmv2_gpspp.yaml
@@ -4,7 +4,7 @@ constants:
   seed: &seed 42
   raise_train_error: true   # Whether the code should raise an error if it crashes during training
   accelerator:
-    type: cpu  #cpu or ipu or gpu
+    type: ipu  #cpu or ipu or gpu
 
 datamodule:
   module_type: "MultitaskFromSmilesDataModule"
@@ -57,12 +57,12 @@ datamodule:
     num_workers: -1
 
     ipu_dataloader_training_opts:
-      max_num_nodes_per_graph: &max_num_nodes 50
-      max_num_edges_per_graph: 150
+      max_num_nodes_per_graph: &max_num_nodes 30
+      max_num_edges_per_graph: &max_num_edges 60
 
     ipu_dataloader_inference_opts:
-      max_num_nodes_per_graph: 50
-      max_num_edges_per_graph: 150
+      max_num_nodes_per_graph: *max_num_nodes
+      max_num_edges_per_graph: *max_num_edges
 
 architecture:
   model_type: FullGraphMultiTaskNetwork
@@ -116,11 +116,10 @@ architecture:
       gaussian_pos:
         encoder_type: "gaussian_kernel"
         input_keys: ["positions_3d"]
-        output_keys: ["feat", "graph_gaussian_bias_3d"] 
+        output_keys: ["feat", "graph_gaussian_bias_3d"]
         num_heads: 1
         num_layers: 1 #2
         embed_dim: 64
-        max_num_nodes_per_graph: *max_num_nodes
         use_input_keys_prefix: False
 
 
@@ -136,9 +135,11 @@ architecture:
     last_normalization: *normalization
     residual_type: simple
     pooling: [sum]
-    virtual_node: 'none'
+    virtual_node: "sum"
+    use_virtual_edges: True
     layer_type: 'pyg:gps' #pyg:gine #'pyg:gps' # pyg:gated-gcn, pyg:gine,pyg:gps
     layer_kwargs:  # Parameters for the model itself. You could define dropout_attn: 0.1
+      out_dim_edges: 32
       node_residual: false
       mpnn_type: 'pyg:mpnnplus'
       mpnn_kwargs:
@@ -146,13 +147,12 @@ architecture:
         out_dim: 64
         in_dim_edges: 32
         out_dim_edges: 32
+        droppath_rate: 0.
       attn_type: "full-attention" # "full-attention", "none"
-      #biased_attention_key: "graph_gaussian_bias_3d" #same key from gaussian_pos/output_keys/"graph_xxx"
+      biased_attention_key: "graph_gaussian_bias_3d" #same key from gaussian_pos/output_keys/"graph_xxx"
       attn_kwargs: null
-      # out_dim_edges: 16
-      # aggregators: [mean, max]
-      # scalers: [identity, amplification, attenuation]
-      # num_heads: 3
+      droppath_rate_attn: 0.3
+      droppath_rate_ffn: 0.3
 
   post_nn:
     out_dim: 32
diff --git a/goli/config/_loader.py b/goli/config/_loader.py
index 450f7619d..fed9da07b 100644
--- a/goli/config/_loader.py
+++ b/goli/config/_loader.py
@@ -7,6 +7,8 @@
 from loguru import logger
 import yaml
 import joblib
+import pathlib
+import warnings
 
 # Torch
 import torch
@@ -90,15 +92,44 @@ def load_datamodule(config: Union[omegaconf.DictConfig, Dict[str, Any]]) -> Base
 
     cfg_data = config["datamodule"]["args"]
 
-    # Default empty values for the IPU configurations
-    ipu_training_opts, ipu_inference_opts = None, None
-    ipu_file = "expts/configs/ipu.config"
-    ipu_dataloader_training_opts = cfg_data.pop("ipu_dataloader_training_opts", {})
-    ipu_dataloader_inference_opts = cfg_data.pop("ipu_dataloader_inference_opts", {})
+    # Instanciate the datamodule
+    module_class = DATAMODULE_DICT[config["datamodule"]["module_type"]]
 
-    if get_accelerator(config) == "ipu":
+    if get_accelerator(config) != "ipu":
+        datamodule = module_class(
+            **config["datamodule"]["args"],
+        )
+        return datamodule
+
+    # IPU specific adjustments
+    else:
+        # Default empty values for the IPU configurations
+        ipu_training_opts, ipu_inference_opts = None, None
+        ipu_training_config_path = "expts/configs/ipu.config"
+        if pathlib.Path(ipu_training_config_path).is_file():
+            ipu_training_config_file = ipu_training_config_path
+        else:
+            raise ValueError(
+                f"IPU configuration path must be specified "
+                "and must be a file, instead got "
+                '"{ipu_training_config_path}"'
+            )
+
+        ipu_inference_config_overrides_path = "expts/configs/ipu_inference.config"
+        if pathlib.Path(ipu_inference_config_overrides_path).is_file():
+            ipu_inference_config_overrides_file = ipu_inference_config_overrides_path
+        else:
+            warnings.warn(
+                "IPU inference overrides configuration either not specified "
+                "or not a file, using same options for training and inference"
+            )
+            ipu_inference_config_overrides_file = None
+
+        ipu_dataloader_training_opts = cfg_data.pop("ipu_dataloader_training_opts", {})
+        ipu_dataloader_inference_opts = cfg_data.pop("ipu_dataloader_inference_opts", {})
         ipu_training_opts, ipu_inference_opts = load_ipu_options(
-            ipu_file=ipu_file,
+            ipu_file=ipu_training_config_file,
+            ipu_inference_overrides=ipu_inference_config_overrides_file,
             seed=config["constants"]["seed"],
             model_name=config["constants"]["name"],
             gradient_accumulation=config["trainer"]["trainer"].get("accumulate_grad_batches", None),
@@ -118,17 +149,15 @@ def load_datamodule(config: Union[omegaconf.DictConfig, Dict[str, Any]]) -> Base
         )
         ipu_dataloader_inference_opts.set_kwargs()
 
-    # Instanciate the datamodule
-    module_class = DATAMODULE_DICT[config["datamodule"]["module_type"]]
-    datamodule = module_class(
-        ipu_training_opts=ipu_training_opts,
-        ipu_inference_opts=ipu_inference_opts,
-        ipu_dataloader_training_opts=ipu_dataloader_training_opts,
-        ipu_dataloader_inference_opts=ipu_dataloader_inference_opts,
-        **config["datamodule"]["args"],
-    )
+        datamodule = module_class(
+            ipu_training_opts=ipu_training_opts,
+            ipu_inference_opts=ipu_inference_opts,
+            ipu_dataloader_training_opts=ipu_dataloader_training_opts,
+            ipu_dataloader_inference_opts=ipu_dataloader_inference_opts,
+            **config["datamodule"]["args"],
+        )
 
-    return datamodule
+        return datamodule
 
 
 def load_metrics(config: Union[omegaconf.DictConfig, Dict[str, Any]]) -> Dict[str, MetricWrapper]:
@@ -313,6 +342,7 @@ def load_trainer(config: Union[omegaconf.DictConfig, Dict[str, Any]], run_name:
             seed=config["constants"]["seed"],
             model_name=config["constants"]["name"],
             gradient_accumulation=config["trainer"]["trainer"].get("accumulate_grad_batches", None),
+            precision=cfg_trainer["trainer"].get("precision", None),
         )
         from goli.ipu.ipu_wrapper import DictIPUStrategy
 
diff --git a/goli/features/featurizer.py b/goli/features/featurizer.py
index 96ed7410f..c0190db70 100644
--- a/goli/features/featurizer.py
+++ b/goli/features/featurizer.py
@@ -821,7 +821,7 @@ def make_pyg_graph(self, **kwargs) -> Data:
 
         # Convert adjacency matrix into edge_index and edge_weight
         edge_index = torch.Tensor(np.stack((self.adj.row, self.adj.col))).to(dtype=torch.int16)
-        edge_weight = torch.Tensor(self.adj.data)
+        edge_weight = torch.from_numpy(self.adj.data)
         num_nodes = self.adj.shape[0]
 
         # Get the node and edge data
diff --git a/goli/ipu/ipu_utils.py b/goli/ipu/ipu_utils.py
index ed3245898..8adc34173 100644
--- a/goli/ipu/ipu_utils.py
+++ b/goli/ipu/ipu_utils.py
@@ -1,6 +1,7 @@
 from copy import deepcopy
 from types import ModuleType
 from typing import Optional, Tuple
+import torch
 
 
 def import_poptorch(raise_error=True) -> Optional[ModuleType]:
@@ -29,11 +30,23 @@ def import_poptorch(raise_error=True) -> Optional[ModuleType]:
         return
 
 
+def is_running_on_ipu() -> bool:
+    """
+    Returns whether the current module is running on ipu.
+    Needs to be used in the `forward` or `backward` pass.
+    """
+    poptorch = import_poptorch(raise_error=False)
+    on_ipu = (poptorch is not None) and (poptorch.isRunningOnIpu())
+    return on_ipu
+
+
 def load_ipu_options(
     ipu_file: str,
     seed: Optional[int] = None,
     model_name: Optional[str] = None,
     gradient_accumulation: Optional[int] = None,
+    precision: Optional[int] = None,
+    ipu_inference_overrides: Optional[str] = None,
 ) -> Tuple["poptorch.Options", "poptorch.Options"]:
     """
     Load the IPU options from the config file.
@@ -69,6 +82,8 @@ def load_ipu_options(
 
         seed: random seed for the IPU
         model_name: Name of the model, to be used for ipu profiling
+        ipu_inference_overrides: optional file path containing IPU configuration overrides for inference.
+            If this file is provided, options in this file override those in `ipu_file` for inference.
 
     Returns:
 
@@ -95,7 +110,9 @@ def load_ipu_options(
         ipu_options.Training.gradientAccumulation(gradient_accumulation)
 
     ipu_options.anchorTensor("input", "input")
-
+    if precision == 16:
+        # IPUOptions.loadFromFile currently doesn't support setting half partials, doing it here
+        ipu_options.Precision.setPartialsType(torch.half)
     training_opts = ipu_options
 
     # Change the inference options to remove gradient accumulation
@@ -103,5 +120,7 @@ def load_ipu_options(
     inference_opts.Training.gradientAccumulation(1)
     if model_name is not None:
         inference_opts.modelName(f"{model_name}_inference")
+    if ipu_inference_overrides is not None:
+        inference_opts.loadFromFile(ipu_inference_overrides)
 
     return training_opts, inference_opts
diff --git a/goli/ipu/ipu_wrapper.py b/goli/ipu/ipu_wrapper.py
index d5689d738..902e9c61b 100644
--- a/goli/ipu/ipu_wrapper.py
+++ b/goli/ipu/ipu_wrapper.py
@@ -13,6 +13,8 @@
 from torch_geometric.data import Data, Batch
 from torch_geometric.data.data import BaseData
 from loguru import logger
+import functools
+import collections
 
 poptorch = import_poptorch()
 
@@ -125,7 +127,8 @@ def compute_loss(
         )
 
     def on_train_batch_end(self, outputs, batch, batch_idx):
-        outputs["loss/train"] = outputs["loss"].mean()
+        outputs = self.convert_from_fp16(outputs)
+        outputs["loss"] = outputs["loss"].mean()
         super().on_train_batch_end(outputs, batch, batch_idx)
 
     def training_step(self, features, labels) -> Dict[str, Any]:
@@ -159,9 +162,27 @@ def predict_step(self, **inputs) -> Dict[str, Any]:
 
         return step_dict
 
+    def validation_epoch_end(self, outputs: Dict[str, Any]):
+        outputs = self.convert_from_fp16(outputs)
+        super().validation_epoch_end(outputs)
+
+    def evaluation_epoch_end(self, outputs: Dict[str, Any]):
+        outputs = self.convert_from_fp16(outputs)
+        super().evaluation_epoch_end(outputs)
+
+    def test_epoch_end(self, outputs: Dict[str, Any]):
+        outputs = self.convert_from_fp16(outputs)
+        super().test_epoch_end(outputs)
+
     def configure_optimizers(self, impl=None):
         if impl is None:
-            impl = self.poptorch.optim.Adam
+            dtype = self.precision_to_dtype(self.trainer.precision)
+            impl = functools.partial(
+                self.poptorch.optim.Adam,
+                accum_type=dtype,
+                first_order_momentum_accum_type=dtype,
+                second_order_momentum_accum_type=torch.float,
+            )
         return super().configure_optimizers(impl=impl)
 
     def squeeze_input_dims(self, features, labels):
@@ -174,6 +195,41 @@ def squeeze_input_dims(self, features, labels):
 
         return features, labels
 
+    def convert_from_fp16(self, data: Any) -> Any:
+        """
+        Converts tensors from FP16 to FP32. Useful to convert the IPU program output data
+        """
+        if isinstance(data, collections.Sequence):
+            for idx in range(len(data)):
+                data[idx] = self.convert_from_fp16(data[idx])
+        elif isinstance(data, collections.Mapping):
+            for key in data:
+                data[key] = self.convert_from_fp16(data[key])
+        elif isinstance(data, torch.Tensor) and data.dtype == torch.float16:
+            data = data.float()
+        return data
+
+    def _convert_features_dtype(self, feats):
+        """
+        Converts features to trainer precision rather than model precision.
+        Necessary to run IPU on FP16.
+        """
+        dtype = self.precision_to_dtype(self.trainer.precision)
+
+        # Convert features to dtype
+        if isinstance(feats, torch.Tensor):
+            feats = feats.to(dtype)
+        elif isinstance(feats, (Data, Batch, dict)):
+            for key, val in feats.items():
+                if isinstance(val, torch.Tensor) and (val.is_floating_point()):
+                    feats[key] = val.to(dtype=dtype)
+        else:
+            raise ValueError(f"Unsupported feats type `{type(feats)}` : {feats}")
+        return feats
+
+    def precision_to_dtype(self, precision):
+        return torch.half if precision in (16, "16") else torch.float
+
     def get_num_graphs(self, data: Batch):
         """
         IPU specific method to compute the number of graphs in a Batch,
diff --git a/goli/nn/architectures/encoder_manager.py b/goli/nn/architectures/encoder_manager.py
index a14cf6528..a3cd9cb80 100644
--- a/goli/nn/architectures/encoder_manager.py
+++ b/goli/nn/architectures/encoder_manager.py
@@ -27,6 +27,7 @@ class EncoderManager(nn.Module):
     def __init__(
         self,
         pe_encoders_kwargs: Optional[Dict[str, Any]] = None,
+        max_num_nodes_per_graph: Optional[int] = None,
         name: str = "encoder_manager",
     ):
         r"""
@@ -45,6 +46,7 @@ def __init__(
 
         super().__init__()
         self.name = name
+        self.max_num_nodes_per_graph = max_num_nodes_per_graph
         self.pe_encoders_kwargs = deepcopy(pe_encoders_kwargs)
         self.pe_encoders = self._initialize_positional_encoders(pe_encoders_kwargs)
 
@@ -104,6 +106,10 @@ def _initialize_positional_encoders(self, pe_encoders_kwargs: Dict[str, Any]) ->
                         f"`in_dim` not understood for encoder {encoder_name}. Provided: {this_in_dims}. Accepted keys are: {accepted_keys}"
                     )
 
+                # Add the max_num_nodes_per_graph if it's in the accepted input keys
+                if "max_num_nodes_per_graph" in accepted_keys:
+                    encoder_kwargs["max_num_nodes_per_graph"] = self.max_num_nodes_per_graph
+
                 # Initialize the pe_encoder layer
                 pe_out_dim2 = encoder_kwargs.pop("out_dim", None)
                 if pe_out_dim2 is not None:
@@ -257,10 +263,3 @@ def out_dim(self) -> int:
             return self.pe_encoders_kwargs["out_dim"]
         else:
             raise ValueError("pe_encoders is not initialized, so there is no output dimension.")
-
-    # @property
-    # def in_dim_edges(self) -> int:
-    #     r"""
-    #     Returns the input edge dimension of the network
-    #     """
-    #     return self.gnn.in_dim_edges
diff --git a/goli/nn/architectures/global_architectures.py b/goli/nn/architectures/global_architectures.py
index 0447dacce..b7740114c 100644
--- a/goli/nn/architectures/global_architectures.py
+++ b/goli/nn/architectures/global_architectures.py
@@ -28,12 +28,7 @@
     gaussian_kernel_pos_encoder,
 )
 
-PE_ENCODERS_DICT = {
-    "laplacian_pe": laplace_pos_encoder.LapPENodeEncoder,
-    "mlp": mlp_encoder.MLPEncoder,
-    "signnet": signnet_pos_encoder.SignNetNodeEncoder,
-    "gaussian_kernel": gaussian_kernel_pos_encoder.GaussianKernelPosEncoder,
-}
+import collections
 
 
 class FeedForwardNN(nn.Module):
@@ -78,8 +73,9 @@ def __init__(
 
             depth:
                 If `hidden_dims` is an integer, `depth` is 1 + the number of
-                hidden layers to use. If `hidden_dims` is a `list`, `depth` must
-                be `None`.
+                hidden layers to use.
+                If `hidden_dims` is a list, then
+                `depth` must be `None` or equal to `len(hidden_dims) + 1`
 
             activation:
                 activation function to use in the hidden layers.
@@ -150,7 +146,9 @@ def __init__(
             self.hidden_dims = [hidden_dims] * (depth - 1)
         else:
             self.hidden_dims = list(hidden_dims)
-            assert depth is None
+            assert (depth is None) or (
+                depth == len(self.hidden_dims) + 1
+            ), "Mismatch between the provided network depth from `hidden_dims` and `depth`"
         self.depth = len(self.hidden_dims) + 1
         self.activation = get_activation(activation)
         self.last_activation = get_activation(last_activation)
@@ -383,6 +381,7 @@ def __init__(
         name: str = "GNN",
         layer_kwargs: Optional[Dict] = None,
         virtual_node: str = "none",
+        use_virtual_edges: bool = False,
         last_layer_is_readout: bool = False,
     ):
         r"""
@@ -521,6 +520,9 @@ def __init__(
                 is "none". Otherwise, it will use a simple ResNet like residual
                 connection.
 
+            use_virtual_edges:
+                A bool flag used to select if the virtual node should use the edges or not
+
             last_layer_is_readout: Whether the last layer should be treated as a readout layer.
                 Allows to use the `mup.MuReadout` from the muTransfer method https://github.com/microsoft/mup
 
@@ -542,6 +544,7 @@ def __init__(
         self.virtual_node = virtual_node.lower() if virtual_node is not None else "none"
         self.pooling = pooling
 
+        self.use_virtual_edges = use_virtual_edges
         self.virtual_node_class = self._parse_virtual_node_class()
 
         # Initialize the parent `FeedForwardNN`
@@ -644,6 +647,8 @@ def _create_layers(self):
                     activation=this_activation,
                     dropout=this_dropout,
                     normalization=this_norm,
+                    layer_idx=ii,
+                    layer_depth=self.depth,
                     **self.layer_kwargs,
                     **this_edge_kwargs,
                 )
@@ -653,13 +658,17 @@ def _create_layers(self):
             if ii < len(residual_out_dims):
                 self.virtual_node_layers.append(
                     self.virtual_node_class(
-                        dim=this_out_dim * self.layers[-1].out_dim_factor,
+                        in_dim=this_out_dim * self.layers[-1].out_dim_factor,
+                        out_dim=this_out_dim * self.layers[-1].out_dim_factor,
+                        in_dim_edges=this_out_dim_edges,
+                        out_dim_edges=this_out_dim_edges,
                         activation=this_activation,
                         dropout=this_dropout,
                         normalization=this_norm,
                         bias=True,
                         vn_type=self.virtual_node,
                         residual=self.residual_type is not None,
+                        use_edges=self.use_virtual_edges,
                     )
                 )
 
@@ -818,7 +827,7 @@ def _parse_pooling_layer(
         raise NotImplementedError("Virtual method must be overwritten by child class")
 
     def _virtual_node_forward(
-        self, g: Union[DGLGraph, Data], h: torch.Tensor, vn_h: torch.Tensor, step_idx: int
+        self, g: Union[DGLGraph, Data], h: torch.Tensor, vn_h: torch.Tensor, step_idx: int, e: torch.Tensor
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         r"""
         Apply the *i-th* virtual node layer, where *i* is the index given by `step_idx`.
@@ -851,13 +860,10 @@ def _virtual_node_forward(
                 `M` is the number of graphs, `Dout` is the output features
 
         """
-
-        if step_idx == 0:
-            vn_h = 0.0
         if step_idx < len(self.virtual_node_layers):
-            h, vn_h = self.virtual_node_layers[step_idx].forward(g=g, h=h, vn_h=vn_h)
+            h, vn_h, e = self.virtual_node_layers[step_idx].forward(g=g, h=h, vn_h=vn_h, e=e)
 
-        return h, vn_h
+        return h, vn_h, e
 
     def forward(self, g) -> torch.Tensor:
         r"""
@@ -891,10 +897,9 @@ def forward(self, g) -> torch.Tensor:
         # Initialize values of the residuals and virtual node
         h_prev = None
         e_prev = None
-        vn_h = 0
+        vn_h = 0.0
         h = self._get_node_feats(g, key="h")
         e = self._get_edge_feats(g, key="edge_attr")
-
         # Apply the normalization before the first network layers
         if self.first_normalization is not None:
             h = self.first_normalization(h)
@@ -906,7 +911,7 @@ def forward(self, g) -> torch.Tensor:
             h, e, h_prev, e_prev = self._graph_layer_forward(
                 layer=layer, g=g, h=h, e=e, h_prev=h_prev, e_prev=e_prev, step_idx=ii
             )
-            h, vn_h = self._virtual_node_forward(g=g, h=h, vn_h=vn_h, step_idx=ii)
+            h, vn_h, e = self._virtual_node_forward(g=g, h=h, e=e, vn_h=vn_h, step_idx=ii)
 
         pooled_h = self._pool_layer_forward(g=g, h=h)
 
@@ -972,9 +977,7 @@ def get_init_kwargs(self) -> Dict[str, Any]:
         kwargs.update(new_kwargs)
         return deepcopy(kwargs)
 
-    def make_mup_base_kwargs(
-        self, divide_factor: float = 2.0, factor_in_dim: bool = False, factor_in_dim_edges: bool = False
-    ) -> Dict[str, Any]:
+    def make_mup_base_kwargs(self, divide_factor: float = 2.0, factor_in_dim: bool = False) -> Dict[str, Any]:
         """
         Create a 'base' model to be used by the `mup` or `muTransfer` scaling of the model.
         The base model is usually identical to the regular model, but with the
@@ -983,17 +986,25 @@ def make_mup_base_kwargs(
         Parameter:
             divide_factor: Factor by which to divide the width.
             factor_in_dim: Whether to factor the input dimension for the nodes
-            factor_in_dim: Whether to factor the input dimension for the edges
         """
         kwargs = self.get_init_kwargs()
         kwargs["hidden_dims"] = [round(dim / divide_factor) for dim in kwargs["hidden_dims"]]
         kwargs["hidden_dims_edges"] = [round(dim / divide_factor) for dim in kwargs["hidden_dims_edges"]]
         if factor_in_dim:
             kwargs["in_dim"] = round(kwargs["in_dim"] / divide_factor)
-        if factor_in_dim_edges:
             kwargs["in_dim_edges"] = round(kwargs["in_dim_edges"] / divide_factor)
         if not self.last_layer_is_readout:
             kwargs["out_dim"] = round(kwargs["out_dim"] / divide_factor)
+
+        def _recursive_divide_dim(x: collections.abc.Mapping):
+            for k, v in x.items():
+                if isinstance(v, collections.abc.Mapping):
+                    _recursive_divide_dim(v)
+                elif k in ["in_dim", "out_dim", "in_dim_edges", "out_dim_edges"]:
+                    x[k] = round(v / divide_factor)
+
+        _recursive_divide_dim(kwargs["layer_kwargs"])
+
         return kwargs
 
     def __repr__(self):
@@ -1035,8 +1046,8 @@ def __init__(
                 - gnn_kwargs["out_dim"] must be equal to post_nn_kwargs["in_dim"]
 
             pe_encoders_kwargs:
-                key-word arguments to use for the initialization of all positional encoding encoders
-                can use the class PE_ENCODERS_DICT: "la_encoder"(tested) , "mlp_encoder" (not tested), "signnet_encoder" (not tested)
+                key-word arguments to use for the initialization of all positional encoding encoders.
+                See the class `EncoderManager` for more details.
 
             pre_nn_kwargs:
                 key-word arguments to use for the initialization of the pre-processing
@@ -1352,18 +1363,16 @@ def make_mup_base_kwargs(self, divide_factor: float = 2.0) -> Dict[str, Any]:
         # For the gnn network, all the dimension are divided, except the input dims if pre-nn are missing
         if self.gnn is not None:
             factor_in_dim = self.pre_nn is not None
-            factor_in_dim_edges = self.pre_nn_edges is not None
             kwargs["gnn_kwargs"] = self.gnn.make_mup_base_kwargs(
                 divide_factor=divide_factor,
                 factor_in_dim=factor_in_dim,
-                factor_in_dim_edges=factor_in_dim_edges,
             )
 
         return kwargs
 
     def set_max_num_nodes_edges_per_graph(self, max_nodes: Optional[int], max_edges: Optional[int]) -> None:
         """
-        Set the maximum number of nodes and edges for all gnn layers
+        Set the maximum number of nodes and edges for all gnn layers and encoder layers
 
         Parameters:
             max_nodes: Maximum number of nodes in the dataset.
@@ -1372,6 +1381,10 @@ def set_max_num_nodes_edges_per_graph(self, max_nodes: Optional[int], max_edges:
             max_edges: Maximum number of edges in the dataset.
                 This will be useful for certain architecture, but ignored by others.
         """
+        if (self.encoder_manager is not None) and (self.encoder_manager.pe_encoders is not None):
+            for encoder in self.encoder_manager.pe_encoders.values():
+                encoder.max_num_nodes_per_graph = max_nodes
+                encoder.max_num_edges_per_graph = max_edges
         if self.gnn is not None:
             for layer in self.gnn.layers:
                 if isinstance(layer, BaseGraphStructure):
@@ -1533,7 +1546,7 @@ def __init__(
         pre_nn_edges_kwargs: Optional[Dict[str, Any]] = None,
         post_nn_kwargs: Optional[Dict[str, Any]] = None,
         num_inference_to_average: int = 1,
-        last_layer_is_readout: bool = False,
+        last_layer_is_readout: bool = True,
         name: str = "Multitask_GNN",
     ):
         r"""
@@ -1586,6 +1599,7 @@ def __init__(
                 purposes.
         """
 
+        # Use last_layer_is_readout=False since readout layers are in task heads
         super().__init__(
             gnn_kwargs=gnn_kwargs,
             pre_nn_kwargs=pre_nn_kwargs,
diff --git a/goli/nn/base_graph_layer.py b/goli/nn/base_graph_layer.py
index 6dfff8a5c..a9a207e40 100644
--- a/goli/nn/base_graph_layer.py
+++ b/goli/nn/base_graph_layer.py
@@ -4,10 +4,10 @@
 
 import torch
 import torch.nn as nn
-from torch import Tensor
+from torch import Tensor, IntTensor
 from torch_sparse import SparseTensor
 
-from goli.nn.base_layers import get_activation
+from goli.nn.base_layers import get_activation, DropPath
 from goli.utils.decorators import classproperty
 
 try:
@@ -28,6 +28,9 @@ def __init__(
         activation: Union[str, Callable] = "relu",
         dropout: float = 0.0,
         normalization: Union[str, Callable] = "none",
+        layer_idx: Optional[int] = None,
+        layer_depth: Optional[int] = None,
+        droppath_rate: float = 0.0,
     ):
         r"""
         Abstract class used to standardize the implementation of DGL layers
@@ -57,6 +60,14 @@ def __init__(
                 - "layer_norm": Layer normalization
                 - `Callable`: Any callable function
 
+            layer_idx:
+                The index of the current layer
+
+            layer_depth:
+                The total depth (number of layers) associated to this specific layer
+
+            droppath_rate:
+                stochastic depth drop rate, between 0 and 1, see https://arxiv.org/abs/1603.09382
         """
 
         super().__init__()
@@ -67,6 +78,9 @@ def __init__(
         self.normalization = normalization
         self.dropout = dropout
         self.activation = activation
+        self.layer_idx = layer_idx
+        self.layer_depth = layer_depth
+        self.droppath_rate = droppath_rate
         self._max_num_nodes_per_graph = None
         self._max_num_edges_per_graph = None
 
@@ -81,6 +95,7 @@ def _initialize_activation_dropout_norm(self):
         self.dropout_layer = self._parse_dropout(self.dropout)
 
         self.norm_layer = self._parse_norm(self.normalization)
+        self.droppath_layer = self._parse_droppath(self.droppath_rate)
 
     def _parse_dropout(self, dropout):
         if callable(dropout):
@@ -89,6 +104,13 @@ def _parse_dropout(self, dropout):
             return nn.Dropout(p=dropout)
         return
 
+    def _parse_droppath(self, droppath_rate):
+        if droppath_rate == 0:
+            return
+
+        droppath_rate = DropPath.get_stochastic_drop_rate(droppath_rate, self.layer_idx, self.layer_depth)
+        return DropPath(drop_rate=droppath_rate)
+
     def _parse_norm(self, normalization, dim=None):
         if dim is None:
             dim = self.out_dim * self.out_dim_factor
@@ -112,6 +134,9 @@ def apply_norm_activation_dropout(
         normalization: bool = True,
         activation: bool = True,
         dropout: bool = True,
+        droppath: bool = True,
+        batch_idx: Optional[IntTensor] = None,
+        batch_size: Optional[int] = None,
     ):
         r"""
         Apply the different normalization and the dropout to the
@@ -122,6 +147,8 @@ def apply_norm_activation_dropout(
             h:
                 Feature tensor, to be normalized
 
+            batch_idx
+
             normalization:
                 Whether to apply the normalization
 
@@ -131,6 +158,9 @@ def apply_norm_activation_dropout(
             dropout:
                 Whether to apply the dropout layer
 
+            droppath:
+                Whether to apply the DropPath layer
+
         Returns:
 
             h:
@@ -147,6 +177,9 @@ def apply_norm_activation_dropout(
         if dropout and (self.dropout_layer is not None):
             h = self.dropout_layer(h)
 
+        if droppath and (self.droppath_layer is not None):
+            h = self.droppath_layer(h, batch_idx=batch_idx, batch_size=batch_size)
+
         return h
 
     @classproperty
@@ -269,6 +302,9 @@ def __init__(
         activation: Union[str, Callable] = "relu",
         dropout: float = 0.0,
         normalization: Union[str, Callable] = "none",
+        layer_idx: Optional[int] = None,
+        layer_depth: Optional[int] = None,
+        droppath_rate: float = 0.0,
     ):
         r"""
         Abstract class used to standardize the implementation of DGL layers
@@ -297,6 +333,15 @@ def __init__(
                 - "batch_norm": Batch normalization
                 - "layer_norm": Layer normalization
                 - `Callable`: Any callable function
+
+            layer_idx:
+                The index of the current layer
+
+            layer_depth:
+                The total depth (number of layers) associated to this specific layer
+
+            droppath_rate:
+                stochastic depth drop rate, between 0 and 1, see https://arxiv.org/abs/1603.09382
         """
 
         super().__init__(
@@ -305,6 +350,9 @@ def __init__(
             normalization=normalization,
             dropout=dropout,
             activation=activation,
+            layer_idx=layer_idx,
+            layer_depth=layer_depth,
+            droppath_rate=droppath_rate,
         )
 
         self._initialize_activation_dropout_norm()
diff --git a/goli/nn/base_layers.py b/goli/nn/base_layers.py
index 67d74a41f..ca2e95216 100644
--- a/goli/nn/base_layers.py
+++ b/goli/nn/base_layers.py
@@ -1,16 +1,16 @@
-from typing import Union, Callable, Optional, Type, Tuple
+from typing import Union, Callable, Optional, Type, Tuple, Iterable
 from copy import deepcopy
 from loguru import logger
 
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from torch import Tensor
+from torch import Tensor, IntTensor
 import mup.init as mupi
 from mup import set_base_shapes, MuReadout
 from torch.nn.functional import linear
 
-from goli.ipu.ipu_utils import import_poptorch
+from goli.ipu.ipu_utils import is_running_on_ipu
 
 SUPPORTED_ACTIVATION_MAP = {
     "ReLU",
@@ -46,12 +46,38 @@ def get_activation(activation: Union[type(None), str, Callable]) -> Optional[Cal
 
     # search in SUPPORTED_ACTIVATION_MAP a torch.nn.modules.activation
     activation = [x for x in SUPPORTED_ACTIVATION_MAP if activation.lower() == x.lower()]
-    assert len(activation) == 1 and isinstance(activation[0], str), "Unhandled activation function"
+    assert len(activation) == 1 and isinstance(
+        activation[0], str
+    ), f"Unhandled activation function {activation} of type {type(activation)}"
     activation = activation[0]
 
     return vars(torch.nn.modules.activation)[activation]()
 
 
+def get_activation_str(activation: Union[type(None), str, Callable]) -> str:
+    r"""
+    returns the string related to the activation function
+
+    Parameters:
+        activation: Callable, `None`, or string with value:
+            "none", "ReLU", "Sigmoid", "Tanh", "ELU", "SELU", "GLU", "LeakyReLU", "Softplus"
+
+    Returns:
+        The name of the activation function
+    """
+
+    if isinstance(activation, str):
+        return activation
+
+    if activation is None:
+        return "None"
+
+    if isinstance(activation, Callable):
+        return activation.__class__._get_name(activation)
+    else:
+        raise ValueError(f"Unhandled activation function {activation} of type {type(activation)}")
+
+
 def get_norm(normalization: Union[Type[None], str, Callable], dim: Optional[int] = None):
     r"""
     returns the normalization function represented by the input string
@@ -141,6 +167,7 @@ def forward(
             # [batch, num_heads, nodes, nodes]
             attn_weights = q @ k.transpose(-1, -2)
             # [batch, num_heads, nodes, nodes]
+
             attn_weights += attn_bias
             # key_padding_mask: [batch, 1, 1, nodes]
             if key_padding_mask is not None:
@@ -301,6 +328,7 @@ def __init__(
         bias: bool = True,
         init_fn: Optional[Callable] = None,
         is_readout_layer: bool = False,
+        droppath_rate: float = 0.0,
     ):
         r"""
         A simple fully connected and customizable layer. This layer is centered around a `torch.nn.Linear` module.
@@ -335,6 +363,8 @@ def __init__(
             is_readout_layer: Whether the layer should be treated as a readout layer by replacing of `torch.nn.Linear`
                 by `mup.MuReadout` from the muTransfer method https://github.com/microsoft/mup
 
+            droppath_rate:
+                stochastic depth drop rate, between 0 and 1, see https://arxiv.org/abs/1603.09382
         Attributes:
             dropout (int):
                 The ratio of units to dropout.
@@ -370,6 +400,10 @@ def __init__(
             self.dropout = nn.Dropout(p=dropout)
         self.activation = get_activation(activation)
 
+        self.drop_path = None
+        if droppath_rate > 0:
+            self.drop_path = DropPath(droppath_rate)
+
         # Linear layer, or MuReadout layer
         if not is_readout_layer:
             self.linear = nn.Linear(in_dim, out_dim, bias=bias)
@@ -437,6 +471,9 @@ def forward(self, h: torch.Tensor) -> torch.Tensor:
             h = self.dropout(h)
         if self.activation is not None:
             h = self.activation(h)
+        if self.drop_path is not None:
+            h = self.drop_path(h)
+
         return h
 
     @property
@@ -461,9 +498,9 @@ class MLP(nn.Module):
     def __init__(
         self,
         in_dim: int,
-        hidden_dim: int,
+        hidden_dims: Union[Iterable[int], int],
         out_dim: int,
-        layers: int,
+        depth: int,
         activation: Union[str, Callable] = "relu",
         last_activation: Union[str, Callable] = "none",
         dropout: float = 0.0,
@@ -472,6 +509,8 @@ def __init__(
         last_normalization: Union[Type[None], str, Callable] = "none",
         first_normalization: Union[Type[None], str, Callable] = "none",
         last_layer_is_readout: bool = False,
+        droppath_rate: float = 0.0,
+        constant_droppath_rate: bool = True,
     ):
         r"""
         Simple multi-layer perceptron, built of a series of FCLayers
@@ -479,13 +518,16 @@ def __init__(
         Parameters:
             in_dim:
                 Input dimension of the MLP
-            hidden_dim:
-                Hidden dimension of the MLP. All hidden dimensions will have
-                the same number of parameters
+            hidden_dims:
+                Either an integer specifying all the hidden dimensions,
+                or a list of dimensions in the hidden layers.
             out_dim:
                 Output dimension of the MLP.
-            layers:
-                Number of hidden layers
+            depth:
+                If `hidden_dims` is an integer, `depth` is 1 + the number of
+                hidden layers to use.
+                If `hidden_dims` is a list, then
+                `depth` must be `None` or equal to `len(hidden_dims) + 1`
             activation:
                 Activation function to use in all the layers except the last.
                 if `layers==1`, this parameter is ignored
@@ -510,61 +552,71 @@ def __init__(
                 The ratio of units to dropout at the last layer.
             last_layer_is_readout: Whether the last layer should be treated as a readout layer.
                 Allows to use the `mup.MuReadout` from the muTransfer method https://github.com/microsoft/mup
-
+            droppath_rate:
+                stochastic depth drop rate, between 0 and 1.
+                See https://arxiv.org/abs/1603.09382
+            constant_droppath_rate:
+                If `True`, drop rates will remain constant accross layers.
+                Otherwise, drop rates will vary stochastically.
+                See `DropPath.get_stochastic_drop_rate`
         """
 
         super().__init__()
 
         self.in_dim = in_dim
-        self.hidden_dim = hidden_dim
         self.out_dim = out_dim
+
+        # Parse the hidden dimensions and depth
+        if isinstance(hidden_dims, int):
+            self.hidden_dims = [hidden_dims] * (depth - 1)
+        else:
+            self.hidden_dims = list(hidden_dims)
+            assert (depth is None) or (
+                depth == len(self.hidden_dims) + 1
+            ), "Mismatch between the provided network depth from `hidden_dims` and `depth`"
+        self.depth = len(self.hidden_dims) + 1
+
+        # Parse the normalization
         self.first_normalization = get_norm(first_normalization, dim=in_dim)
 
+        all_dims = [in_dim] + self.hidden_dims + [out_dim]
         fully_connected = []
-        if layers == 0:
+        if depth == 0:
             self.fully_connected = None
             return
-        elif layers == 1:
-            fully_connected.append(
-                FCLayer(
-                    in_dim,
-                    out_dim,
-                    activation=last_activation,
-                    normalization=last_normalization,
-                    dropout=last_dropout,
-                    is_readout_layer=last_layer_is_readout,
-                )
-            )
-        elif layers > 1:
-            fully_connected.append(
-                FCLayer(
-                    in_dim,
-                    hidden_dim,
-                    activation=activation,
-                    normalization=normalization,
-                    dropout=dropout,
-                )
-            )
-            for _ in range(layers - 2):
+        else:
+            for ii in range(depth):
+                if ii < (depth - 1):
+                    # Define the parameters for all intermediate layers
+                    this_activation = activation
+                    this_normalization = normalization
+                    this_dropout = dropout
+                    is_readout_layer = False
+                else:
+                    # Define the parameters for the last layer
+                    this_activation = last_activation
+                    this_normalization = last_normalization
+                    this_dropout = last_dropout
+                    is_readout_layer = last_layer_is_readout
+
+                if constant_droppath_rate:
+                    this_drop_rate = droppath_rate
+                else:
+                    this_drop_rate = DropPath.get_stochastic_drop_rate(droppath_rate, ii, depth)
+
+                # Add a fully-connected layer
                 fully_connected.append(
                     FCLayer(
-                        hidden_dim,
-                        hidden_dim,
-                        activation=activation,
-                        normalization=normalization,
-                        dropout=dropout,
+                        all_dims[ii],
+                        all_dims[ii + 1],
+                        activation=this_activation,
+                        normalization=this_normalization,
+                        dropout=this_dropout,
+                        is_readout_layer=is_readout_layer,
+                        droppath_rate=this_drop_rate,
                     )
                 )
-            fully_connected.append(
-                FCLayer(
-                    hidden_dim,
-                    out_dim,
-                    activation=last_activation,
-                    normalization=last_normalization,
-                    dropout=last_dropout,
-                    is_readout_layer=last_layer_is_readout,
-                )
-            )
+
         self.fully_connected = nn.Sequential(*fully_connected)
 
     def forward(self, h: torch.Tensor) -> torch.Tensor:
@@ -653,3 +705,96 @@ def forward(self, x, y):
         x = self.gru(x, y)[1]
         x = x.reshape(B, N, -1)
         return x
+
+
+class DropPath(nn.Module):
+    def __init__(self, drop_rate: float):
+        r"""
+        DropPath class for stochastic depth
+        Deep Networks with Stochastic Depth
+        Gao Huang, Yu Sun, Zhuang Liu, Daniel Sedra and Kilian Weinberger
+        https://arxiv.org/abs/1603.09382
+
+        Parameters:
+            drop_rate:
+                Drop out probability
+        """
+
+        super().__init__()
+        self.drop_rate = drop_rate
+
+    @staticmethod
+    def get_stochastic_drop_rate(
+        drop_rate: float, layer_idx: Optional[int] = None, layer_depth: Optional[int] = None
+    ):
+        """
+        Get the stochastic drop rate from the nominal drop rate, the layer index, and the layer depth.
+
+        `return drop_rate * (layer_idx / (layer_depth - 1))`
+
+        Parameters:
+            drop_rate:
+                Drop out nominal probability
+
+            layer_idx:
+                The index of the current layer
+
+            layer_depth:
+                The total depth (number of layers) associated to this specific layer
+
+        """
+        if drop_rate == 0:
+            return 0
+        else:
+            assert (layer_idx is not None) and (
+                layer_depth is not None
+            ), f"layer_idx={layer_idx} and layer_depth={layer_depth} should be integers when `droppath_rate>0`"
+            return drop_rate * (layer_idx / (layer_depth - 1))
+
+    def forward(
+        self,
+        input: Tensor,
+        batch_idx: IntTensor,
+        batch_size: Optional[int] = None,
+    ) -> Tensor:
+        r"""
+        Parameters:
+            input:  `torch.Tensor[total_num_nodes, hidden]`
+            batch: batch attribute of the batch object, batch.batch
+            batch_size: The batch size. Must be provided when working on IPU
+
+        Returns:
+            torch.Tensor: `torch.Tensor[total_num_nodes, hidde]`
+
+        """
+        on_ipu = is_running_on_ipu()
+
+        if self.drop_rate > 0:
+            keep_prob = 1 - self.drop_rate
+
+            # Parse the batch size
+            if batch_size is None:
+                if on_ipu:
+                    raise ValueError(
+                        "When using the IPU the batch size must be "
+                        "provided during compilation instead of determined at runtime"
+                    )
+                else:
+                    batch_size = int(batch_idx.max()) + 1
+
+            # mask shape: [num_graphs, 1]
+            mask = input.new_empty(batch_size, 1).bernoulli_(keep_prob)
+            # if on_ipu, the last graph is a padded fake graph
+            if on_ipu:
+                mask[-1] = 0
+            # using gather to extend mask to [total_num_nodes, 1]
+            node_mask = mask[batch_idx]
+            if keep_prob == 0:
+                # avoid dividing by 0
+                input_scaled = input
+            else:
+                input_scaled = input / keep_prob
+            out = input_scaled * node_mask
+        else:
+            out = input
+        return out
diff --git a/goli/nn/dgl_layers/gat_dgl.py b/goli/nn/dgl_layers/gat_dgl.py
index 715db6577..e069dd3bb 100644
--- a/goli/nn/dgl_layers/gat_dgl.py
+++ b/goli/nn/dgl_layers/gat_dgl.py
@@ -23,6 +23,7 @@ def __init__(
         activation="elu",
         dropout: float = 0.0,
         normalization: Union[str, Callable] = "none",
+        **kwargs,
     ):
         r"""
         GAT: Graph Attention Network
@@ -65,6 +66,7 @@ def __init__(
             activation=activation,
             dropout=dropout,
             normalization=normalization,
+            **kwargs,
         )
 
         self.gatconv = GATConv(
diff --git a/goli/nn/dgl_layers/gated_gcn_dgl.py b/goli/nn/dgl_layers/gated_gcn_dgl.py
index a6a3cad98..3632a431e 100644
--- a/goli/nn/dgl_layers/gated_gcn_dgl.py
+++ b/goli/nn/dgl_layers/gated_gcn_dgl.py
@@ -24,6 +24,7 @@ def __init__(
         activation: Union[Callable, str] = "relu",
         dropout: float = 0.0,
         normalization: Union[str, Callable] = "none",
+        **kwargs,
     ):
         r"""
         ResGatedGCN: Residual Gated Graph ConvNets
@@ -63,6 +64,7 @@ def __init__(
             activation=activation,
             dropout=dropout,
             normalization=normalization,
+            **kwargs,
         )
 
         self.out_dim_edges = out_dim_edges
diff --git a/goli/nn/dgl_layers/gcn_dgl.py b/goli/nn/dgl_layers/gcn_dgl.py
index fe640f3b5..6b806db29 100644
--- a/goli/nn/dgl_layers/gcn_dgl.py
+++ b/goli/nn/dgl_layers/gcn_dgl.py
@@ -22,6 +22,7 @@ def __init__(
         activation: Union[str, Callable] = "relu",
         dropout: float = 0.0,
         normalization: Union[str, Callable] = "none",
+        **kwargs,
     ):
         r"""
         Graph convolutional network (GCN) layer from
@@ -57,6 +58,7 @@ def __init__(
             activation=activation,
             dropout=dropout,
             normalization=normalization,
+            **kwargs,
         )
 
         self.conv = GraphConv(
diff --git a/goli/nn/dgl_layers/gin_dgl.py b/goli/nn/dgl_layers/gin_dgl.py
index 40dbdba2b..fb3d356c2 100644
--- a/goli/nn/dgl_layers/gin_dgl.py
+++ b/goli/nn/dgl_layers/gin_dgl.py
@@ -24,6 +24,7 @@ def __init__(
         normalization: Union[str, Callable] = "none",
         init_eps: float = 0.0,
         learn_eps: bool = True,
+        **kwargs,
     ):
         r"""
         GIN: Graph Isomorphism Networks
@@ -68,6 +69,7 @@ def __init__(
             activation=activation,
             dropout=dropout,
             normalization=normalization,
+            **kwargs,
         )
 
         # Specify to consider the edges weight in the aggregation
@@ -81,9 +83,9 @@ def __init__(
         # The weights of the model, applied after the aggregation
         self.mlp = MLP(
             in_dim=self.in_dim,
-            hidden_dim=self.in_dim,
+            hidden_dims=self.in_dim,
             out_dim=self.out_dim,
-            layers=2,
+            depth=2,
             activation=self.activation_layer,
             last_activation="none",
             normalization=self.normalization,
diff --git a/goli/nn/dgl_layers/pna_dgl.py b/goli/nn/dgl_layers/pna_dgl.py
index 56e51724b..f4541efd8 100644
--- a/goli/nn/dgl_layers/pna_dgl.py
+++ b/goli/nn/dgl_layers/pna_dgl.py
@@ -29,6 +29,7 @@ def __init__(
         avg_d: float = 1.0,
         last_activation: Union[Callable, str] = "none",
         in_dim_edges: int = 0,
+        **kwargs,
     ):
         r"""
         Abstract class used to standardize the implementation of PNA layers
@@ -89,6 +90,7 @@ def __init__(
             activation=activation,
             dropout=dropout,
             normalization=normalization,
+            **kwargs,
         )
 
         # Edge dimensions
@@ -273,6 +275,7 @@ def __init__(
         last_activation: Union[Callable, str] = "none",
         posttrans_layers: int = 1,
         in_dim_edges: int = 0,
+        **kwargs,
     ):
         r"""
 
@@ -333,14 +336,15 @@ def __init__(
             normalization="none",
             last_activation=last_activation,
             in_dim_edges=in_dim_edges,
+            **kwargs,
         )
 
         # MLP used on the aggregated messages of the neighbours
         self.posttrans = MLP(
             in_dim=(len(aggregators) * len(scalers)) * (self.in_dim + self.in_dim_edges),
-            hidden_dim=self.out_dim,
+            hidden_dims=self.out_dim,
             out_dim=self.out_dim,
-            layers=posttrans_layers,
+            depth=posttrans_layers,
             activation=self.activation,
             last_activation=self.last_activation,
             dropout=dropout,
@@ -502,14 +506,15 @@ def __init__(
             normalization="none",
             last_activation=last_activation,
             in_dim_edges=in_dim_edges,
+            **kwargs,
         )
 
         # MLP used on each pair of nodes with their edge MLP(h_u, h_v, e_uv)
         self.pretrans = MLP(
             in_dim=2 * in_dim + in_dim_edges,
-            hidden_dim=in_dim,
+            hidden_dims=in_dim,
             out_dim=in_dim,
-            layers=pretrans_layers,
+            depth=pretrans_layers,
             activation=self.activation,
             last_activation=self.last_activation,
             dropout=dropout,
@@ -520,9 +525,9 @@ def __init__(
         # MLP used on the aggregated messages MLP(h'_u)
         self.posttrans = MLP(
             in_dim=(len(self.aggregators) * len(self.scalers) + 1) * in_dim,
-            hidden_dim=out_dim,
+            hidden_dims=out_dim,
             out_dim=out_dim,
-            layers=posttrans_layers,
+            depth=posttrans_layers,
             activation=self.activation,
             last_activation=self.last_activation,
             dropout=dropout,
diff --git a/goli/nn/dgl_layers/pooling_dgl.py b/goli/nn/dgl_layers/pooling_dgl.py
index 2691b136a..8196523ab 100644
--- a/goli/nn/dgl_layers/pooling_dgl.py
+++ b/goli/nn/dgl_layers/pooling_dgl.py
@@ -36,9 +36,9 @@ def __init__(
         # fully connected layers
         self.mlp = MLP(
             in_dim=2 * in_dim,
-            hidden_dim=hidden_dim,
+            hidden_dims=hidden_dim,
             out_dim=out_dim,
-            layers=fc_layers,
+            depth=fc_layers,
             activation="relu",
             last_activation=final_activation,
             normalization="batch_norm",
@@ -247,13 +247,17 @@ def parse_pooling_layer_dgl(in_dim: int, pooling: Union[str, List[str]], **kwarg
 class VirtualNodeDgl(nn.Module):
     def __init__(
         self,
-        dim: int,
+        in_dim: int,
+        out_dim: int,
+        in_dim_edges: int,
+        out_dim_edges: int,
         vn_type: Union[type(None), str] = "sum",
         activation: Union[str, Callable] = "relu",
         dropout: float = 0.0,
         normalization: Union[str, Callable] = "none",
         bias: bool = True,
         residual: bool = True,
+        use_edges: bool = False,
     ):
         r"""
         The VirtualNode is a layer that pool the features of the graph,
@@ -262,7 +266,7 @@ def __init__(
 
         Parameters:
 
-            dim:
+            in_dim:
                 Input and output feature dimensions of the virtual node layer
 
             activation:
@@ -286,6 +290,17 @@ def __init__(
                 Whether all virtual nodes should be connected together
                 via a residual connection
 
+            out_dim:
+                NOT USED: Out node dimmension for virtual node pooling
+            in_dim_edges:
+                NOT USED: Edge dimmension for virtual node pooling
+            out_dim_edges:
+                NOT USED: Edge dimmension for virtual node pooling
+
+            use_edges:
+                NOT USED: boolean to choose using edges or not in virtual node
+
+
         """
         super().__init__()
         if (vn_type is None) or (vn_type.lower() == "none"):
@@ -297,8 +312,8 @@ def __init__(
         self.vn_type = vn_type.lower()
         self.residual = residual
         self.fc_layer = FCLayer(
-            in_dim=dim,
-            out_dim=dim,
+            in_dim=in_dim,
+            out_dim=in_dim,
             activation=activation,
             dropout=dropout,
             normalization=normalization,
@@ -306,7 +321,7 @@ def __init__(
         )
 
     def forward(
-        self, g: dgl.DGLGraph, h: torch.Tensor, vn_h: torch.Tensor
+        self, g: dgl.DGLGraph, h: torch.Tensor, vn_h: torch.Tensor, e: torch.Tensor
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         r"""
         Apply the virtual node layer.
@@ -341,7 +356,7 @@ def forward(
 
         # Pool the features
         if self.vn_type is None:
-            return h, vn_h
+            return h, vn_h, e
         elif self.vn_type == "mean":
             pool = mean_nodes(g, "h")
         elif self.vn_type == "max":
@@ -371,4 +386,4 @@ def forward(
         )
         h = h + temp_h
 
-        return h, vn_h
+        return h, vn_h, e
diff --git a/goli/nn/encoders/gaussian_kernel_pos_encoder.py b/goli/nn/encoders/gaussian_kernel_pos_encoder.py
index c20b59620..c4faed7f0 100644
--- a/goli/nn/encoders/gaussian_kernel_pos_encoder.py
+++ b/goli/nn/encoders/gaussian_kernel_pos_encoder.py
@@ -2,7 +2,7 @@
 from torch_geometric.data import Batch
 
 from goli.nn.pyg_layers.utils import PreprocessPositions
-from goli.ipu.ipu_utils import import_poptorch
+from goli.ipu.ipu_utils import is_running_on_ipu
 from goli.nn.encoders.base_encoder import BaseEncoder
 
 
@@ -21,7 +21,7 @@ def __init__(
         out_dim: int,
         embed_dim: int,
         num_layers: int,
-        max_num_nodes_per_graph: int,
+        max_num_nodes_per_graph: Optional[int] = None,
         activation: Union[str, Callable] = "gelu",
         first_normalization="none",
         use_input_keys_prefix: bool = True,
@@ -73,8 +73,7 @@ def parse_output_keys(self, output_keys):
     def forward(self, batch: Batch, key_prefix: Optional[str] = None) -> Dict[str, Any]:
         input_keys = self.parse_input_keys_with_prefix(key_prefix)
 
-        poptorch = import_poptorch(raise_error=False)  # TODO: Change to `is_running_on_ipu` after merge
-        on_ipu = (poptorch is not None) and (poptorch.isRunningOnIpu())
+        on_ipu = is_running_on_ipu()
         max_num_nodes_per_graph = None
         if on_ipu:
             max_num_nodes_per_graph = self.max_num_nodes_per_graph
diff --git a/goli/nn/encoders/laplace_pos_encoder.py b/goli/nn/encoders/laplace_pos_encoder.py
index 305a8ea88..1260b332b 100644
--- a/goli/nn/encoders/laplace_pos_encoder.py
+++ b/goli/nn/encoders/laplace_pos_encoder.py
@@ -74,9 +74,9 @@ def __init__(
             # DeepSet model for LapPE (this will be followed by a sum pooling)
             self.pe_encoder = MLP(
                 in_dim=hidden_dim,
-                hidden_dim=hidden_dim,
+                hidden_dims=hidden_dim,
                 out_dim=hidden_dim,
-                layers=num_layers,
+                depth=num_layers,
                 dropout=dropout,
                 **model_kwargs,
             )
@@ -84,9 +84,9 @@ def __init__(
             # MLP that will mix all eigenvalues and eigenvectors
             self.pe_encoder = MLP(
                 in_dim=self.in_dim * hidden_dim,
-                hidden_dim=hidden_dim,
+                hidden_dims=hidden_dim,
                 out_dim=hidden_dim,
-                layers=num_layers_post,
+                depth=num_layers_post,
                 dropout=dropout,
                 activation=activation,
                 last_activation="none",
@@ -100,9 +100,9 @@ def __init__(
             # MLP to apply post pooling
             self.post_mlp = MLP(
                 in_dim=hidden_dim,
-                hidden_dim=hidden_dim,
+                hidden_dims=hidden_dim,
                 out_dim=out_dim,
-                layers=num_layers_post,
+                depth=num_layers_post,
                 dropout=dropout,
                 activation=activation,
                 last_activation="none",
@@ -136,7 +136,7 @@ def forward(self, batch: Batch, key_prefix: Optional[str] = None) -> Dict[str, t
 
         # Random flipping to the Laplacian encoder
         if self.training:
-            sign_flip = torch.rand(eigvecs.size(1), device=eigvecs.device)
+            sign_flip = torch.rand(eigvecs.size(1), device=eigvecs.device, dtype=eigvecs.dtype)
             sign_flip[sign_flip >= 0.5] = 1.0
             sign_flip[sign_flip < 0.5] = -1.0
             eigvecs = eigvecs * sign_flip.unsqueeze(0)
diff --git a/goli/nn/encoders/mlp_encoder.py b/goli/nn/encoders/mlp_encoder.py
index ac1adb583..530ba5688 100644
--- a/goli/nn/encoders/mlp_encoder.py
+++ b/goli/nn/encoders/mlp_encoder.py
@@ -57,9 +57,9 @@ def __init__(
         # Initialize the MLP
         self.pe_encoder = MLP(
             in_dim=in_dim,
-            hidden_dim=hidden_dim,
+            hidden_dims=hidden_dim,
             out_dim=out_dim,
-            layers=num_layers,
+            depth=num_layers,
             dropout=dropout,
             activation=activation,
             last_activation=activation,
diff --git a/goli/nn/pyg_layers/gated_gcn_pyg.py b/goli/nn/pyg_layers/gated_gcn_pyg.py
index e52772f50..88e2e57bc 100644
--- a/goli/nn/pyg_layers/gated_gcn_pyg.py
+++ b/goli/nn/pyg_layers/gated_gcn_pyg.py
@@ -72,6 +72,7 @@ def __init__(
             activation=activation,
             dropout=dropout,
             normalization=normalization,
+            **kwargs,
         )
 
         self._initialize_activation_dropout_norm()
@@ -110,7 +111,7 @@ def forward(self, batch: Union[Data, Batch]):
 
         # Propagate, and apply norm, activation, dropout
         x, e = self.propagate(edge_index, Bx=Bx, Dx=Dx, Ex=Ex, Ce=Ce, e=e, Ax=Ax)
-        x = self.apply_norm_activation_dropout(x)
+        x = self.apply_norm_activation_dropout(x, batch_idx=batch.batch)
         e = self.edge_out(e)
 
         # Output
diff --git a/goli/nn/pyg_layers/gin_pyg.py b/goli/nn/pyg_layers/gin_pyg.py
index 47572814d..510765eac 100644
--- a/goli/nn/pyg_layers/gin_pyg.py
+++ b/goli/nn/pyg_layers/gin_pyg.py
@@ -73,13 +73,14 @@ def __init__(
             activation=activation,
             dropout=dropout,
             normalization=normalization,
+            **kwargs,
         )
 
         gin_nn = MLP(
             in_dim=self.in_dim,
-            hidden_dim=self.in_dim,
+            hidden_dims=self.in_dim,
             out_dim=self.out_dim,
-            layers=2,
+            depth=2,
             activation=self.activation_layer,
             last_activation="none",
             normalization=self.normalization,
@@ -91,7 +92,7 @@ def __init__(
 
     def forward(self, batch: Union[Data, Batch]):
         batch.h = self.model(batch.h, batch.edge_index)
-        batch.h = self.apply_norm_activation_dropout(batch.h)
+        batch.h = self.apply_norm_activation_dropout(batch.h, batch_idx=batch.batch)
 
         return batch
 
@@ -217,9 +218,9 @@ def __init__(
 
         gin_nn = MLP(
             in_dim=self.in_dim,
-            hidden_dim=self.in_dim,
+            hidden_dims=self.in_dim,
             out_dim=self.out_dim,
-            layers=2,
+            depth=2,
             activation=self.activation_layer,
             last_activation="none",
             normalization=self.normalization,
@@ -231,7 +232,7 @@ def __init__(
 
     def forward(self, batch):
         batch.h = self.model(batch.h, batch.edge_index, batch.edge_attr)
-        batch.h = self.apply_norm_activation_dropout(batch.h)
+        batch.h = self.apply_norm_activation_dropout(batch.h, batch_idx=batch.batch)
 
         return batch
 
diff --git a/goli/nn/pyg_layers/gps_pyg.py b/goli/nn/pyg_layers/gps_pyg.py
index fdda74f21..380b380f3 100644
--- a/goli/nn/pyg_layers/gps_pyg.py
+++ b/goli/nn/pyg_layers/gps_pyg.py
@@ -3,10 +3,13 @@
 """
 
 from copy import deepcopy
-from typing import Callable, Union, Optional
+from typing import Callable, Union, Optional, Dict, Any
+from torch.nn import Module
+from torch import Tensor
+from torch_geometric.data import Batch
 
 from goli.nn.base_graph_layer import BaseGraphModule
-from goli.nn.base_layers import FCLayer, MultiheadAttentionMup
+from goli.nn.base_layers import FCLayer, MultiheadAttentionMup, MLP, get_activation_str
 from goli.nn.pyg_layers import (
     GatedGCNPyg,
     GINConvPyg,
@@ -16,7 +19,7 @@
 )
 from goli.utils.decorators import classproperty
 from goli.ipu.to_dense_batch import to_dense_batch, to_sparse_batch
-from goli.ipu.ipu_utils import import_poptorch
+from goli.ipu.ipu_utils import is_running_on_ipu
 
 PYG_LAYERS_DICT = {
     "pyg:gin": GINConvPyg,
@@ -43,28 +46,38 @@ def __init__(
         dropout: float = 0.0,
         node_residual: Optional[bool] = True,
         normalization: Union[str, Callable] = "none",
-        num_gaussian_kernels: Optional[int] = 128,
         mpnn_type: str = "pyg:gine",
         mpnn_kwargs=None,
         attn_type: str = "full-attention",
         biased_attention_key: Optional[str] = None,
         attn_kwargs=None,
+        droppath_rate_attn: float = 0.0,
+        droppath_rate_ffn: float = 0.0,
+        hidden_dim_scaling: float = 4.0,
+        **kwargs,
     ):
         r"""
-        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
+        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
 
-        [!] code uses the pytorch-geometric implementation of GINEConv
+        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 layer
+                Input node feature dimensions of the layer
 
             out_dim:
-                Output feature dimensions of the layer
+                Output node feature dimensions of the layer
+
+            in_dim:
+                Input edge feature dimensions of the layer
+
+            out_dim:
+                Output edge feature dimensions of the layer
 
             activation:
                 activation function to use in the layer
@@ -83,8 +96,23 @@ def __init__(
                 - "layer_norm": Layer normalization
                 - `Callable`: Any callable function
 
-            num_gaussian_kernels
-                Number of Gaussian kernels to use for the attention layer
+            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
+
+            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
@@ -100,63 +128,47 @@ def __init__(
             out_dim=out_dim,
             activation=activation,
             dropout=dropout,
-            normalization=None,
+            normalization=normalization,
+            droppath_rate=droppath_rate_attn,
+            **kwargs,
         )
+        # Set the other attributes
+        self.in_dim_edges = in_dim_edges
+        self.out_dim_edges = out_dim_edges
 
         # Dropout layers
         self.dropout_local = self.dropout_layer
         self.dropout_attn = self._parse_dropout(dropout=self.dropout)
-        self.ff_dropout1 = self._parse_dropout(dropout=self.dropout)
-        self.ff_dropout2 = self._parse_dropout(dropout=self.dropout)
+
+        # DropPath layers
+        self.droppath_ffn = self._parse_droppath(droppath_rate_ffn)
 
         # Residual connections
         self.node_residual = node_residual
 
-        # linear layers
-        self.ff_linear1 = FCLayer(in_dim, in_dim * 2, activation=None)
-        self.ff_linear2 = FCLayer(in_dim * 2, in_dim, activation=None)
-        self.ff_out = FCLayer(in_dim, out_dim, activation=None)
+        # MLP applied at the end of the GPS layer
+        self.mlp = MLP(
+            in_dim=in_dim,
+            hidden_dims=int(hidden_dim_scaling * in_dim),
+            out_dim=in_dim,
+            depth=2,
+            activation=activation,
+            dropout=self.dropout,
+        )
+        self.f_out = FCLayer(in_dim, out_dim, normalization=normalization)
 
         # Normalization layers
         self.norm_layer_local = self._parse_norm(normalization=self.normalization, dim=in_dim)
         self.norm_layer_attn = self._parse_norm(normalization=self.normalization, dim=in_dim)
         self.norm_layer_ff = self._parse_norm(self.normalization)
 
-        # Set the default values for the Attention layer
-        if attn_kwargs is None:
-            attn_kwargs = {}
-        attn_kwargs.setdefault("embed_dim", in_dim)
-        attn_kwargs.setdefault("num_heads", 1)
-        attn_kwargs.setdefault("dropout", dropout)
-        attn_kwargs.setdefault("batch_first", True)
-
-        self.num_gaussian_kernels = num_gaussian_kernels
         self.biased_attention_key = biased_attention_key
-        # Set the default values for the MPNN layer
-        if mpnn_kwargs is None:
-            mpnn_kwargs = {}
-        mpnn_kwargs = deepcopy(mpnn_kwargs)
-        mpnn_kwargs.setdefault("in_dim", in_dim)
-        mpnn_kwargs.setdefault("out_dim", in_dim)
-        mpnn_kwargs.setdefault("in_dim_edges", in_dim_edges)
-        mpnn_kwargs.setdefault("out_dim_edges", out_dim_edges)
-        # TODO: The rest of default values
 
-        # Initialize the MPNN layer
-        mpnn_class = PYG_LAYERS_DICT[mpnn_type]
-        self.mpnn = mpnn_class(**mpnn_kwargs)
-
-        # Initialize the Attention layer
-        self.attn_layer = self._parse_attn_layer(attn_type, self.biased_attention_key, **attn_kwargs)
-
-    def forward(self, batch):
-        # Check whether the model runs on IPU, if so define a maximal number of nodes per graph when reshaping
-        poptorch = import_poptorch(raise_error=False)
-        on_ipu = (poptorch is not None) and (poptorch.isRunningOnIpu())
-        max_num_nodes_per_graph = None
-        if on_ipu:
-            max_num_nodes_per_graph = self.max_num_nodes_per_graph
+        # Initialize the MPNN and Attention layers
+        self.mpnn = self._parse_mpnn_layer(mpnn_type, mpnn_kwargs)
+        self.attn_layer = self._parse_attn_layer(attn_type, self.biased_attention_key, attn_kwargs)
 
+    def forward(self, batch: Batch) -> Batch:
         # pe, h, edge_index, edge_attr = batch.pos_enc_feats_sign_flip, batch.h, batch.edge_index, batch.edge_attr
         h = batch.h
 
@@ -171,44 +183,64 @@ def forward(self, batch):
             h_local = h_in + h_local  # Residual connection for nodes, not used in gps++.
         if self.norm_layer_local is not None:
             h_local = self.norm_layer_local(h_local)
-        h = h_local
+
         # Multi-head attention.
-        # * batch.batch is the indicator vector for nodes of which graph it belongs to
-        # * h_dense
         if self.attn_layer is not None:
-            # Convert the tensor to a dense batch, then back to a sparse batch
-            batch_size = None if h.device.type != "ipu" else batch.graph_is_true.shape[0]
-            h_dense, mask, idx = to_dense_batch(
-                h,
-                batch=batch.batch,
-                batch_size=batch_size,
-                max_num_nodes_per_graph=max_num_nodes_per_graph,
-                drop_nodes_last_graph=on_ipu,
-            )
-            attn_bias = None
-            if self.biased_attention_key is not None:
-                attn_bias = getattr(batch, self.biased_attention_key)
-            h_attn = self._sa_block(h_dense, attn_bias)
-            h_attn = to_sparse_batch(h_attn, idx)
-
-            # Dropout, residual, norm
-            if self.dropout_attn is not None:
-                h_attn = self.dropout_attn(h_attn)
-            h_attn = h_in + h_attn
-            if self.norm_layer_attn is not None:
-                h_attn = self.norm_layer_attn(h_attn)
-
+            h_attn = self._self_attention_block(h, h_in, batch)
             # Combine local and global outputs.
-            h = h + h_attn
-
-        # Feed Forward block.
-        h = self._ff_block(h)
+            h = h_local + h_attn
+        else:
+            h = h_local
+
+        # MLP block, with skip connection
+        h = h + self.mlp(h)
+        h = self.f_out(h)
+        # Add the droppath to the output of the MLP
+        batch_size = None if h.device.type != "ipu" else batch.graph_is_true.shape[0]
+        if self.droppath_ffn is not None:
+            h = self.droppath_ffn(h, batch.batch, batch_size=batch_size)
 
         batch_out.h = h
 
         return batch_out
 
-    def _parse_attn_layer(self, attn_type, biased_attention_key, **attn_kwargs):
+    def _parse_mpnn_layer(self, mpnn_type, mpnn_kwargs: Dict[str, Any]) -> Optional[Module]:
+        """Parse the MPNN layer."""
+
+        mpnn_kwargs = deepcopy(mpnn_kwargs)
+        if mpnn_kwargs is None:
+            mpnn_kwargs = {}
+
+        # Set the default values
+        mpnn_kwargs = deepcopy(mpnn_kwargs)
+        mpnn_kwargs.setdefault("in_dim", self.in_dim)
+        mpnn_kwargs.setdefault("out_dim", self.in_dim)
+        mpnn_kwargs.setdefault("in_dim_edges", self.in_dim_edges)
+        mpnn_kwargs.setdefault("out_dim_edges", self.out_dim_edges)
+        # TODO: The rest of default values
+        self.mpnn_kwargs = mpnn_kwargs
+
+        # Initialize the MPNN layer
+        mpnn_class = PYG_LAYERS_DICT[mpnn_type]
+        mpnn_layer = mpnn_class(**mpnn_kwargs, layer_depth=self.layer_depth, layer_idx=self.layer_idx)
+
+        return mpnn_layer
+
+    def _parse_attn_layer(
+        self, attn_type, biased_attention_key: str, attn_kwargs: Dict[str, Any]
+    ) -> Optional[Module]:
+        """Parse the attention layer."""
+
+        # Set the default values for the Attention layer
+        if attn_kwargs is None:
+            attn_kwargs = {}
+        attn_kwargs.setdefault("embed_dim", self.in_dim)
+        attn_kwargs.setdefault("num_heads", 1)
+        attn_kwargs.setdefault("dropout", self.dropout)
+        attn_kwargs.setdefault("batch_first", True)
+        self.attn_kwargs = attn_kwargs
+
+        # Initialize the Attention layer
         attn_layer, attn_class = None, None
         if attn_type is not None:
             attn_class = ATTENTION_LAYERS_DICT[attn_type]
@@ -216,29 +248,53 @@ def _parse_attn_layer(self, attn_type, biased_attention_key, **attn_kwargs):
             attn_layer = attn_class(biased_attention_key, **attn_kwargs)
         return attn_layer
 
-    def _ff_block(self, h):
-        """Feed Forward block."""
-        h_in = h
-        # First linear layer + activation + dropout
-        h = self.ff_linear1(h)
-        if self.activation_layer is not None:
-            h = self.activation_layer(h)
-        if self.ff_dropout1 is not None:
-            h = self.ff_dropout1(h)
-
-        # Second linear layer + dropout
-        h = self.ff_linear2(h)
-        if self.ff_dropout2 is not None:
-            h = self.ff_dropout2(h)
-
-        # Residual
-        h = h + h_in
-
-        # Third linear layer + norm
-        h = self.ff_out(h)
-        if self.norm_layer_ff is not None:
-            h = self.norm_layer_ff(h)
-        return h
+    def _self_attention_block(self, h: Tensor, h_in: Tensor, batch: Batch) -> Tensor:
+        """
+        Applying the multi-head self-attention to the batch of graphs.
+        First the batch is converted from [num_nodes, hidden_dim] to [num_graphs, max_num_nodes, hidden_dim]
+        Then the self-attention is applied on each graph
+        Then the batch is converted again to [num_nodes, hidden_dim]
+        """
+
+        # Multi-head attention.
+        on_ipu = is_running_on_ipu()
+        max_num_nodes_per_graph = None
+        if on_ipu:
+            max_num_nodes_per_graph = self.max_num_nodes_per_graph
+
+        # Convert the tensor to a dense batch, then back to a sparse batch
+        batch_size = None if h.device.type != "ipu" else batch.graph_is_true.shape[0]
+
+        # h[num_nodes, hidden_dim] -> h_dense[num_graphs, max_num_nodes, hidden_dim]
+        h_dense, mask, idx = to_dense_batch(
+            h,
+            batch=batch.batch,  # The batch index as a vector that indicates for nodes of which graph it belongs to
+            batch_size=batch_size,
+            max_num_nodes_per_graph=max_num_nodes_per_graph,
+            drop_nodes_last_graph=on_ipu,
+        )
+
+        attn_bias = None
+        if self.biased_attention_key is not None:
+            attn_bias = batch[self.biased_attention_key]
+
+        # h_dense[num_graphs, max_num_nodes, hidden_dim] -> h_attn[num_graphs, max_num_nodes, hidden_dim]
+        h_attn = self._sa_block(h_dense, attn_bias=attn_bias, attn_mask=None, key_padding_mask=~mask)
+
+        # h_attn[num_graphs, max_num_nodes, hidden_dim] -> h_attn[num_nodes, hidden_dim]
+        h_attn = to_sparse_batch(h_attn, idx)
+
+        # Dropout, residual, norm
+        if self.dropout_attn is not None:
+            h_attn = self.dropout_attn(h_attn)
+        h_attn = h_in + h_attn
+        if self.norm_layer_attn is not None:
+            h_attn = self.norm_layer_attn(h_attn)
+        if self.droppath_layer is not None:
+            self.droppath_layer(h_attn, batch.batch, batch_size=batch_size)
+
+        # Combine local and global outputs.
+        return h + h_attn
 
     def _sa_block(self, x, attn_bias, attn_mask=None, key_padding_mask=None):
         """Self-attention block."""
diff --git a/goli/nn/pyg_layers/mpnn_pyg.py b/goli/nn/pyg_layers/mpnn_pyg.py
index 1af39b55e..494f18754 100644
--- a/goli/nn/pyg_layers/mpnn_pyg.py
+++ b/goli/nn/pyg_layers/mpnn_pyg.py
@@ -26,6 +26,7 @@ def __init__(
         out_dim_edges: Optional[int] = 32,
         aggregation_method: Optional[List[Union[str, Aggregation]]] = ["sum"],
         num_edge_mlp: Optional[int] = 2,
+        use_globals: bool = True,
         edge_dropout_rate: Optional[float] = 0.0035,
         **kwargs,
     ):
@@ -111,6 +112,7 @@ def __init__(
             activation=activation,
             dropout=dropout,
             normalization=normalization,
+            **kwargs,
         )
 
         self.gather_from = gather_from
@@ -131,9 +133,9 @@ def __init__(
         node_model_hidden_dim = 4 * self.in_dim
         self.node_model = MLP(
             in_dim=node_model_in_dim,
-            hidden_dim=node_model_hidden_dim,
+            hidden_dims=node_model_hidden_dim,
             out_dim=self.out_dim,
-            layers=self.num_node_mlp,
+            depth=self.num_node_mlp,
             activation=self.activation_layer,
             normalization=self.normalization,
         )
@@ -143,14 +145,16 @@ def __init__(
         edge_model_hidden_dim = 4 * self.in_dim_edges
         self.edge_model = MLP(
             in_dim=edge_model_in_dim,
-            hidden_dim=edge_model_hidden_dim,
+            hidden_dims=edge_model_hidden_dim,
             out_dim=self.out_dim_edges,
-            layers=self.num_edge_mlp,
+            depth=self.num_edge_mlp,
             activation=self.activation_layer,
             last_dropout=self.edge_dropout_rate,
             normalization=self.normalization,
         )
 
+        self.use_globals = use_globals
+
     def gather_features(
         self,
         input_features: Tensor,
@@ -263,7 +267,6 @@ def aggregate_features(
     def forward(self, batch: Batch) -> Batch:
         senders = batch.edge_index[0]
         receivers = batch.edge_index[1]
-
         # ---------------EDGE step---------------
         edge_model_input, sender_nodes, receiver_nodes = self.gather_features(batch.h, senders, receivers)
 
@@ -287,7 +290,9 @@ def forward(self, batch: Batch) -> Batch:
 
         # ---------------Apply norm activation and dropout---------------
         # use dropout value of the layer (default 0.3)
-        batch.h = self.apply_norm_activation_dropout(batch.h, normalization=False, activation=False)
+        batch.h = self.apply_norm_activation_dropout(
+            batch.h, normalization=False, activation=False, batch_idx=batch.batch, batch_size=batch.num_graphs
+        )
 
         return batch
 
diff --git a/goli/nn/pyg_layers/pna_pyg.py b/goli/nn/pyg_layers/pna_pyg.py
index 270a8f850..fcb8b3e93 100644
--- a/goli/nn/pyg_layers/pna_pyg.py
+++ b/goli/nn/pyg_layers/pna_pyg.py
@@ -110,6 +110,7 @@ def __init__(
             activation=activation,
             dropout=dropout,
             normalization=normalization,
+            **kwargs,
         )
 
         # Allow int32 as edge index
@@ -131,9 +132,9 @@ def __init__(
         # MLP used on each pair of nodes with their edge MLP(h_u, h_v, e_uv)
         self.pretrans = MLP(
             in_dim=2 * in_dim + in_dim_edges,
-            hidden_dim=in_dim,
+            hidden_dims=in_dim,
             out_dim=in_dim,
-            layers=pretrans_layers,
+            depth=pretrans_layers,
             activation=self.activation,
             last_activation=self.last_activation,
             dropout=dropout,
@@ -144,9 +145,9 @@ def __init__(
         # MLP used on the aggregated messages of the neighbours
         self.posttrans = MLP(
             in_dim=(len(aggregators) * len(scalers)) * self.in_dim,
-            hidden_dim=self.out_dim,
+            hidden_dims=self.out_dim,
             out_dim=self.out_dim,
-            layers=posttrans_layers,
+            depth=posttrans_layers,
             activation=self.activation,
             last_activation=self.last_activation,
             dropout=dropout,
diff --git a/goli/nn/pyg_layers/pooling_pyg.py b/goli/nn/pyg_layers/pooling_pyg.py
index 430a5c0e5..949e818c2 100644
--- a/goli/nn/pyg_layers/pooling_pyg.py
+++ b/goli/nn/pyg_layers/pooling_pyg.py
@@ -1,14 +1,15 @@
 import torch
 import torch.nn as nn
 from torch import Tensor, LongTensor
-from typing import List, Union, Callable, Tuple, Optional
+from typing import List, Union, Callable, Tuple, Optional, Dict
+from copy import deepcopy
 
 from torch_scatter import scatter
 from torch_geometric.data import Data, Batch
 
 from goli.nn.base_layers import MLP, FCLayer
 from goli.utils.tensor import ModuleListConcat, ModuleWrap
-
+from goli.nn.base_layers import MuReadoutGoli
 
 EPS = 1e-6
 
@@ -71,12 +72,22 @@ def scatter_std_pool(x: Tensor, batch: LongTensor, dim: int = 0, dim_size: Optio
 
 
 class PoolingWrapperPyg(ModuleWrap):
-    def forward(self, g, h, *args, **kwargs):
+    def __init__(self, func, feat_type, *args, **kwargs) -> None:
+        super().__init__(func, *args, **kwargs)
+        self.feat_type = feat_type
+
+    def forward(self, g, feature, *args, **kwargs):
         dim_size = g.num_graphs
-        return self.func(h, g.batch, dim_size=dim_size, *args, **kwargs, **self.kwargs)
+        if self.feat_type is "node":
+            index = g.batch
+        elif self.feat_type is "edge":
+            index = g.batch[g.edge_index][0]
+        else:
+            index = g.batch
+        return self.func(feature, index, dim_size=dim_size, *args, **kwargs, **self.kwargs)
 
 
-def parse_pooling_layer_pyg(in_dim: int, pooling: Union[str, List[str]], **kwargs):
+def parse_pooling_layer_pyg(in_dim: int, pooling: Union[str, List[str]], feat_type: str = "node", **kwargs):
     r"""
     Select the pooling layers from a list of strings, and put them
     in a Module that concatenates their outputs.
@@ -104,22 +115,22 @@ def parse_pooling_layer_pyg(in_dim: int, pooling: Union[str, List[str]], **kwarg
     out_pool_dim = 0
     if isinstance(pooling, str):
         pooling = [pooling]
-
+    assert feat_type in ["node", "edge", "global"]
     for this_pool in pooling:
         this_pool = None if this_pool is None else this_pool.lower()
         out_pool_dim += in_dim
         if this_pool == "sum":
-            pool_layer.append(PoolingWrapperPyg(scatter, dim=0, reduce="add", **kwargs))
+            pool_layer.append(PoolingWrapperPyg(scatter, dim=0, reduce="add", feat_type=feat_type, **kwargs))
         elif this_pool == "mean":
-            pool_layer.append(PoolingWrapperPyg(scatter, dim=0, reduce="mean", **kwargs))
+            pool_layer.append(PoolingWrapperPyg(scatter, dim=0, reduce="mean", feat_type=feat_type, **kwargs))
         elif this_pool == "logsum":
-            pool_layer.append(PoolingWrapperPyg(scatter_logsum_pool, dim=0, **kwargs))
+            pool_layer.append(PoolingWrapperPyg(scatter_logsum_pool, dim=0, feat_type=feat_type, **kwargs))
         elif this_pool == "max":
-            pool_layer.append(PoolingWrapperPyg(scatter, dim=0, reduce="max", **kwargs))
+            pool_layer.append(PoolingWrapperPyg(scatter, dim=0, reduce="max", feat_type=feat_type, **kwargs))
         elif this_pool == "min":
-            pool_layer.append(PoolingWrapperPyg(scatter, dim=0, reduce="min", **kwargs))
+            pool_layer.append(PoolingWrapperPyg(scatter, dim=0, reduce="min", feat_type=feat_type, **kwargs))
         elif this_pool == "std":
-            pool_layer.append(PoolingWrapperPyg(scatter_std_pool, dim=0, **kwargs))
+            pool_layer.append(PoolingWrapperPyg(scatter_std_pool, dim=0, feat_type=feat_type, **kwargs))
         elif (this_pool == "none") or (this_pool is None):
             pass
         else:
@@ -131,13 +142,18 @@ def parse_pooling_layer_pyg(in_dim: int, pooling: Union[str, List[str]], **kwarg
 class VirtualNodePyg(nn.Module):
     def __init__(
         self,
-        dim: int,
+        in_dim: int,
+        out_dim: int,
+        in_dim_edges: Optional[int],
+        out_dim_edges: Optional[int],
         vn_type: Union[type(None), str] = "sum",
         activation: Union[str, Callable] = "relu",
         dropout: float = 0.0,
         normalization: Union[str, Callable] = "none",
         bias: bool = True,
         residual: bool = True,
+        use_edges: bool = False,
+        **kwargs,
     ):
         r"""
         The VirtualNode is a layer that pool the features of the graph,
@@ -146,8 +162,17 @@ def __init__(
 
         Parameters:
 
-            dim:
-                Input and output feature dimensions of the virtual node layer
+            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:
@@ -160,7 +185,6 @@ def __init__(
                 - "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.
 
@@ -182,6 +206,10 @@ def __init__(
                 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
+
         """
         super().__init__()
         if (vn_type is None) or (vn_type.lower() == "none"):
@@ -189,20 +217,43 @@ def __init__(
             self.fc_layer = None
             self.residual = None
             return
-
+        # Layer sizes stored here
+        self.in_dim_edges = in_dim_edges
+        self.out_dim_edges = out_dim_edges
+        self.in_dim_nodes = in_dim
+        self.out_dim_nodes = out_dim
+
+        # Use edge features in pooling
+        self.use_edges = use_edges
         self.vn_type = vn_type.lower()
-        self.layer, out_pool_dim = parse_pooling_layer_pyg(in_dim=dim, pooling=self.vn_type)
+        self.layer, out_pool_dim = parse_pooling_layer_pyg(
+            in_dim=self.in_dim_nodes, pooling=self.vn_type, feat_type="node"
+        )
         self.residual = residual
+
+        if self.use_edges:
+            self.edge_layer, out_edge_pool_dim = parse_pooling_layer_pyg(
+                in_dim=self.in_dim_edges, pooling=self.vn_type, feat_type="edge"
+            )
+
+            out_pool_dim = out_pool_dim + out_edge_pool_dim
+
         self.fc_layer = FCLayer(
             in_dim=out_pool_dim,
-            out_dim=dim,
+            out_dim=out_pool_dim,
             activation=activation,
             dropout=dropout,
             normalization=normalization,
             bias=bias,
         )
 
-    def forward(self, g: Union[Data, Batch], h: Tensor, vn_h: LongTensor) -> Tuple[Tensor, Tensor]:
+        # Projection layers from the pooling layer to node and edge feature sizes
+        self.node_projection = MuReadoutGoli(out_pool_dim, self.out_dim_nodes)
+        self.edge_projection = MuReadoutGoli(out_pool_dim, self.out_dim_edges)
+
+    def forward(
+        self, g: Union[Data, Batch], h: Tensor, vn_h: LongTensor, e: Tensor
+    ) -> Tuple[Tensor, Tensor, Tensor]:
         r"""
         Apply the virtual node layer.
 
@@ -220,7 +271,9 @@ def forward(self, g: Union[Data, Batch], h: Tensor, vn_h: LongTensor) -> Tuple[T
                 `M` is the number of graphs, `Din` is the input features.
                 It is added to the result after the MLP, as a residual connection
 
-            batch
+            e (torch.Tensor[..., E, Din]):
+                Edge feature tensor, before convolution.
+                `E` is the number of edges, `Din` is the input features
 
         Returns:
 
@@ -232,22 +285,31 @@ def forward(self, g: Union[Data, Batch], h: Tensor, vn_h: LongTensor) -> Tuple[T
                 Graph feature tensor to be used at the next virtual node, or `None`
                 `M` is the number of graphs, `Dout` is the output features
 
+            `e = 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
+
         """
 
         # Pool the features
         if self.vn_type is None:
-            return h, vn_h
-        else:
-            pool = self.layer(g, h)
+            return h, vn_h, e
+        pool = self.layer(g, h)
+        if self.use_edges:
+            edge_pool = self.edge_layer(g, e)
+            pool = torch.cat((pool, edge_pool), -1)
 
-        # Compute the new virtual node features
         vn_h_temp = self.fc_layer.forward(vn_h + pool)
+
         if self.residual:
             vn_h = vn_h + vn_h_temp
         else:
             vn_h = vn_h_temp
 
-        # Add the virtual node value to the graph features
-        h = h + vn_h[g.batch]
+        # Add the virtual node projections to the node features
+        h = h + self.node_projection(vn_h[g.batch])
 
-        return h, vn_h
+        if self.use_edges:
+            # Add the virtual node projections to the edge features
+            e = e + self.edge_projection(vn_h[g.batch[g.edge_index][0]])
+        return h, vn_h, e
diff --git a/goli/nn/pyg_layers/utils.py b/goli/nn/pyg_layers/utils.py
index 5793d91f6..121441c9f 100644
--- a/goli/nn/pyg_layers/utils.py
+++ b/goli/nn/pyg_layers/utils.py
@@ -45,9 +45,9 @@ def __init__(
         self.gaussian = GaussianLayer(self.num_kernel, in_dim=in_dim)
         self.gaussian_proj = MLP(
             in_dim=self.num_kernel,
-            hidden_dim=self.num_kernel,
+            hidden_dims=self.num_kernel,
             out_dim=self.num_heads,
-            layers=num_layers,
+            depth=num_layers,
             activation=activation,
             last_layer_is_readout=True,  # Since the output is not proportional to the hidden dim, but to the number of heads
         )
diff --git a/mup_base_params.yaml b/mup_base_params.yaml
index bfe20b441..bfc4dd213 100644
--- a/mup_base_params.yaml
+++ b/mup_base_params.yaml
@@ -50,53 +50,73 @@ gnn.layers.0.attn_layer.out_proj.bias:
 gnn.layers.0.attn_layer.out_proj.weight:
 - 32
 - 32
-gnn.layers.0.ff_linear1.linear.bias:
-- 64
-gnn.layers.0.ff_linear1.linear.weight:
-- 64
+gnn.layers.0.f_out.linear.bias:
 - 32
-gnn.layers.0.ff_linear2.linear.bias:
+gnn.layers.0.f_out.linear.weight:
 - 32
-gnn.layers.0.ff_linear2.linear.weight:
 - 32
-- 64
-gnn.layers.0.ff_out.linear.bias:
+gnn.layers.0.f_out.normalization.bias:
+- 32
+gnn.layers.0.f_out.normalization.weight:
 - 32
-gnn.layers.0.ff_out.linear.weight:
+gnn.layers.0.mlp.fully_connected.0.linear.bias:
+- 128
+gnn.layers.0.mlp.fully_connected.0.linear.weight:
+- 128
 - 32
+gnn.layers.0.mlp.fully_connected.1.linear.bias:
 - 32
+gnn.layers.0.mlp.fully_connected.1.linear.weight:
+- 32
+- 128
 gnn.layers.0.mpnn.edge_model.fully_connected.0.linear.bias:
-- null
+- 64
 gnn.layers.0.mpnn.edge_model.fully_connected.0.linear.weight:
-- null
-- null
+- 64
+- 80
 gnn.layers.0.mpnn.edge_model.fully_connected.0.normalization.bias:
-- null
+- 64
 gnn.layers.0.mpnn.edge_model.fully_connected.0.normalization.weight:
-- null
+- 64
 gnn.layers.0.mpnn.edge_model.fully_connected.1.linear.bias:
-- null
+- 16
 gnn.layers.0.mpnn.edge_model.fully_connected.1.linear.weight:
-- null
-- null
+- 16
+- 64
 gnn.layers.0.mpnn.node_model.fully_connected.0.linear.bias:
-- null
+- 128
 gnn.layers.0.mpnn.node_model.fully_connected.0.linear.weight:
-- null
-- null
+- 128
+- 128
 gnn.layers.0.mpnn.node_model.fully_connected.0.normalization.bias:
-- null
+- 128
 gnn.layers.0.mpnn.node_model.fully_connected.0.normalization.weight:
-- null
+- 128
 gnn.layers.0.mpnn.node_model.fully_connected.1.linear.bias:
-- null
+- 32
 gnn.layers.0.mpnn.node_model.fully_connected.1.linear.weight:
-- null
-- null
+- 32
+- 128
 gnn.layers.0.mpnn.norm_layer.bias:
-- null
+- 32
 gnn.layers.0.mpnn.norm_layer.weight:
-- null
+- 32
+gnn.layers.0.norm_layer.bias:
+- 32
+gnn.layers.0.norm_layer.weight:
+- 32
+gnn.layers.0.norm_layer_attn.bias:
+- 32
+gnn.layers.0.norm_layer_attn.weight:
+- 32
+gnn.layers.0.norm_layer_ff.bias:
+- 32
+gnn.layers.0.norm_layer_ff.weight:
+- 32
+gnn.layers.0.norm_layer_local.bias:
+- 32
+gnn.layers.0.norm_layer_local.weight:
+- 32
 gnn.layers.1.attn_layer.in_proj_bias:
 - 96
 gnn.layers.1.attn_layer.in_proj_weight:
@@ -107,53 +127,73 @@ gnn.layers.1.attn_layer.out_proj.bias:
 gnn.layers.1.attn_layer.out_proj.weight:
 - 32
 - 32
-gnn.layers.1.ff_linear1.linear.bias:
-- 64
-gnn.layers.1.ff_linear1.linear.weight:
-- 64
+gnn.layers.1.f_out.linear.bias:
 - 32
-gnn.layers.1.ff_linear2.linear.bias:
+gnn.layers.1.f_out.linear.weight:
 - 32
-gnn.layers.1.ff_linear2.linear.weight:
 - 32
-- 64
-gnn.layers.1.ff_out.linear.bias:
+gnn.layers.1.f_out.normalization.bias:
 - 32
-gnn.layers.1.ff_out.linear.weight:
+gnn.layers.1.f_out.normalization.weight:
 - 32
+gnn.layers.1.mlp.fully_connected.0.linear.bias:
+- 128
+gnn.layers.1.mlp.fully_connected.0.linear.weight:
+- 128
 - 32
+gnn.layers.1.mlp.fully_connected.1.linear.bias:
+- 32
+gnn.layers.1.mlp.fully_connected.1.linear.weight:
+- 32
+- 128
 gnn.layers.1.mpnn.edge_model.fully_connected.0.linear.bias:
-- null
+- 64
 gnn.layers.1.mpnn.edge_model.fully_connected.0.linear.weight:
-- null
-- null
+- 64
+- 80
 gnn.layers.1.mpnn.edge_model.fully_connected.0.normalization.bias:
-- null
+- 64
 gnn.layers.1.mpnn.edge_model.fully_connected.0.normalization.weight:
-- null
+- 64
 gnn.layers.1.mpnn.edge_model.fully_connected.1.linear.bias:
-- null
+- 16
 gnn.layers.1.mpnn.edge_model.fully_connected.1.linear.weight:
-- null
-- null
+- 16
+- 64
 gnn.layers.1.mpnn.node_model.fully_connected.0.linear.bias:
-- null
+- 128
 gnn.layers.1.mpnn.node_model.fully_connected.0.linear.weight:
-- null
-- null
+- 128
+- 128
 gnn.layers.1.mpnn.node_model.fully_connected.0.normalization.bias:
-- null
+- 128
 gnn.layers.1.mpnn.node_model.fully_connected.0.normalization.weight:
-- null
+- 128
 gnn.layers.1.mpnn.node_model.fully_connected.1.linear.bias:
-- null
+- 32
 gnn.layers.1.mpnn.node_model.fully_connected.1.linear.weight:
-- null
-- null
+- 32
+- 128
 gnn.layers.1.mpnn.norm_layer.bias:
-- null
+- 32
 gnn.layers.1.mpnn.norm_layer.weight:
-- null
+- 32
+gnn.layers.1.norm_layer.bias:
+- 32
+gnn.layers.1.norm_layer.weight:
+- 32
+gnn.layers.1.norm_layer_attn.bias:
+- 32
+gnn.layers.1.norm_layer_attn.weight:
+- 32
+gnn.layers.1.norm_layer_ff.bias:
+- 32
+gnn.layers.1.norm_layer_ff.weight:
+- 32
+gnn.layers.1.norm_layer_local.bias:
+- 32
+gnn.layers.1.norm_layer_local.weight:
+- 32
 gnn.layers.2.attn_layer.in_proj_bias:
 - 96
 gnn.layers.2.attn_layer.in_proj_weight:
@@ -164,53 +204,73 @@ gnn.layers.2.attn_layer.out_proj.bias:
 gnn.layers.2.attn_layer.out_proj.weight:
 - 32
 - 32
-gnn.layers.2.ff_linear1.linear.bias:
-- 64
-gnn.layers.2.ff_linear1.linear.weight:
-- 64
+gnn.layers.2.f_out.linear.bias:
 - 32
-gnn.layers.2.ff_linear2.linear.bias:
+gnn.layers.2.f_out.linear.weight:
 - 32
-gnn.layers.2.ff_linear2.linear.weight:
 - 32
-- 64
-gnn.layers.2.ff_out.linear.bias:
+gnn.layers.2.f_out.normalization.bias:
 - 32
-gnn.layers.2.ff_out.linear.weight:
+gnn.layers.2.f_out.normalization.weight:
 - 32
+gnn.layers.2.mlp.fully_connected.0.linear.bias:
+- 128
+gnn.layers.2.mlp.fully_connected.0.linear.weight:
+- 128
 - 32
+gnn.layers.2.mlp.fully_connected.1.linear.bias:
+- 32
+gnn.layers.2.mlp.fully_connected.1.linear.weight:
+- 32
+- 128
 gnn.layers.2.mpnn.edge_model.fully_connected.0.linear.bias:
-- null
+- 64
 gnn.layers.2.mpnn.edge_model.fully_connected.0.linear.weight:
-- null
-- null
+- 64
+- 80
 gnn.layers.2.mpnn.edge_model.fully_connected.0.normalization.bias:
-- null
+- 64
 gnn.layers.2.mpnn.edge_model.fully_connected.0.normalization.weight:
-- null
+- 64
 gnn.layers.2.mpnn.edge_model.fully_connected.1.linear.bias:
-- null
+- 16
 gnn.layers.2.mpnn.edge_model.fully_connected.1.linear.weight:
-- null
-- null
+- 16
+- 64
 gnn.layers.2.mpnn.node_model.fully_connected.0.linear.bias:
-- null
+- 128
 gnn.layers.2.mpnn.node_model.fully_connected.0.linear.weight:
-- null
-- null
+- 128
+- 128
 gnn.layers.2.mpnn.node_model.fully_connected.0.normalization.bias:
-- null
+- 128
 gnn.layers.2.mpnn.node_model.fully_connected.0.normalization.weight:
-- null
+- 128
 gnn.layers.2.mpnn.node_model.fully_connected.1.linear.bias:
-- null
+- 32
 gnn.layers.2.mpnn.node_model.fully_connected.1.linear.weight:
-- null
-- null
+- 32
+- 128
 gnn.layers.2.mpnn.norm_layer.bias:
-- null
+- 32
 gnn.layers.2.mpnn.norm_layer.weight:
-- null
+- 32
+gnn.layers.2.norm_layer.bias:
+- 32
+gnn.layers.2.norm_layer.weight:
+- 32
+gnn.layers.2.norm_layer_attn.bias:
+- 32
+gnn.layers.2.norm_layer_attn.weight:
+- 32
+gnn.layers.2.norm_layer_ff.bias:
+- 32
+gnn.layers.2.norm_layer_ff.weight:
+- 32
+gnn.layers.2.norm_layer_local.bias:
+- 32
+gnn.layers.2.norm_layer_local.weight:
+- 32
 gnn.layers.3.attn_layer.in_proj_bias:
 - 96
 gnn.layers.3.attn_layer.in_proj_weight:
@@ -221,53 +281,73 @@ gnn.layers.3.attn_layer.out_proj.bias:
 gnn.layers.3.attn_layer.out_proj.weight:
 - 32
 - 32
-gnn.layers.3.ff_linear1.linear.bias:
-- 64
-gnn.layers.3.ff_linear1.linear.weight:
-- 64
+gnn.layers.3.f_out.linear.bias:
 - 32
-gnn.layers.3.ff_linear2.linear.bias:
+gnn.layers.3.f_out.linear.weight:
 - 32
-gnn.layers.3.ff_linear2.linear.weight:
 - 32
-- 64
-gnn.layers.3.ff_out.linear.bias:
+gnn.layers.3.f_out.normalization.bias:
+- 32
+gnn.layers.3.f_out.normalization.weight:
 - 32
-gnn.layers.3.ff_out.linear.weight:
+gnn.layers.3.mlp.fully_connected.0.linear.bias:
+- 128
+gnn.layers.3.mlp.fully_connected.0.linear.weight:
+- 128
 - 32
+gnn.layers.3.mlp.fully_connected.1.linear.bias:
 - 32
+gnn.layers.3.mlp.fully_connected.1.linear.weight:
+- 32
+- 128
 gnn.layers.3.mpnn.edge_model.fully_connected.0.linear.bias:
-- null
+- 64
 gnn.layers.3.mpnn.edge_model.fully_connected.0.linear.weight:
-- null
-- null
+- 64
+- 80
 gnn.layers.3.mpnn.edge_model.fully_connected.0.normalization.bias:
-- null
+- 64
 gnn.layers.3.mpnn.edge_model.fully_connected.0.normalization.weight:
-- null
+- 64
 gnn.layers.3.mpnn.edge_model.fully_connected.1.linear.bias:
-- null
+- 16
 gnn.layers.3.mpnn.edge_model.fully_connected.1.linear.weight:
-- null
-- null
+- 16
+- 64
 gnn.layers.3.mpnn.node_model.fully_connected.0.linear.bias:
-- null
+- 128
 gnn.layers.3.mpnn.node_model.fully_connected.0.linear.weight:
-- null
-- null
+- 128
+- 128
 gnn.layers.3.mpnn.node_model.fully_connected.0.normalization.bias:
-- null
+- 128
 gnn.layers.3.mpnn.node_model.fully_connected.0.normalization.weight:
-- null
+- 128
 gnn.layers.3.mpnn.node_model.fully_connected.1.linear.bias:
-- null
+- 32
 gnn.layers.3.mpnn.node_model.fully_connected.1.linear.weight:
-- null
-- null
+- 32
+- 128
 gnn.layers.3.mpnn.norm_layer.bias:
-- null
+- 32
 gnn.layers.3.mpnn.norm_layer.weight:
-- null
+- 32
+gnn.layers.3.norm_layer.bias:
+- 32
+gnn.layers.3.norm_layer.weight:
+- 32
+gnn.layers.3.norm_layer_attn.bias:
+- 32
+gnn.layers.3.norm_layer_attn.weight:
+- 32
+gnn.layers.3.norm_layer_ff.bias:
+- 32
+gnn.layers.3.norm_layer_ff.weight:
+- 32
+gnn.layers.3.norm_layer_local.bias:
+- 32
+gnn.layers.3.norm_layer_local.weight:
+- 32
 gnn.out_linear.linear.bias:
 - 32
 gnn.out_linear.linear.weight:
@@ -277,6 +357,63 @@ gnn.out_linear.normalization.bias:
 - 32
 gnn.out_linear.normalization.weight:
 - 32
+gnn.virtual_node_layers.0.edge_projection.bias:
+- 16
+gnn.virtual_node_layers.0.edge_projection.weight:
+- 16
+- 32
+gnn.virtual_node_layers.0.fc_layer.linear.bias:
+- 32
+gnn.virtual_node_layers.0.fc_layer.linear.weight:
+- 32
+- 32
+gnn.virtual_node_layers.0.fc_layer.normalization.bias:
+- 32
+gnn.virtual_node_layers.0.fc_layer.normalization.weight:
+- 32
+gnn.virtual_node_layers.0.node_projection.bias:
+- 32
+gnn.virtual_node_layers.0.node_projection.weight:
+- 32
+- 32
+gnn.virtual_node_layers.1.edge_projection.bias:
+- 16
+gnn.virtual_node_layers.1.edge_projection.weight:
+- 16
+- 32
+gnn.virtual_node_layers.1.fc_layer.linear.bias:
+- 32
+gnn.virtual_node_layers.1.fc_layer.linear.weight:
+- 32
+- 32
+gnn.virtual_node_layers.1.fc_layer.normalization.bias:
+- 32
+gnn.virtual_node_layers.1.fc_layer.normalization.weight:
+- 32
+gnn.virtual_node_layers.1.node_projection.bias:
+- 32
+gnn.virtual_node_layers.1.node_projection.weight:
+- 32
+- 32
+gnn.virtual_node_layers.2.edge_projection.bias:
+- 16
+gnn.virtual_node_layers.2.edge_projection.weight:
+- 16
+- 32
+gnn.virtual_node_layers.2.fc_layer.linear.bias:
+- 32
+gnn.virtual_node_layers.2.fc_layer.linear.weight:
+- 32
+- 32
+gnn.virtual_node_layers.2.fc_layer.normalization.bias:
+- 32
+gnn.virtual_node_layers.2.fc_layer.normalization.weight:
+- 32
+gnn.virtual_node_layers.2.node_projection.bias:
+- 32
+gnn.virtual_node_layers.2.node_projection.weight:
+- 32
+- 32
 post_nn.layers.0.linear.bias:
 - 200
 post_nn.layers.0.linear.weight:
@@ -319,7 +456,7 @@ task_heads.task_heads.homo.layers.0.normalization.bias:
 task_heads.task_heads.homo.layers.0.normalization.weight:
 - 16
 task_heads.task_heads.homo.layers.1.linear.bias:
-- 0
+- null
 task_heads.task_heads.homo.layers.1.linear.weight:
-- 0
+- null
 - 16
diff --git a/news/ipu_fp16.rst b/news/ipu_fp16.rst
new file mode 100644
index 000000000..5f7a85d77
--- /dev/null
+++ b/news/ipu_fp16.rst
@@ -0,0 +1,3 @@
+**Added:**
+
+* better support for training models in FP16 on the IPU 
diff --git a/news/simplified_gps.rst b/news/simplified_gps.rst
new file mode 100644
index 000000000..b6d0717c2
--- /dev/null
+++ b/news/simplified_gps.rst
@@ -0,0 +1,8 @@
+**Added:**
+
+* Methods to pre-process the `attn_kwargs` and `mpnn_kwargs` in `GPSLayerPyg`
+
+**Changed:**
+
+* Using an `MLP` object instead of the `_ff_block` to simplify the code
+
diff --git a/news/stochastic_depth.rst b/news/stochastic_depth.rst
new file mode 100644
index 000000000..0d2bbeb91
--- /dev/null
+++ b/news/stochastic_depth.rst
@@ -0,0 +1,6 @@
+**Added:**
+* stochastic depth
+* unit test for stochastic depth
+
+**Changed:**
+* GPS layer with stochastic depth for attention and ffn module


Test patch:
diff --git a/tests/test_base_layers.py b/tests/test_base_layers.py
new file mode 100644
index 000000000..b5fee49f8
--- /dev/null
+++ b/tests/test_base_layers.py
@@ -0,0 +1,53 @@
+"""
+Unit tests for the different layers of goli/nn/base_layers
+"""
+
+import torch
+import unittest as ut
+from torch_geometric.data import Data, Batch
+from copy import deepcopy
+
+from goli.nn.base_layers import DropPath
+
+
+class test_Base_Layers(ut.TestCase):
+    in_dim = 21
+    out_dim = 11
+    in_dim_edges = 13
+    out_dim_edges = 17
+
+    edge_idx1 = torch.stack([torch.tensor([0, 1, 2, 3, 2]), torch.tensor([1, 2, 3, 0, 0])])
+    edge_idx2 = torch.stack([torch.tensor([0, 0, 0, 1]), torch.tensor([0, 1, 2, 0])])
+    x1 = torch.randn(edge_idx1.max() + 1, in_dim, dtype=torch.float32)
+    e1 = torch.randn(edge_idx1.shape[-1], in_dim_edges, dtype=torch.float32)
+    x2 = torch.randn(edge_idx2.max() + 1, in_dim, dtype=torch.float32)
+    e2 = torch.randn(edge_idx2.shape[-1], in_dim_edges, dtype=torch.float32)
+
+    g1 = Data(h=x1, edge_index=edge_idx1, edge_attr=e1)
+    g2 = Data(h=x2, edge_index=edge_idx2, edge_attr=e2)
+    bg = Batch.from_data_list([g1, g2])
+
+    # for drop_rate=0.5, test if the output shape is correct
+    def test_droppath_layer_0p5(self):
+        bg = deepcopy(self.bg)
+        h_in = bg.h
+        layer = DropPath(drop_rate=0.5)
+        h_out = layer.forward(h_in, bg.batch)
+        self.assertEqual(h_out.shape, h_in.shape)
+
+    # for drop_rate=1.0, test if the output are all zeros
+    def test_droppath_layer_1p0(self):
+        bg = deepcopy(self.bg)
+        h_in = bg.h
+        zero_tesor = torch.zeros(h_in.shape)
+        layer = DropPath(drop_rate=1.0)
+        h_out = layer.forward(h_in, bg.batch)
+        self.assertTrue(torch.allclose(zero_tesor, h_out.detach()))
+
+    # for drop_rate=0.0, test if the output matches the original output
+    def test_droppath_layer_0p0(self):
+        bg = deepcopy(self.bg)
+        h_in = bg.h
+        layer = DropPath(drop_rate=0.0)
+        h_out = layer.forward(h_in, bg.batch)
+        self.assertTrue(torch.allclose(h_in.detach(), h_out.detach()))
diff --git a/tests/test_pyg_layers.py b/tests/test_pyg_layers.py
index cace23979..670155cf1 100644
--- a/tests/test_pyg_layers.py
+++ b/tests/test_pyg_layers.py
@@ -4,14 +4,12 @@
 The layers are not thoroughly tested due to the difficulty of testing them
 """
 
-from ast import Assert
 import numpy as np
 import torch
 import unittest as ut
 from torch_geometric.data import Data, Batch
 from copy import deepcopy
 
-from goli.ipu.to_dense_batch import to_dense_batch
 from goli.nn.pyg_layers import (
     GINConvPyg,
     GINEConvPyg,
@@ -19,6 +17,7 @@
     GatedGCNPyg,
     PNAMessagePassingPyg,
     GPSLayerPyg,
+    VirtualNodePyg,
 )
 
 from goli.nn.pyg_layers.utils import (
@@ -49,12 +48,21 @@ class test_Pyg_Layers(ut.TestCase):
         "activation": "relu",
         "dropout": 0.1,
         "normalization": "batch_norm",
+        "droppath_rate": 0.1,
+        "layer_idx": 1,
+        "layer_depth": 10,
     }
 
     def test_gpslayer(self):
         bg = deepcopy(self.bg)
         h_in = bg.h
-        layer = GPSLayerPyg(in_dim=self.in_dim, out_dim=self.out_dim, **self.kwargs)
+        kwargs = deepcopy(self.kwargs)
+        kwargs.pop("droppath_rate")
+        kwargs["droppath_rate_attn"] = 0.2
+        kwargs["droppath_rate_ffn"] = 0.3
+        kwargs["mpnn_kwargs"] = {"droppath_rate_ffn": 0.4}
+
+        layer = GPSLayerPyg(in_dim=self.in_dim, out_dim=self.out_dim, **kwargs)
 
         # Check the re-implementation of abstract methods
         self.assertTrue(layer.layer_supports_edges)
@@ -237,6 +245,44 @@ def test_gaussianlayer(self):
         tensor_with_kernel = layer.forward(input)
         self.assertEqual(tensor_with_kernel.size(), torch.Size([2, 4, 4, 3]))
 
+    def test_pooling_virtual_node(self):
+        bg = deepcopy(self.bg)
+        h_in = bg.h
+        e_in = bg.edge_attr
+        vn_h = 0
+
+        vn_types = ["sum", "mean"]
+        for vn_type, use_edges, expected_v_node in zip(vn_types, [False, True], [(2, 21), (2, 34)]):
+            with self.subTest(
+                vn_type=vn_type,
+                use_edges=use_edges,
+                expected_v_node=expected_v_node,
+            ):
+                vn_h = 0.0
+                print(vn_h, self.out_dim, h_in.size())
+                layer = VirtualNodePyg(
+                    in_dim=self.in_dim,
+                    out_dim=self.in_dim,
+                    in_dim_edges=self.in_dim_edges,
+                    out_dim_edges=self.in_dim_edges,
+                    vn_type=vn_type,
+                    use_edges=use_edges,
+                    **self.kwargs,
+                )
+
+                h_out, vn_out, e_out = layer.forward(bg, h_in, vn_h, e=e_in)
+                assert vn_out.shape == expected_v_node
+                assert h_out.shape == (7, 21)
+                assert e_out.shape == (9, 13)
+                if use_edges is False:
+                    # i.e. that the node features have been updated
+                    assert torch.equal(h_out, h_in) == False
+                    # And that the edge features have not
+                    assert torch.equal(e_out, e_in) == True
+                else:
+                    assert torch.equal(h_out, h_in) == False
+                    assert torch.equal(e_out, e_in) == False
+
 
 if __name__ == "__main__":
     ut.main()