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

Issue:
Creating examples for the GNN layers and architectures
Wrote 3 examples in jupyter notebook and added them to the documentation.

- **Using GNN layers:** Examples with GCN, GIN, GAT, PNA
- **Implementing GNN layers:** How to write our own GNN
- **Making GNN networks:** Using the architecture classes to easily implement a full neural network with GNN layers.

Fixed some minor bugs in the models that were detected while writing the examples, and added the appropriate checks in the unit-tests.


Discussion / hints:

Fix patch:
diff --git a/docs/examples/basics/implementing_gnn_layers.ipynb b/docs/examples/basics/implementing_gnn_layers.ipynb
new file mode 100644
index 000000000..d7916fb46
--- /dev/null
+++ b/docs/examples/basics/implementing_gnn_layers.ipynb
@@ -0,0 +1,325 @@
+{
+ "metadata": {
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.7.9-final"
+  },
+  "orig_nbformat": 2,
+  "kernelspec": {
+   "name": "python3",
+   "display_name": "Python 3.7.9 64-bit ('goli_env': conda)",
+   "metadata": {
+    "interpreter": {
+     "hash": "ce752a05bc046fbf91e84c99a25ae7ca133d76ead2eb04f6c646180f827f8f44"
+    }
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2,
+ "cells": [
+  {
+   "source": [
+    "# Creating GNN layers\n",
+    "\n",
+    "In this example, you will learn how to create a custom GNN class such that it is compatible with the flexible network architecture `goli.dgl.architectures.FeedForwardDGL`. We will first start by a simple layer that does not use edges, to a more complex layer that uses edges.\n",
+    "\n",
+    "Since these examples are built on top of DGL, we recommend looking at their [library](https://docs.dgl.ai/en/0.5.x/index.html) for more info. "
+   ],
+   "cell_type": "markdown",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stderr",
+     "text": [
+      "Using backend: pytorch\n"
+     ]
+    }
+   ],
+   "source": [
+    "%load_ext autoreload\n",
+    "%autoreload 2\n",
+    "\n",
+    "import torch\n",
+    "import dgl\n",
+    "from copy import deepcopy\n",
+    "\n",
+    "from goli.dgl.dgl_layers import BaseDGLLayer\n",
+    "from goli.dgl.base_layers import FCLayer\n",
+    "from goli.commons.decorators import classproperty\n",
+    "\n",
+    "\n",
+    "_ = torch.manual_seed(42)"
+   ]
+  },
+  {
+   "source": [
+    "## Pre-defining test variables\n",
+    "\n",
+    "We define below a small batched graph on which we can test the created layers"
+   ],
+   "cell_type": "markdown",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "Graph(num_nodes=7, num_edges=14,\n      ndata_schemes={'h': Scheme(shape=(5,), dtype=torch.float64)}\n      edata_schemes={'e': Scheme(shape=(13,), dtype=torch.float64)})\n"
+     ]
+    }
+   ],
+   "source": [
+    "in_dim = 5          # Input node-feature dimensions\n",
+    "out_dim = 11        # Desired output node-feature dimensions\n",
+    "in_dim_edges = 13   # Input edge-feature dimensions\n",
+    "\n",
+    "# Let's create 2 simple graphs. Here the tensors represent the connectivity between nodes\n",
+    "g1 = dgl.graph((torch.tensor([0, 1, 2]), torch.tensor([1, 2, 3])))\n",
+    "g2 = dgl.graph((torch.tensor([0, 0, 0, 1]), torch.tensor([0, 1, 2, 0])))\n",
+    "\n",
+    "# We add some node features to the graphs\n",
+    "g1.ndata[\"h\"] = torch.rand(g1.num_nodes(), in_dim, dtype=float)\n",
+    "g2.ndata[\"h\"] = torch.rand(g2.num_nodes(), in_dim, dtype=float)\n",
+    "\n",
+    "# We also add some edge features to the graphs\n",
+    "g1.edata[\"e\"] = torch.rand(g1.num_edges(), in_dim_edges, dtype=float)\n",
+    "g2.edata[\"e\"] = torch.rand(g2.num_edges(), in_dim_edges, dtype=float)\n",
+    "\n",
+    "# Finally we batch the graphs in a way compatible with the DGL library\n",
+    "bg = dgl.batch([g1, g2])\n",
+    "bg = dgl.add_self_loop(bg)\n",
+    "\n",
+    "# The batched graph will show as a single graph with 7 nodes\n",
+    "print(bg)"
+   ]
+  },
+  {
+   "source": [
+    "## Creating a simple layer\n",
+    "\n",
+    "Here, we will show how to create a GNN layer that does a mean aggregation on the neighbouring features.\n",
+    "\n",
+    "First, for the layer to be fully compatible with the flexible architecture provided by `goli.dgl.architectures.FeedForwardDGL`, it needs to inhering from the class `goli.dgl.dgl_layers.BaseDGLLayer`. This base-layer has multiple virtual methods that must be implemented in any class that inherits from it.\n",
+    "\n",
+    "The virtual methods are below\n",
+    "\n",
+    "- `layer_supports_edges`: We want to return `False` since our layer doesn't support edges\n",
+    "- `layer_inputs_edges`: We want to return `False` since our layer doesn't input edges\n",
+    "- `layer_outputs_edges`: We want to return `False` since our layer doesn't output edges\n",
+    "- `layer_outdim_factor`: We want to return `1` since the output dimension does not depend on internal parameters.\n",
+    "\n",
+    "The example is given below"
+   ],
+   "cell_type": "markdown",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class SimpleMeanLayer(BaseDGLLayer):\n",
+    "    def __init__(self, in_dim, out_dim, activation, dropout, batch_norm):\n",
+    "        # Initialize the parent class\n",
+    "        super().__init__(   in_dim=in_dim, out_dim=out_dim, activation=activation,\n",
+    "                            dropout=dropout, batch_norm=batch_norm)\n",
+    "\n",
+    "        # Create the layer with learned parameters\n",
+    "        self.layer = FCLayer(in_dim=in_dim, out_dim=out_dim)\n",
+    "\n",
+    "    def forward(self, g, h):\n",
+    "        # We first apply the mean aggregation\n",
+    "        g.ndata[\"h\"] = h\n",
+    "        g.update_all(message_func=dgl.function.copy_u(\"h\", \"m\"), \n",
+    "                    reduce_func=dgl.function.mean(\"m\", \"h\"))\n",
+    "\n",
+    "        # Then we apply the FCLayer, and the non-linearities\n",
+    "        h = g.ndata[\"h\"]\n",
+    "        h = self.layer(h)\n",
+    "        h = self.apply_norm_activation_dropout(h)\n",
+    "        return h\n",
+    "\n",
+    "    # Finally, we define all the virtual properties according to how\n",
+    "    # the class works\n",
+    "    @classproperty\n",
+    "    def layer_supports_edges(cls):\n",
+    "        return False\n",
+    "\n",
+    "    @property\n",
+    "    def layer_inputs_edges(self):\n",
+    "        return False\n",
+    "\n",
+    "    @property\n",
+    "    def layer_outputs_edges(self):\n",
+    "        return False\n",
+    "\n",
+    "    @property\n",
+    "    def out_dim_factor(self):\n",
+    "        return 1   "
+   ]
+  },
+  {
+   "source": [
+    "Now, we are ready to test the `SimpleMeanLayer` on some DGL graphs. Note that in this example, we **ignore** the edge features since they are not supported."
+   ],
+   "cell_type": "markdown",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "torch.Size([7, 5])\ntorch.Size([7, 11])\n"
+     ]
+    }
+   ],
+   "source": [
+    "graph = deepcopy(bg)\n",
+    "h_in = graph.ndata[\"h\"]\n",
+    "layer = SimpleMeanLayer(\n",
+    "            in_dim=in_dim, out_dim=out_dim, \n",
+    "            activation=\"relu\", dropout=.3, batch_norm=True).to(float)\n",
+    "h_out = layer(graph, h_in)\n",
+    "\n",
+    "print(h_in.shape)\n",
+    "print(h_out.shape)"
+   ]
+  },
+  {
+   "source": [
+    "## Creating a complex layer with edges\n",
+    "\n",
+    "Here, we will show how to create a GNN layer that does a mean aggregation on the neighbouring features, concatenated to the edge features with their neighbours. In that case, only the node features will change, and the network will not update the edge features.\n",
+    "\n",
+    "The virtual methods will have different outputs\n",
+    "\n",
+    "- `layer_supports_edges`: We want to return `True` since our layer does support edges\n",
+    "- `layer_inputs_edges`: We want to return `True` since our layer does input edges\n",
+    "- `layer_outputs_edges`: We want to return `False` since our layer will not output new edges\n",
+    "- `layer_outdim_factor`: We want to return `1` since the output dimension does not depend on internal parameters.\n",
+    "\n",
+    "The example is given below"
+   ],
+   "cell_type": "markdown",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class ComplexMeanLayer(BaseDGLLayer):\n",
+    "    def __init__(self, in_dim, out_dim, in_dim_edges, activation, dropout, batch_norm):\n",
+    "        # Initialize the parent class\n",
+    "        super().__init__(   in_dim=in_dim, out_dim=out_dim, activation=activation,\n",
+    "                            dropout=dropout, batch_norm=batch_norm)\n",
+    "\n",
+    "        # Create the layer with learned parameters. Note the addition\n",
+    "        self.layer = FCLayer(in_dim=in_dim + in_dim_edges, out_dim=out_dim)\n",
+    "\n",
+    "    def cat_nodes_edges(self, edges):\n",
+    "        # Create a message \"m\" by concatenating \"h\" and \"e\" for each pair of nodes\n",
+    "        nodes_edges = torch.cat([edges.src[\"h\"], edges.data[\"e\"]], dim=-1)\n",
+    "        return {\"m\": nodes_edges}\n",
+    "\n",
+    "    def get_edges_messages(self, edges): # Simply return the messages on the edges\n",
+    "        return {\"m\": edges.data[\"m\"]}\n",
+    "\n",
+    "    def forward(self, g, h, e):\n",
+    "\n",
+    "        # We first concatenate both the node and edge features on the edges\n",
+    "        g.ndata[\"h\"] = h\n",
+    "        g.edata[\"e\"] = e\n",
+    "        g.apply_edges(self.cat_nodes_edges)\n",
+    "\n",
+    "        # Then we apply the mean aggregation to generate a message \"m\"\n",
+    "        g.update_all(message_func=self.get_edges_messages, \n",
+    "                    reduce_func=dgl.function.mean(\"m\", \"h\"))\n",
+    "\n",
+    "        # Finally we apply the FCLayer, and the non-linearities\n",
+    "        h = g.ndata[\"h\"]\n",
+    "        h = self.layer(h)\n",
+    "        h = self.apply_norm_activation_dropout(h)\n",
+    "        return h\n",
+    "\n",
+    "    # Finally, we define all the virtual properties according to how\n",
+    "    # the class works\n",
+    "    @classproperty\n",
+    "    def layer_supports_edges(cls):\n",
+    "        return True\n",
+    "\n",
+    "    @property\n",
+    "    def layer_inputs_edges(self):\n",
+    "        return True\n",
+    "\n",
+    "    @property\n",
+    "    def layer_outputs_edges(self):\n",
+    "        return False\n",
+    "\n",
+    "    @property\n",
+    "    def out_dim_factor(self):\n",
+    "        return 1   "
+   ]
+  },
+  {
+   "source": [
+    "Now, we are ready to test the `ComplexMeanLayer` on some DGL graphs. Note that in this example, we **use** the edge features since they are mandatory."
+   ],
+   "cell_type": "markdown",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "torch.Size([7, 5])\ntorch.Size([7, 11])\n"
+     ]
+    }
+   ],
+   "source": [
+    "graph = deepcopy(bg)\n",
+    "h_in = graph.ndata[\"h\"]\n",
+    "e_in = graph.edata[\"e\"]\n",
+    "layer = ComplexMeanLayer(\n",
+    "            in_dim=in_dim, out_dim=out_dim, in_dim_edges=in_dim_edges,\n",
+    "            activation=\"relu\", dropout=.3, batch_norm=True).to(float)\n",
+    "h_out = layer(graph, h_in, e_in)\n",
+    "\n",
+    "print(h_in.shape)\n",
+    "print(h_out.shape)"
+   ]
+  }
+ ]
+}
\ No newline at end of file
diff --git a/docs/examples/basics/making_gnn_networks.ipynb b/docs/examples/basics/making_gnn_networks.ipynb
new file mode 100644
index 000000000..fb263d7c6
--- /dev/null
+++ b/docs/examples/basics/making_gnn_networks.ipynb
@@ -0,0 +1,215 @@
+{
+ "metadata": {
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.7.9-final"
+  },
+  "orig_nbformat": 2,
+  "kernelspec": {
+   "name": "python3",
+   "display_name": "Python 3.7.9 64-bit ('goli_env': conda)",
+   "metadata": {
+    "interpreter": {
+     "hash": "ce752a05bc046fbf91e84c99a25ae7ca133d76ead2eb04f6c646180f827f8f44"
+    }
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2,
+ "cells": [
+  {
+   "source": [
+    "# Making GNN Networks\n",
+    "\n",
+    "In this example, you will learn how to easily build a full GNN network using any kind of GNN layer."
+   ],
+   "cell_type": "markdown",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stderr",
+     "text": [
+      "Using backend: pytorch\n"
+     ]
+    }
+   ],
+   "source": [
+    "%load_ext autoreload\n",
+    "%autoreload 2\n",
+    "\n",
+    "import torch\n",
+    "import dgl\n",
+    "from copy import deepcopy\n",
+    "\n",
+    "from goli.dgl.dgl_layers import PNAMessagePassingLayer\n",
+    "from goli.dgl.architectures import FullDGLNetwork\n",
+    "\n",
+    "_ = torch.manual_seed(42)"
+   ]
+  },
+  {
+   "source": [
+    "We will first create some simple batched graphs that will be used accross the examples."
+   ],
+   "cell_type": "markdown",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "Graph(num_nodes=7, num_edges=14,\n      ndata_schemes={'h': Scheme(shape=(5,), dtype=torch.float64)}\n      edata_schemes={'e': Scheme(shape=(13,), dtype=torch.float64)})\n"
+     ]
+    }
+   ],
+   "source": [
+    "in_dim = 5          # Input node-feature dimensions\n",
+    "out_dim = 11        # Desired output node-feature dimensions\n",
+    "in_dim_edges = 13   # Input edge-feature dimensions\n",
+    "\n",
+    "# Let's create 2 simple graphs. Here the tensors represent the connectivity between nodes\n",
+    "g1 = dgl.graph((torch.tensor([0, 1, 2]), torch.tensor([1, 2, 3])))\n",
+    "g2 = dgl.graph((torch.tensor([0, 0, 0, 1]), torch.tensor([0, 1, 2, 0])))\n",
+    "\n",
+    "# We add some node features to the graphs\n",
+    "g1.ndata[\"h\"] = torch.rand(g1.num_nodes(), in_dim, dtype=float)\n",
+    "g2.ndata[\"h\"] = torch.rand(g2.num_nodes(), in_dim, dtype=float)\n",
+    "\n",
+    "# We also add some edge features to the graphs\n",
+    "g1.edata[\"e\"] = torch.rand(g1.num_edges(), in_dim_edges, dtype=float)\n",
+    "g2.edata[\"e\"] = torch.rand(g2.num_edges(), in_dim_edges, dtype=float)\n",
+    "\n",
+    "# Finally we batch the graphs in a way compatible with the DGL library\n",
+    "bg = dgl.batch([g1, g2])\n",
+    "bg = dgl.add_self_loop(bg)\n",
+    "\n",
+    "# The batched graph will show as a single graph with 7 nodes\n",
+    "print(bg)\n"
+   ]
+  },
+  {
+   "source": [
+    "## Building the network\n",
+    "\n",
+    "To build the network, we must define the arguments to pass at the different steps:\n",
+    "\n",
+    "- `pre_nn_kwargs`: The parameters used by a feed-forward neural network on the input node-features, before passing to the convolutional layers. See class `FeedForwardNN` for details on the required parameters. Will be ignored if set to `None`.\n",
+    "\n",
+    "- `gnn_kwargs`: The parameters used by a feed-forward **graph** neural network on the features after it has passed through the pre-processing neural network. See class `FeedForwardDGL` for details on the required parameters.\n",
+    "\n",
+    "- `post_nn_kwargs`: The parameters used by a feed-forward neural network on the features after the GNN layers. See class `FeedForwardNN` for details on the required parameters. Will be ignored if set to `None`."
+   ],
+   "cell_type": "markdown",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "temp_dim_1 = 23\n",
+    "temp_dim_2 = 17\n",
+    "\n",
+    "pre_nn_kwargs = {\n",
+    "        \"in_dim\": in_dim,\n",
+    "        \"out_dim\": temp_dim_1,\n",
+    "        \"hidden_dims\": [4, 4, 4],\n",
+    "        \"activation\": \"relu\",\n",
+    "        \"last_activation\": \"none\",\n",
+    "        \"batch_norm\": True,\n",
+    "        \"dropout\": 0.2,    }\n",
+    "\n",
+    "post_nn_kwargs = {\n",
+    "        \"in_dim\": temp_dim_2,\n",
+    "        \"out_dim\": out_dim,\n",
+    "        \"hidden_dims\": [6, 6],\n",
+    "        \"activation\": \"relu\",\n",
+    "        \"last_activation\": \"sigmoid\",\n",
+    "        \"batch_norm\": False,\n",
+    "        \"dropout\": 0.,    }\n",
+    "\n",
+    "layer_kwargs = {\n",
+    "    \"aggregators\": [\"mean\", \"max\", \"sum\"], \n",
+    "    \"scalers\": [\"identity\", \"amplification\"],\n",
+    "    \"in_dim_edges\": in_dim_edges}\n",
+    "\n",
+    "gnn_kwargs = {\n",
+    "    \"in_dim\": temp_dim_1,\n",
+    "    \"out_dim\": temp_dim_2,\n",
+    "    \"hidden_dims\": [5, 5, 5, 5, 5, 5],\n",
+    "    \"residual_type\": \"densenet\",\n",
+    "    \"residual_skip_steps\": 2,\n",
+    "    \"layer_type\": PNAMessagePassingLayer,\n",
+    "    \"pooling\": [\"sum\"],\n",
+    "    \"activation\": \"relu\",\n",
+    "    \"last_activation\": \"none\",\n",
+    "    \"batch_norm\": False,\n",
+    "    \"dropout\": 0.2,\n",
+    "    **layer_kwargs,\n",
+    "}\n",
+    "\n",
+    "gnn_net = FullDGLNetwork(\n",
+    "    pre_nn_kwargs=pre_nn_kwargs, \n",
+    "    gnn_kwargs=gnn_kwargs, \n",
+    "    post_nn_kwargs=post_nn_kwargs).to(float)"
+   ]
+  },
+  {
+   "source": [
+    "## Applying the network\n",
+    "\n",
+    "Once the network is defined, we only need to run the forward pass on the input graphs to get a prediction.\n",
+    "\n",
+    "The network will handle the node and edge features depending on it's parameters and layer type."
+   ],
+   "cell_type": "markdown",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "torch.Size([7, 5])\ntorch.Size([1, 11])\n\n\nDGL_GNN\n---------\n    pre-trans-NN(depth=4, ResidualConnectionNone(skip_steps=1))\n        [FCLayer[5 -> 4 -> 4 -> 4 -> 23] -> Linear({self.out_dim})\n    \n    main-GNN(depth=7, ResidualConnectionDenseNet(skip_steps=2))\n        PNAMessagePassingLayer[23 -> 5 -> 5 -> 5 -> 5 -> 5 -> 5 -> 17]\n        -> Pooling(['sum']) -> FCLayer(17 -> 17, activation=None)\n    \n    post-trans-NN(depth=3, ResidualConnectionNone(skip_steps=1))\n        [FCLayer[17 -> 6 -> 6 -> 11] -> Linear({self.out_dim})\n"
+     ]
+    }
+   ],
+   "source": [
+    "graph = deepcopy(bg)\n",
+    "h_in = graph.ndata[\"h\"]\n",
+    "\n",
+    "h_out = gnn_net(graph)\n",
+    "\n",
+    "print(h_in.shape)\n",
+    "print(h_out.shape)\n",
+    "print(\"\\n\")\n",
+    "print(gnn_net)\n"
+   ]
+  }
+ ]
+}
\ No newline at end of file
diff --git a/docs/examples/basics/using_gnn_layers.ipynb b/docs/examples/basics/using_gnn_layers.ipynb
new file mode 100644
index 000000000..e8b7fa667
--- /dev/null
+++ b/docs/examples/basics/using_gnn_layers.ipynb
@@ -0,0 +1,379 @@
+{
+ "metadata": {
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.7.9-final"
+  },
+  "orig_nbformat": 2,
+  "kernelspec": {
+   "name": "python3",
+   "display_name": "Python 3.7.9 64-bit ('goli_env': conda)",
+   "metadata": {
+    "interpreter": {
+     "hash": "ce752a05bc046fbf91e84c99a25ae7ca133d76ead2eb04f6c646180f827f8f44"
+    }
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2,
+ "cells": [
+  {
+   "source": [
+    "# Using GNN layers\n",
+    "\n",
+    "In this example, you will learn how to use the GCN, GIN, Gated-GCN and PNA layers in a simple `forward` context."
+   ],
+   "cell_type": "markdown",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stderr",
+     "text": [
+      "Using backend: pytorch\n"
+     ]
+    }
+   ],
+   "source": [
+    "%load_ext autoreload\n",
+    "%autoreload 2\n",
+    "\n",
+    "import torch\n",
+    "import dgl\n",
+    "from copy import deepcopy\n",
+    "\n",
+    "from goli.dgl.dgl_layers import (\n",
+    "    GCNLayer,\n",
+    "    GINLayer,\n",
+    "    GATLayer,\n",
+    "    GatedGCNLayer,\n",
+    "    PNAConvolutionalLayer,\n",
+    "    PNAMessagePassingLayer,\n",
+    ")\n",
+    "\n",
+    "_ = torch.manual_seed(42)"
+   ]
+  },
+  {
+   "source": [
+    "We will first create some simple batched graphs that will be used accross the examples."
+   ],
+   "cell_type": "markdown",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "Graph(num_nodes=7, num_edges=14,\n      ndata_schemes={'h': Scheme(shape=(5,), dtype=torch.float64)}\n      edata_schemes={'e': Scheme(shape=(13,), dtype=torch.float64)})\n"
+     ]
+    }
+   ],
+   "source": [
+    "in_dim = 5          # Input node-feature dimensions\n",
+    "out_dim = 11        # Desired output node-feature dimensions\n",
+    "in_dim_edges = 13   # Input edge-feature dimensions\n",
+    "\n",
+    "# Let's create 2 simple graphs. Here the tensors represent the connectivity between nodes\n",
+    "g1 = dgl.graph((torch.tensor([0, 1, 2]), torch.tensor([1, 2, 3])))\n",
+    "g2 = dgl.graph((torch.tensor([0, 0, 0, 1]), torch.tensor([0, 1, 2, 0])))\n",
+    "\n",
+    "# We add some node features to the graphs\n",
+    "g1.ndata[\"h\"] = torch.rand(g1.num_nodes(), in_dim, dtype=float)\n",
+    "g2.ndata[\"h\"] = torch.rand(g2.num_nodes(), in_dim, dtype=float)\n",
+    "\n",
+    "# We also add some edge features to the graphs\n",
+    "g1.edata[\"e\"] = torch.rand(g1.num_edges(), in_dim_edges, dtype=float)\n",
+    "g2.edata[\"e\"] = torch.rand(g2.num_edges(), in_dim_edges, dtype=float)\n",
+    "\n",
+    "# Finally we batch the graphs in a way compatible with the DGL library\n",
+    "bg = dgl.batch([g1, g2])\n",
+    "bg = dgl.add_self_loop(bg)\n",
+    "\n",
+    "# The batched graph will show as a single graph with 7 nodes\n",
+    "print(bg)\n"
+   ]
+  },
+  {
+   "source": [
+    "## GCN Layer\n",
+    "\n",
+    "To use the GCN layer from the *Kipf et al.* paper, the steps are very simple. We create the layer with the desired attributes, and apply it to the graph.\n",
+    "\n",
+    "<sub>Kipf, Thomas N., and Max Welling. \"Semi-supervised classification with graph convolutional networks.\" arXiv preprint arXiv:1609.02907 (2016).</sub>"
+   ],
+   "cell_type": "markdown",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "GCNLayer(5 -> 11, activation=relu)\ntorch.Size([7, 5])\ntorch.Size([7, 11])\n"
+     ]
+    }
+   ],
+   "source": [
+    "# We first need to extract the node features from the graph.\n",
+    "# The GCN method doesn't support edge features, so we ignore them\n",
+    "graph = deepcopy(bg)\n",
+    "h_in = graph.ndata[\"h\"]\n",
+    "\n",
+    "# We create the layer\n",
+    "layer = GCNLayer(\n",
+    "            in_dim=in_dim, out_dim=out_dim, \n",
+    "            activation=\"relu\", dropout=.3, batch_norm=True).to(float)\n",
+    "\n",
+    "# We apply the forward loop on the node features\n",
+    "h_out = layer(graph, h_in)\n",
+    "\n",
+    "# 7 is the number of nodes, 5 number of input features and 11 number of output features\n",
+    "print(layer)\n",
+    "print(h_in.shape)\n",
+    "print(h_out.shape)"
+   ]
+  },
+  {
+   "source": [
+    "## GIN Layer\n",
+    "\n",
+    "To use the GIN layer from the *Xu et al.* paper, the steps are identical to GCN.\n",
+    "\n",
+    "<sub>Xu, Keyulu, et al. \"How powerful are graph neural networks?.\" arXiv preprint arXiv:1810.00826 (2018).</sub>"
+   ],
+   "cell_type": "markdown",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "GINLayer(5 -> 11, activation=relu)\ntorch.Size([7, 5])\ntorch.Size([7, 11])\n"
+     ]
+    }
+   ],
+   "source": [
+    "graph = deepcopy(bg)\n",
+    "h_in = graph.ndata[\"h\"]\n",
+    "layer = GINLayer(\n",
+    "            in_dim=in_dim, out_dim=out_dim, \n",
+    "            activation=\"relu\", dropout=.3, batch_norm=True).to(float)\n",
+    "h_out = layer(graph, h_in)\n",
+    "\n",
+    "print(layer)\n",
+    "print(h_in.shape)\n",
+    "print(h_out.shape)"
+   ]
+  },
+  {
+   "source": [
+    "## GAT Layer\n",
+    "\n",
+    "To use the GAT layer from the *Velickovic et al.* paper, the steps are identical to GCN, but the output dimension is multiplied by the number of heads.\n",
+    "\n",
+    "<sub>Velickovic, Petar, et al. \"Graph attention networks.\" arXiv preprint arXiv:1710.10903 (2017).</sub>"
+   ],
+   "cell_type": "markdown",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "GATLayer(5 -> 11 * 5, activation=elu)\ntorch.Size([7, 5])\ntorch.Size([7, 55])\n"
+     ]
+    }
+   ],
+   "source": [
+    "graph = deepcopy(bg)\n",
+    "h_in = graph.ndata[\"h\"]\n",
+    "layer = GATLayer(\n",
+    "            in_dim=in_dim, out_dim=out_dim, num_heads=5,\n",
+    "            activation=\"elu\", dropout=.3, batch_norm=True).to(float)\n",
+    "h_out = layer(graph, h_in)\n",
+    "\n",
+    "print(layer)\n",
+    "print(h_in.shape)\n",
+    "print(h_out.shape)"
+   ]
+  },
+  {
+   "source": [
+    "## Gated-GCN Layer\n",
+    "\n",
+    "To use the Gated-GCN layer from the *Bresson et al.* paper, the steps are different since the layer requires edge features as inputs, and outputs new edge features.\n",
+    "\n",
+    "<sub>Bresson, Xavier, and Thomas Laurent. \"Residual gated graph convnets.\" arXiv preprint arXiv:1711.07553 (2017).</sub>"
+   ],
+   "cell_type": "markdown",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "GatedGCNLayer(5 -> 11, activation=relu)\ntorch.Size([7, 5])\ntorch.Size([7, 11])\ntorch.Size([14, 13])\ntorch.Size([14, 11])\n"
+     ]
+    }
+   ],
+   "source": [
+    "# We first need to extract the node and edge features from the graph.\n",
+    "graph = deepcopy(bg)\n",
+    "h_in = graph.ndata[\"h\"]\n",
+    "e_in = graph.edata[\"e\"]\n",
+    "\n",
+    "# We create the layer\n",
+    "layer = GatedGCNLayer(\n",
+    "        in_dim=in_dim, out_dim=out_dim, \n",
+    "        in_dim_edges=in_dim_edges, out_dim_edges=out_dim,\n",
+    "        activation=\"relu\", dropout=.3, batch_norm=True).to(float)\n",
+    "\n",
+    "# We apply the forward loop on the node features\n",
+    "h_out, e_out = layer(graph, h_in, e_in)\n",
+    "\n",
+    "# 7 is the number of nodes, 5 number of input features and 11 number of output features\n",
+    "# 13 is the number of input edge features\n",
+    "print(layer)\n",
+    "print(h_in.shape)\n",
+    "print(h_out.shape)\n",
+    "print(e_in.shape)\n",
+    "print(e_out.shape)"
+   ]
+  },
+  {
+   "source": [
+    "# PNA\n",
+    "\n",
+    "PNA is a multi-aggregator method proposed by *Corso et al.*. It supports 2 types of aggregations, convolutional *PNA-conv* or message passing *PNA-msgpass*.\n",
+    "\n",
+    "<sub>PNA: Principal Neighbourhood Aggregation \n",
+    "Gabriele Corso, Luca Cavalleri, Dominique Beaini, Pietro Lio, Petar Velickovic\n",
+    "https://arxiv.org/abs/2004.05718</sub>\n",
+    "\n",
+    "## PNA-conv\n",
+    "\n",
+    "First, let's focus on the PNA-conv. In this case, it works exactly as the GCN and GIN methods."
+   ],
+   "cell_type": "markdown",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "PNAConvolutionalLayer(5 -> 11, activation=relu)\ntorch.Size([7, 5])\ntorch.Size([7, 11])\n"
+     ]
+    }
+   ],
+   "source": [
+    "graph = deepcopy(bg)\n",
+    "h_in = graph.ndata[\"h\"]\n",
+    "\n",
+    "# We create the layer, and need to specify the aggregators and scalers\n",
+    "layer = PNAConvolutionalLayer(\n",
+    "    in_dim=in_dim, out_dim=out_dim, \n",
+    "    aggregators=[\"mean\", \"max\", \"min\", \"std\"],\n",
+    "    scalers=[\"identity\", \"amplification\", \"attenuation\"],\n",
+    "    activation=\"relu\", dropout=.3, batch_norm=True).to(float)\n",
+    "\n",
+    "h_out = layer(graph, h_in)\n",
+    "\n",
+    "print(layer)\n",
+    "print(h_in.shape)\n",
+    "print(h_out.shape)"
+   ]
+  },
+  {
+   "source": [
+    "## PNA-msgpass\n",
+    "\n",
+    "The PNA message passing is typically more powerful that the convolutional one, and it supports edges as inputs, but doesn't output edges. It's usage is very similar to the *PNA-conv*, but we need to specify the edge dimensions and features."
+   ],
+   "cell_type": "markdown",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "PNAMessagePassingLayer(5 -> 11, activation=relu)\ntorch.Size([7, 5])\ntorch.Size([7, 11])\n"
+     ]
+    }
+   ],
+   "source": [
+    "graph = deepcopy(bg)\n",
+    "h_in = graph.ndata[\"h\"]\n",
+    "e_in = graph.edata[\"e\"]\n",
+    "\n",
+    "# We create the layer, and need to specify the aggregators and scalers\n",
+    "layer = PNAMessagePassingLayer(\n",
+    "    in_dim=in_dim, out_dim=out_dim, in_dim_edges=in_dim_edges,\n",
+    "    aggregators=[\"mean\", \"max\", \"min\", \"std\"],\n",
+    "    scalers=[\"identity\", \"amplification\", \"attenuation\"],\n",
+    "    activation=\"relu\", dropout=.3, batch_norm=True).to(float)\n",
+    "\n",
+    "h_out = layer(graph, h_in, e_in)\n",
+    "\n",
+    "print(layer)\n",
+    "print(h_in.shape)\n",
+    "print(h_out.shape)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ]
+}
\ No newline at end of file
diff --git a/goli/commons/config_loader.py b/goli/commons/config_loader.py
index cfec490a3..a5b8db32c 100644
--- a/goli/commons/config_loader.py
+++ b/goli/commons/config_loader.py
@@ -35,7 +35,7 @@
 
 
 class DimensionalityReduction:
-    """
+    r"""
     Dimensionality reduction helper. It wraps around a standart reduction like pca, tsne, umap and adds an inverse
     transform either by pre-implemented inverse transform or a knn based inverse transform.
 
diff --git a/goli/dgl/architectures.py b/goli/dgl/architectures.py
index def69bc37..b17c4abe4 100644
--- a/goli/dgl/architectures.py
+++ b/goli/dgl/architectures.py
@@ -5,13 +5,14 @@
 from copy import deepcopy
 import dgl
 from dgl.nn.pytorch.glob import mean_nodes, sum_nodes
-from typing import List, Dict
+from typing import List, Dict, Tuple, Union, Callable, Any
 import inspect
 
 from goli.dgl.base_layers import FCLayer, get_activation
 from goli.dgl.dgl_layers.pooling import parse_pooling_layer, VirtualNode
 
 from goli.dgl.dgl_layers import (
+    BaseDGLLayer,
     GATLayer,
     GCNLayer,
     GINLayer,
@@ -21,6 +22,7 @@
 )
 
 from goli.dgl.residual_connections import (
+    ResidualConnectionBase,
     ResidualConnectionConcat,
     ResidualConnectionDenseNet,
     ResidualConnectionNone,
@@ -52,19 +54,76 @@
 class FeedForwardNN(nn.Module):
     def __init__(
         self,
-        in_dim,
-        out_dim,
-        hidden_dims,
-        activation="relu",
-        last_activation="none",
-        batch_norm=False,
-        dropout=0.25,
-        residual_type="none",
-        residual_skip_steps=1,
-        name="LNN",
-        layer_type=FCLayer,
+        in_dim: int,
+        out_dim: int,
+        hidden_dims: List[int],
+        activation: Union[str, Callable] = "relu",
+        last_activation: Union[str, Callable] = "none",
+        dropout: float = 0.0,
+        batch_norm: bool = False,
+        residual_type: str = "none",
+        residual_skip_steps: int = 1,
+        name: str = "LNN",
+        layer_type: Union[str, nn.Module] = "fc",
         **layer_kwargs,
     ):
+        r"""
+        A flexible neural network architecture, with variable hidden dimensions,
+        support for multiple layer types, and support for different residual
+        connections.
+
+        Parameters:
+
+            in_dim:
+                Input feature dimensions of the layer
+
+            out_dim:
+                Output feature dimensions of the layer
+
+            hidden_dims:
+                List of dimensions in the hidden layers.
+                Be careful, the "simple" residual type only supports
+                hidden dimensions of the same value.
+
+            activation:
+                activation function to use in the hidden layers.
+
+            last_activation:
+                activation function to use in the last layer.
+
+            dropout:
+                The ratio of units to dropout. Must be between 0 and 1
+
+            batch_norm:
+                Whether to use batch normalization
+
+            residual_type:
+                - "none": No residual connection
+                - "simple": Residual connection similar to the ResNet architecture.
+                  See class `ResidualConnectionSimple`
+                - "weighted": Residual connection similar to the Resnet architecture,
+                  but with weights applied before the summation. See class `ResidualConnectionWeighted`
+                - "concat": Residual connection where the residual is concatenated instead
+                  of being added.
+                - "densenet": Residual connection where the residual of all previous layers
+                  are concatenated. This leads to a strong increase in the number of parameters
+                  if there are multiple hidden layers.
+
+            residual_skip_steps:
+                The number of steps to skip between each residual connection.
+                If `1`, all the layers are connected. If `2`, half of the
+                layers are connected.
+
+            name:
+                Name attributed to the current network, for display and printing
+                purposes.
+
+            layer_type:
+                The type of layers to use in the network.
+                Either "fc" as the `FCLayer`, or a class representing the `nn.Module`
+                to use.
+
+        """
 
         super().__init__()
 
@@ -93,14 +152,19 @@ def __init__(
         self._check_bad_arguments()
 
     def _check_bad_arguments(self):
-        if (self.residual_type == "simple") and (self.hidden_dims[:-1] == self.hidden_dims[1:]):
+        r"""
+        Raise comprehensive errors if the arguments seem wrong
+        """
+        if (self.residual_type == "simple") and not (self.hidden_dims[:-1] == self.hidden_dims[1:]):
             raise ValueError(
                 f"When using the residual_type={self.residual_type}"
                 + f", all elements in the hidden_dims must be equal. Provided:{self.hidden_dims}"
             )
 
     def _register_hparams(self):
-        # Register the Hyper-parameters to be compatible with Pytorch-Lightning and Tensorboard
+        r"""
+        Register the hyperparameters for tracking by Pytorch-lightning
+        """
         self.hparams = {
             f"{self.name}.out_dim": self.out_dim,
             f"{self.name}.hidden_dims": self.hidden_dims,
@@ -115,11 +179,16 @@ def _register_hparams(self):
             f"{self.name}.layer_kwargs": str(self.layer_kwargs),
         }
 
-    def _parse_class_from_dict(self, name_or_class, class_dict):
+    def _parse_class_from_dict(
+        self, name_or_class: Union[type, str], class_dict: Dict[str, type]
+    ) -> Tuple[type, str]:
+        r"""
+        Register the hyperparameters for tracking by Pytorch-lightning
+        """
         if isinstance(name_or_class, str):
             obj_name = name_or_class.lower()
             obj_class = class_dict[obj_name]
-        elif isinstance(callable):
+        elif callable(name_or_class):
             obj_name = str(name_or_class)
             obj_class = name_or_class
         else:
@@ -127,16 +196,19 @@ def _parse_class_from_dict(self, name_or_class, class_dict):
 
         return obj_class, obj_name
 
-    def _create_residual_connection(self, out_dims):
-
+    def _create_residual_connection(self, out_dims: List[int]) -> Tuple[ResidualConnectionBase, List[int]]:
+        r"""
+        Create the residual connection classes.
+        The out_dims is only used if the residual classes requires weights
+        """
         if self.residual_class.has_weights:
             residual_layer = self.residual_class(
                 skip_steps=self.residual_skip_steps,
-                full_dims=out_dims,
+                out_dims=out_dims,
                 dropout=self.dropout,
                 activation=self.activation,
                 batch_norm=self.batch_norm,
-                bias=self.bias,
+                bias=False,
             )
         else:
             residual_layer = self.residual_class(skip_steps=self.residual_skip_steps)
@@ -146,6 +218,13 @@ def _create_residual_connection(self, out_dims):
         return residual_layer, residual_out_dims
 
     def _create_layers(self):
+        r"""
+        Create all the necessary layers for the network.
+        It's a bit complicated to explain what's going on in this function,
+        but it must manage the varying features sizes caused by:
+
+        - The presence of different types of residual connections
+        """
 
         self.residual_layer, residual_out_dims = self._create_residual_connection(out_dims=self.full_dims[1:])
 
@@ -171,17 +250,31 @@ def _create_layers(self):
                 )
             )
 
-            this_in_dim = residual_out_dims[ii]
+            if ii < len(residual_out_dims):
+                this_in_dim = residual_out_dims[ii]
 
-        self.out_linear = nn.Linear(in_features=this_in_dim, out_features=self.out_dim)
+    def forward(self, h: torch.Tensor) -> torch.Tensor:
+        r"""
+        Apply the neural network on the input features.
+
+        Parameters:
+
+            h: `torch.Tensor[..., Din]`:
+                Input feature tensor, before the network.
+                `Din` is the number of input features
+
+        Returns:
+
+            `torch.Tensor[..., Dout]`:
+                Output feature tensor, after the network.
+                `Dout` is the number of output features
 
-    def forward(self, h):
+        """
         h_prev = None
         for ii, layer in enumerate(self.layers):
             h = layer.forward(h)
-            h, h_prev = self.residual.forward(h, h_prev, step_idx=ii)
-
-        h = self.out_linear.forward(h)
+            if ii < len(self.layers) - 1:
+                h, h_prev = self.residual_layer.forward(h, h_prev, step_idx=ii)
 
         return h
 
@@ -189,8 +282,8 @@ def __repr__(self):
         r"""
         Controls how the class is printed
         """
-        class_str = f"{self.__class__.__name__}(depth={self.depth})"
-        layer_str = f"[{self.layer_name}[{' -> '.join(map(str, self.full_dims))}]"
+        class_str = f"{self.name}(depth={self.depth}, {self.residual_layer})\n    "
+        layer_str = f"[{self.layer_class.__name__}[{' -> '.join(map(str, self.full_dims))}]"
         out_str = " -> Linear({self.out_dim})"
 
         return class_str + layer_str + out_str
@@ -199,32 +292,132 @@ def __repr__(self):
 class FeedForwardDGL(FeedForwardNN):
     def __init__(
         self,
-        in_dim,
-        out_dim,
-        hidden_dims,
-        activation="relu",
-        last_activation="none",
-        batch_norm=False,
-        dropout=0.25,
-        residual_type="none",
-        residual_skip_steps=1,
-        in_dim_edges=0,
-        hidden_dims_edges=[],
-        pooling=["sum"],
-        name="GNN",
-        layer_class="gcn",
-        virtual_node="none",
+        in_dim: int,
+        out_dim: int,
+        hidden_dims: List[int],
+        activation: Union[str, Callable] = "relu",
+        last_activation: Union[str, Callable] = "none",
+        dropout: float = 0.0,
+        batch_norm: bool = False,
+        residual_type: str = "none",
+        residual_skip_steps: int = 1,
+        in_dim_edges: int = 0,
+        hidden_dims_edges: List[int] = [],
+        pooling: List[Union[str, Callable]] = ["sum"],
+        name: str = "GNN",
+        layer_type: Union[str, nn.Module] = "gcn",
+        virtual_node: str = "none",
         **layer_kwargs,
     ):
+        r"""
+        A flexible neural network architecture, with variable hidden dimensions,
+        support for multiple layer types, and support for different residual
+        connections.
+
+        This class is meant to work with different DGL-based graph neural networks
+        layers. Any layer must inherit from `goli.dgl.dgl_layers.BaseDGLLayer`.
+
+        Parameters:
+
+            in_dim:
+                Input feature dimensions of the layer
+
+            out_dim:
+                Output feature dimensions of the layer
+
+            hidden_dims:
+                List of dimensions in the hidden layers.
+                Be careful, the "simple" residual type only supports
+                hidden dimensions of the same value.
+
+            activation:
+                activation function to use in the hidden layers.
+
+            last_activation:
+                activation function to use in the last layer.
+
+            dropout:
+                The ratio of units to dropout. Must be between 0 and 1
+
+            batch_norm:
+                Whether to use batch normalization
+
+            residual_type:
+                - "none": No residual connection
+                - "simple": Residual connection similar to the ResNet architecture.
+                  See class `ResidualConnectionSimple`
+                - "weighted": Residual connection similar to the Resnet architecture,
+                  but with weights applied before the summation. See class `ResidualConnectionWeighted`
+                - "concat": Residual connection where the residual is concatenated instead
+                  of being added.
+                - "densenet": Residual connection where the residual of all previous layers
+                  are concatenated. This leads to a strong increase in the number of parameters
+                  if there are multiple hidden layers.
+
+            residual_skip_steps:
+                The number of steps to skip between each residual connection.
+                If `1`, all the layers are connected. If `2`, half of the
+                layers are connected.
+
+            in_dim_edges:
+                Input edge-feature dimensions of the network. Keep at 0 if not using
+                edge features, or if the layer doesn't support edges.
+
+            hidden_dims_edges:
+                Hidden dimensions for the edges. Most models don't support it, so it
+                should only be used for those that do, i.e. `GatedGCNLayer`
+
+            pooling:
+                The pooling types to use. Multiple pooling can be used, and their
+                results will be concatenated.
+                For node feature predictions, use `["none"]`.
+                For graph feature predictions see `goli.dgl.dgl_layers.pooling.parse_pooling_layer`.
+                The list must either contain Callables, or the string below
+
+                - "none": No pooling is applied
+                - "sum": `SumPooling`
+                - "mean": `MeanPooling`
+                - "max": `MaxPooling`
+                - "min": `MinPooling`
+                - "std": `StdPooling`
+                - "s2s": `Set2Set`
+
+            name:
+                Name attributed to the current network, for display and printing
+                purposes.
+
+            layer_type:
+                The type of layers to use in the network.
+                Either a class that inherits from `goli.dgl.dgl_layers.BaseDGLLayer`,
+                or one of the following strings
+
+                - "gcn": `GCNLayer`
+                - "gin": `GINLayer`
+                - "gat": `GATLayer`
+                - "gated-gcn": `GatedGCNLayer`
+                - "pna-conv": `PNAConvolutionalLayer`
+                - "pna-msgpass": `PNAMessagePassingLayer`
+
+            virtual_node:
+                A string associated to the type of virtual node to use,
+                either `None`, "none", "mean", "sum", "max", "logsum".
+                See `goli.dgl.dgl_layers.VirtualNode`.
+
+                The virtual node will not use any residual connection if `residual_type`
+                is "none". Otherwise, it will use a simple ResNet like residual
+                connection.
+
+        """
 
         # Initialize the additional attributes
         self.in_dim_edges = in_dim_edges
         self.hidden_dims_edges = hidden_dims_edges
-        self.edge_features = len(self.hidden_dims_edges) > 0
         self.full_dims_edges = None
-        if self.edge_features:
+        if len(self.hidden_dims_edges) > 0:
             self.full_dims_edges = [self.in_dim_edges] + self.hidden_dims_edges + [self.hidden_dims_edges[-1]]
+
         self.virtual_node = virtual_node.lower()
+        self.pooling = pooling
 
         # Initialize the parent `FeedForwardNN`
         super().__init__(
@@ -237,26 +430,42 @@ def __init__(
             residual_type=residual_type,
             residual_skip_steps=residual_skip_steps,
             name=name,
-            layer_class=layer_class,
+            layer_type=layer_type,
             dropout=dropout,
             **layer_kwargs,
         )
 
-        # Initialize input and output linear layers
-        self.in_linear = nn.Linear(in_features=self.in_dim, out_features=self.hidden_dims[0])
-
     def _check_bad_arguments(self):
+        r"""
+        Raise comprehensive errors if the arguments seem wrong
+        """
         super()._check_bad_arguments()
-        if self.edge_features and not self.layer_class.layer_supports_edges:
+        if (
+            (self.in_dim_edges > 0) or (self.full_dims_edges is not None)
+        ) and not self.layer_class.layer_supports_edges:
             raise ValueError(f"Cannot use edge features with class `{self.layer_class}`")
 
     def _register_hparams(self):
+        r"""
+        Register the hyperparameters for tracking by Pytorch-lightning
+        """
         return super()._register_hparams()
         self.hparams["hidden_edge_dim"] = self.hidden_edge_dim
         self.hparams["pooling"] = self.pooling
         self.hparams["intermittent_pooling"] = self.intermittent_pooling
 
     def _create_layers(self):
+        r"""
+        Create all the necessary layers for the network.
+        It's a bit complicated to explain what's going on in this function,
+        but it must manage the varying features sizes caused by:
+
+        - The presence of different types of residual connections
+        - The presence or absence of edges
+        - The output dimensions varying for different networks i.e. `GatLayer` outputs different feature sizes according to the number of heads
+        - The presence or absence of virtual nodes
+        - The different possible pooling, and the concatenation of multiple pooling together.
+        """
 
         residual_layer_temp, residual_out_dims = self._create_residual_connection(out_dims=self.full_dims[1:])
 
@@ -264,23 +473,28 @@ def _create_layers(self):
         self.layers = nn.ModuleList()
         self.virtual_node_layers = nn.ModuleList()
         this_in_dim = self.full_dims[0]
+        this_activation = self.activation
 
+        # Find the appropriate edge dimensions, depending if edges are used,
+        # And if the residual is required for the edges
         this_in_dim_edges, this_out_dim_edges = None, None
         if self.full_dims_edges is not None:
             this_in_dim_edges, this_out_dim_edges = self.full_dims_edges[0:2]
-            residual_out_dims_edges = residual_layer_temp.get_true_out_dims(full_dims_edges[1:])
-
-        this_activation = self.activation
+            residual_out_dims_edges = residual_layer_temp.get_true_out_dims(self.full_dims_edges[1:])
+        elif self.in_dim_edges > 0:
+            this_in_dim_edges = self.in_dim_edges
         layer_out_dims_edges = []
 
+        # Create all the layers in a loop
         for ii in range(self.depth):
             this_out_dim = self.full_dims[ii + 1]
 
             if ii == self.depth - 1:
                 this_activation = self.last_activation
 
+            # Find the edge key-word arguments depending on the layer type and residual connection
             this_edge_kwargs = {}
-            if self.layer_class.layer_supports_edges() and self.edge_features:
+            if self.layer_class.layer_supports_edges and self.in_dim_edges > 0:
                 this_edge_kwargs["in_dim_edges"] = this_in_dim_edges
                 if "out_dim_edges" in inspect.signature(self.layer_class.__init__).parameters.keys():
                     layer_out_dims_edges.append(self.full_dims_edges[ii + 1])
@@ -299,42 +513,72 @@ def _create_layers(self):
                 )
             )
 
-            # Create the Virtual Node layer
-            self.virtual_node_layers.append(
-                VirtualNode(
-                    dim=this_out_dim,
-                    activation=this_activation,
-                    dropout=self.dropout,
-                    batch_norm=self.batch_norm,
-                    bias=True,
-                    vn_type=self.virtual_node,
-                    residual=self.residual,
+            # Create the Virtual Node layer, except at the last layer
+            if ii < len(residual_out_dims):
+                self.virtual_node_layers.append(
+                    VirtualNode(
+                        dim=this_out_dim,
+                        activation=this_activation,
+                        dropout=self.dropout,
+                        batch_norm=self.batch_norm,
+                        bias=True,
+                        vn_type=self.virtual_node,
+                        residual=self.residual_type is not None,
+                    )
                 )
-            )
 
             # Get the true input dimension of the next layer,
             # by factoring both the residual connection and GNN layer type
-            this_in_dim = residual_out_dims[ii] * layers[ii - 1].out_dim_factor
-            this_in_dim_edges = residual_out_dims_edges[ii] * layers[ii - 1].out_dim_factor
+            if ii < len(residual_out_dims):
+                this_in_dim = residual_out_dims[ii] * self.layers[ii - 1].out_dim_factor
+                if self.full_dims_edges is not None:
+                    this_in_dim_edges = residual_out_dims_edges[ii] * self.layers[ii - 1].out_dim_factor
 
         layer_out_dims = [layer.out_dim_factor * layer.out_dim for layer in self.layers]
 
-        # Initialize residual, pooling and output layers
+        # Initialize residual and pooling layers
         self.residual_layer, _ = self._create_residual_connection(out_dims=layer_out_dims)
         if len(layer_out_dims_edges) > 0:
             self.residual_edges_layer, _ = self._create_residual_connection(out_dims=layer_out_dims_edges)
         else:
             self.residual_edges_layer = None
         self.global_pool_layer, out_pool_dim = parse_pooling_layer(layer_out_dims[-1], self.pooling)
-        self.out_linear = nn.Linear(in_features=out_pool_dim, out_features=self.out_dim)
 
-    def _pool_layer_forward(self, graph, h):
+        # Output linear layer
+        self.out_linear = FCLayer(
+            in_dim=out_pool_dim,
+            out_dim=self.out_dim,
+            activation="none",
+            dropout=self.dropout,
+            batch_norm=self.batch_norm,
+        )
+
+    def _pool_layer_forward(self, g, h):
+        r"""
+        Apply the graph pooling layer, followed by the linear output layer.
+
+        Parameters:
+
+            g: dgl.DGLGraph
+                graph on which the convolution is done
+
+            h (torch.Tensor[..., N, Din]):
+                Node feature tensor, before convolution.
+                `N` is the number of nodes, `Din` is the output size of the last DGL layer
+
+        Returns:
+
+            torch.Tensor[..., M, Din] or torch.Tensor[..., N, Din]:
+                Node feature tensor, after convolution.
+                `N` is the number of nodes, `M` is the number of graphs, `Dout` is the output dimension ``self.out_dim``
+                If the pooling is `None`, the dimension is `N`, otherwise it is `M`
+
+        """
 
-        # Pool the nodes together
         if len(self.global_pool_layer) > 0:
             pooled_h = []
             for this_pool in self.global_pool_layer:
-                pooled_h.append(this_pool(graph, h))
+                pooled_h.append(this_pool(g, h))
             pooled_h = torch.cat(pooled_h, dim=-1)
         else:
             pooled_h = h
@@ -343,51 +587,172 @@ def _pool_layer_forward(self, graph, h):
 
         return pooled_h
 
-    def _dgl_layer_forward(self, layer, graph, h, e, h_prev, e_prev, step_idx):
+    def _dgl_layer_forward(
+        self,
+        layer: BaseDGLLayer,
+        g: dgl.DGLGraph,
+        h: torch.Tensor,
+        e: Union[torch.Tensor, None],
+        h_prev: Union[torch.Tensor, None],
+        e_prev: Union[torch.Tensor, None],
+        step_idx: int,
+    ) -> Tuple[torch.Tensor, Union[torch.Tensor, None], Union[torch.Tensor, None], Union[torch.Tensor, None]]:
+        r"""
+        Apply the *i-th* DGL graph layer, where *i* is the index given by `step_idx`.
+        The layer is applied differently depending if there are edge features or not.
+
+        Then, the residual is also applied on both the features and the edges (if applicable)
+
+        Parameters:
+
+            layer:
+                The DGL layer used for the convolution
+
+            g:
+                graph on which the convolution is done
+
+            h (torch.Tensor[..., N, Din]):
+                Node feature tensor, before convolution.
+                `N` is the number of nodes, `Din` is the input features
+
+            e (torch.Tensor[..., N, Ein]):
+                Edge feature tensor, before convolution.
+                `N` is the number of nodes, `Ein` is the input edge features
+
+            h_prev:
+                Node feature of the previous residual connection, or `None`
+
+            e_prev:
+                Edge feature of the previous residual connection, or `None`
+
+            step_idx:
+                The current step idx in the forward loop
+
+        Returns:
+
+            h (torch.Tensor[..., N, Dout]):
+                Node feature tensor, after convolution and residual.
+                `N` is the number of nodes, `Dout` is the output features of the layer and residual
+
+            e:
+                Edge feature tensor, after convolution and residual.
+                `N` is the number of nodes, `Ein` is the input edge features
+
+            h_prev:
+                Node feature tensor to be used at the next residual connection, or `None`
+
+            e_prev:
+                Edge feature tensor to be used at the next residual connection, or `None`
+
+        """
 
         # Apply the GNN layer with the right inputs/outputs
         if layer.layer_inputs_edges and layer.layer_outputs_edges:
-            h, e = layer(g=graph, h=h, e=e)
+            h, e = layer(g=g, h=h, e=e)
         elif layer.layer_inputs_edges:
-            h = layer(g=graph, h=h, e=e)
+            h = layer(g=g, h=h, e=e)
         elif layer.layer_outputs_edges:
-            h, e = layer(g=graph, h=h)
+            h, e = layer(g=g, h=h)
         else:
-            h = layer(g=graph, h=h)
+            h = layer(g=g, h=h)
 
         # Apply the residual layers on the features and edges (if applicable)
-        h, h_prev = self.residual_layer.forward(h, h_prev, step_idx=step_idx)
-        if (self.residual_edges_layer is not None) and (layer.layer_outputs_edges):
-            e, e_prev = self.residual_layer.forward(e, e_prev, step_idx=step_idx)
+        if step_idx < len(self.layers) - 1:
+            h, h_prev = self.residual_layer.forward(h, h_prev, step_idx=step_idx)
+            if (self.residual_edges_layer is not None) and (layer.layer_outputs_edges):
+                e, e_prev = self.residual_edges_layer.forward(e, e_prev, step_idx=step_idx)
 
         return h, e, h_prev, e_prev
 
-    def _virtual_node_forward(self, g, h, vn_h, step_idx):
-        # Apply the Virtual Node
+    def _virtual_node_forward(
+        self, g: dgl.DGLGraph, h: torch.Tensor, vn_h: torch.Tensor, step_idx: int
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        r"""
+        Apply the *i-th* virtual node layer, where *i* is the index given by `step_idx`.
+
+        Parameters:
+
+            g:
+                graph on which the convolution is done
+
+            h (torch.Tensor[..., N, Din]):
+                Node feature tensor, before convolution.
+                `N` is the number of nodes, `Din` is the input features
+
+            vn_h (torch.Tensor[..., M, Din]):
+                Graph feature of the previous virtual node, or `None`
+                `M` is the number of graphs, `Din` is the input features
+                It is added to the result after the MLP, as a residual connection
+
+            step_idx:
+                The current step idx in the forward loop
+
+        Returns:
+
+            `h = torch.Tensor[..., N, Dout]`:
+                Node feature tensor, after convolution and residual.
+                `N` is the number of nodes, `Dout` is the output features of the layer and residual
+
+            `vn_h = torch.Tensor[..., M, Dout]`:
+                Graph feature tensor to be used at the next virtual node, or `None`
+                `M` is the number of graphs, `Dout` is the output features
+
+        """
+
         if step_idx == 0:
-            vn_h = 0
+            vn_h = 0.0
         if step_idx < len(self.virtual_node_layers):
-            vn_h, h = self.virtual_node_layers[step_idx].forward(g, h, vn_h)
+            h, vn_h = self.virtual_node_layers[step_idx].forward(g=g, h=h, vn_h=vn_h)
+
+        return h, vn_h
+
+    def forward(self, g: dgl.DGLGraph) -> torch.Tensor:
+        r"""
+        Apply the full graph neural network on the input graph and node features.
+
+        Parameters:
+
+            g:
+                graph on which the convolution is done.
+                Must contain the following elements:
+
+                - `g.ndata["h"]`: `torch.Tensor[..., N, Din]`.
+                  Input node feature tensor, before the network.
+                  `N` is the number of nodes, `Din` is the input features
 
-        return vn_h, h
+                - `g.edata["e"]`: `torch.Tensor[..., N, Ein]` **Optional**.
+                  The edge features to use. It will be ignored if the
+                  model doesn't supporte edge features or if
+                  `self.in_dim_edges==0`.
 
-    def forward(self, graph):
+        Returns:
 
-        # Get graph features and apply linear layer
-        h = graph.ndata["h"]
-        e = graph.edata["e"] if self.edge_features else None
-        h = self.in_linear(h)
+            `torch.Tensor[..., M, Dout]` or `torch.Tensor[..., N, Dout]`:
+                Node or graph feature tensor, after the network.
+                `N` is the number of nodes, `M` is the number of graphs,
+                `Dout` is the output dimension ``self.out_dim``
+                If the `self.pooling` is [`None`], then it returns node features and the output dimension is `N`,
+                otherwise it returns graph features and the output dimension is `M`
 
+        """
+
+        # Get node and edge features
+        h = g.ndata["h"]
+        e = g.edata["e"] if (self.in_dim_edges > 0) else None
+
+        # Initialize values of the residuals and virtual node
         h_prev = None
         e_prev = None
         vn_h = 0
+
+        # Apply the forward loop of the layers, residuals and virtual nodes
         for ii, layer in enumerate(self.layers):
             h, e, h_prev, e_prev = self._dgl_layer_forward(
-                layer=layer, graph=graph, h=h, e=e, h_prev=h_prev, e_prev=e_prev, step_idx=ii
+                layer=layer, g=g, h=h, e=e, h_prev=h_prev, e_prev=e_prev, step_idx=ii
             )
-            vn_h, h = self._virtual_node_forward(graph, h, vn_h, step_idx=ii)
+            h, vn_h = self._virtual_node_forward(g=g, h=h, vn_h=vn_h, step_idx=ii)
 
-        pooled_h = self._pool_layer_forward(graph, h)
+        pooled_h = self._pool_layer_forward(g=g, h=h)
 
         return pooled_h
 
@@ -395,10 +760,10 @@ def __repr__(self):
         r"""
         Controls how the class is printed
         """
-        class_str = f"{self.__class__.__name__}(depth={self.depth})"
-        layer_str = f"[{self.layer_name}[{' -> '.join(map(str, self.full_dims))}]"
+        class_str = f"{self.name}(depth={self.depth}, {self.residual_layer})\n    "
+        layer_str = f"{self.layer_class.__name__}[{' -> '.join(map(str, self.full_dims))}]\n    "
         pool_str = f"-> Pooling({self.pooling})"
-        out_str = " -> Linear({self.out_dim})"
+        out_str = f" -> {self.out_linear}"
 
         return class_str + layer_str + pool_str + out_str
 
@@ -406,36 +771,134 @@ def __repr__(self):
 class FullDGLNetwork(nn.Module):
     def __init__(
         self,
-        pre_nn_kwargs,
-        gnn_kwargs,
-        post_nn_kwargs,
-        name="DGL_GNN",
+        gnn_kwargs: Dict[str, Any],
+        pre_nn_kwargs: Union[type(None), Dict[str, Any]] = None,
+        post_nn_kwargs: Union[type(None), Dict[str, Any]] = None,
+        name: str = "DGL_GNN",
     ):
+        r"""
+        Class that allows to implement a full graph neural network architecture,
+        including the pre-processing MLP and the post processing MLP.
+
+        Parameters:
+
+            gnn_kwargs:
+                key-word arguments to use for the initialization of the pre-processing
+                GNN network using the class `FeedForwardDGL`.
+                It must respect the following criteria:
+
+                - gnn_kwargs["in_dim"] must be equal to pre_nn_kwargs["out_dim"]
+                - gnn_kwargs["out_dim"] must be equal to post_nn_kwargs["in_dim"]
+
+            pre_nn_kwargs:
+                key-word arguments to use for the initialization of the pre-processing
+                MLP network of the node features before the GNN, using the class `FeedForwardNN`.
+                If `None`, there won't be a pre-processing MLP.
+
+            post_nn_kwargs:
+                key-word arguments to use for the initialization of the post-processing
+                MLP network after the GNN, using the class `FeedForwardNN`.
+                If `None`, there won't be a post-processing MLP.
+
+            name:
+                Name attributed to the current network, for display and printing
+                purposes.
+        """
 
-        # Initialize the parent nn.Module
         super().__init__()
         self.name = name
 
+        # Initialize the networks
         self.pre_nn, self.post_nn = None, None
         if pre_nn_kwargs is not None:
-            self.pre_nn = FeedForwardNN(**pre_nn_kwargs, name="NN-pre-trans")
-        self.gnn = FeedForwardDGL(**gnn_kwargs, name="main-GNN")
+            name = pre_nn_kwargs.pop("name", "pre-trans-NN")
+            self.pre_nn = FeedForwardNN(**pre_nn_kwargs, name=name)
+        name = gnn_kwargs.pop("name", "main-GNN")
+        self.gnn = FeedForwardDGL(**gnn_kwargs, name=name)
         if post_nn_kwargs is not None:
-            self.post_nn = FeedForwardNN(**post_nn_kwargs, name="NN-post-trans")
+            name = post_nn_kwargs.pop("name", "post-trans-NN")
+            self.post_nn = FeedForwardNN(**post_nn_kwargs, name=name)
 
+        # Initialize the hyper-parameter variable to be accessible by Pytorch Lightning
         hparams_temp = {**self.pre_nn.hparams, **self.pre_nn.hparams, **self.pre_nn.hparams}
         self.hparams = {f"{self.name}.{key}": elem for key, elem in hparams_temp.items()}
 
-    def forward(self, graph):
+    def _check_bad_arguments(self):
+        r"""
+        Raise comprehensive errors if the arguments seem wrong
+        """
         if self.pre_nn is not None:
-            h = graph.ndata["h"]
+            if self.pre_nn["out_dim"] != self.gnn["in_dim"]:
+                raise ValueError(
+                    f"`self.pre_nn.out_dim` must be equal to `self.gnn.in_dim`."
+                    + 'Provided" {self.pre_nn.out_dim} and {self.gnn.in_dim}'
+                )
+
+        if self.post_nn is not None:
+            if self.gnn["out_dim"] != self.post_nn["in_dim"]:
+                raise ValueError(
+                    f"`self.gnn.out_dim` must be equal to `self.post_nn.in_dim`."
+                    + 'Provided" {self.gnn.out_dim} and {self.post_nn.in_dim}'
+                )
+
+    def forward(self, g: dgl.DGLGraph) -> torch.Tensor:
+        r"""
+        Apply the pre-processing neural network, the graph neural network,
+        and the post-processing neural network on the graph features.
+
+        Parameters:
+
+            g:
+                graph on which the convolution is done.
+                Must contain the following elements:
+
+                - `g.ndata["h"]`: `torch.Tensor[..., N, Din]`.
+                  Input node feature tensor, before the network.
+                  `N` is the number of nodes, `Din` is the input features dimension ``self.pre_nn.in_dim``
+
+                - `g.edata["e"]`: `torch.Tensor[..., N, Ein]` **Optional**.
+                  The edge features to use. It will be ignored if the
+                  model doesn't supporte edge features or if
+                  `self.in_dim_edges==0`.
+
+        Returns:
+
+            `torch.Tensor[..., M, Dout]` or `torch.Tensor[..., N, Dout]`:
+                Node or graph feature tensor, after the network.
+                `N` is the number of nodes, `M` is the number of graphs,
+                `Dout` is the output dimension ``self.post_nn.out_dim``
+                If the `self.gnn.pooling` is [`None`], then it returns node features and the output dimension is `N`,
+                otherwise it returns graph features and the output dimension is `M`
+
+        """
+
+        if self.pre_nn is not None:
+            h = g.ndata["h"]
             h = self.pre_nn.forward(h)
-            graph.ndata["h"] = h
-        h = self.gnn.forward(graph)
+            g.ndata["h"] = h
+        h = self.gnn.forward(g)
         if self.post_nn is not None:
             h = self.post_nn.forward(h)
         return h
 
+    def __repr__(self):
+        r"""
+        Controls how the class is printed
+        """
+        pre_nn_str, post_nn_str = "", ""
+        if self.pre_nn is not None:
+            pre_nn_str = self.pre_nn.__repr__() + "\n\n"
+        gnn_str = self.gnn.__repr__() + "\n\n"
+        if self.post_nn is not None:
+            post_nn_str = self.post_nn.__repr__()
+
+        child_str = "    " + pre_nn_str + gnn_str + post_nn_str
+        child_str = "    ".join(child_str.splitlines(True))
+
+        full_str = self.name + "\n" + "-" * (len(self.name) + 2) + "\n" + child_str
+
+        return full_str
+
 
 class FullDGLSiameseNetwork(FullDGLNetwork):
     def __init__(self, pre_nn_kwargs, gnn_kwargs, post_nn_kwargs, dist_method, name="Siamese_DGL_GNN"):
diff --git a/goli/dgl/base_layers.py b/goli/dgl/base_layers.py
index 1ecd0f1ae..ce9ff8050 100644
--- a/goli/dgl/base_layers.py
+++ b/goli/dgl/base_layers.py
@@ -1,16 +1,24 @@
-import abc
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
+from typing import Union, Callable
 
 SUPPORTED_ACTIVATION_MAP = {"ReLU", "Sigmoid", "Tanh", "ELU", "SELU", "GLU", "LeakyReLU", "Softplus", "None"}
-EPS = 1e-5
 
 
-def get_activation(activation):
-    r""" returns the activation function represented by the input string """
-    if activation and callable(activation):
+def get_activation(activation: Union[type(None), str, Callable]) -> Union[type(None), Callable]:
+    r"""
+    returns the activation function represented by the input string
+
+    Parameters:
+        activation: Callable, `None`, or string with value:
+            "none", "ReLU", "Sigmoid", "Tanh", "ELU", "SELU", "GLU", "LeakyReLU", "Softplus"
+
+    Returns:
+        Callable or None: The activation function
+    """
+    if (activation is not None) and callable(activation):
         # activation is already a function
         return activation
 
@@ -30,11 +38,11 @@ def __init__(
         self,
         in_dim: int,
         out_dim: int,
-        activation="relu",
-        dropout=0.0,
-        batch_norm=False,
-        bias=True,
-        init_fn=None,
+        activation: Union[str, Callable] = "relu",
+        dropout: float = 0.0,
+        batch_norm: bool = False,
+        bias: bool = True,
+        init_fn: Union[type(None), Callable] = None,
     ):
 
         r"""
@@ -47,36 +55,36 @@ def __init__(
         - Batch Normalization (if applicable)
 
         Parameters:
-            in_dim: int
+            in_dim:
                 Input dimension of the layer (the torch.nn.Linear)
-            out_dim: int
+            out_dim:
                 Output dimension of the layer.
             dropout:
                 The ratio of units to dropout. No dropout by default.
-            activation: str or callable
+            activation:
                 Activation function to use.
-            batch_norm: bool
+            batch_norm:
                 Whether to use batch normalization
-            bias: bool
+            bias:
                 Whether to enable bias in for the linear layer.
-            init_fn: callable
+            init_fn:
                 Initialization function to use for the weight of the layer. Default is
                 $$\mathcal{U}(-\sqrt{k}, \sqrt{k})$$ with $$k=\frac{1}{ \text{in_dim}}$$
 
         Attributes:
-            dropout: int
+            dropout (int):
                 The ratio of units to dropout.
-            batch_norm: int
+            batch_norm (int):
                 Whether to use batch normalization
-            linear: torch.nn.Linear
+            linear (torch.nn.Linear):
                 The linear layer
-            activation: the torch.nn.Module
+            activation (torch.nn.Module):
                 The activation layer
-            init_fn: function
+            init_fn (Callable):
                 Initialization function used for the weight of the layer
-            in_dim: int
+            in_dim (int):
                 Input dimension of the linear layer
-            out_dim: int
+            out_dim (int):
                 Output dimension of the linear layer
         """
 
@@ -107,8 +115,25 @@ def reset_parameters(self, init_fn=None):
         if self.bias:
             self.linear.bias.data.zero_()
 
-    def forward(self, x):
-        h = self.linear(x)
+    def forward(self, h: torch.Tensor) -> torch.Tensor:
+        r"""
+        Apply the FC layer on the input features.
+
+        Parameters:
+
+            h: `torch.Tensor[..., Din]`:
+                Input feature tensor, before the FC.
+                `Din` is the number of input features
+
+        Returns:
+
+            `torch.Tensor[..., Dout]`:
+                Output feature tensor, after the FC.
+                `Dout` is the number of output features
+
+        """
+
+        h = self.linear(h)
 
         if self.batch_norm is not None:
             if h.shape[1] != self.out_dim:
@@ -123,26 +148,50 @@ def forward(self, x):
         return h
 
     def __repr__(self):
-        return f"{self.__class__.__name__}(in_dim={self.in_dim}, out_dim={self.out_dim}, activation={self.activation})"
+        return f"{self.__class__.__name__}({self.in_dim} -> {self.out_dim}, activation={self.activation})"
 
 
 class MLP(nn.Module):
-    r"""
-    Simple multi-layer perceptron, built of a series of FCLayers
-    """
-
     def __init__(
         self,
-        in_dim,
-        hidden_dim,
-        out_dim,
-        layers,
-        mid_activation="relu",
-        last_activation="none",
+        in_dim: int,
+        hidden_dim: int,
+        out_dim: int,
+        layers: int,
+        actibation: Union[str, Callable] = "relu",
+        last_activation: Union[str, Callable] = "none",
         dropout=0.0,
-        mid_batch_norm=False,
+        batch_norm=False,
         last_batch_norm=False,
     ):
+        r"""
+        Simple multi-layer perceptron, built of a series of FCLayers
+
+        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
+            out_dim:
+                Output dimension of the MLP.
+            layers:
+                Number of hidden layers
+            activation:
+                Activation function to use in all the layers except the last.
+                if `layers==1`, this parameter is ignored
+            last_activation:
+                Activation function to use in the last layer.
+            dropout:
+                The ratio of units to dropout. Must be between 0 and 1
+            batch_norm:
+                Whether to use batch normalization in the hidden layers.
+                if `layers==1`, this parameter is ignored
+            last_batch_norm:
+                Whether to use batch normalization in the last layer
+
+        """
+
         super().__init__()
 
         self.in_dim = in_dim
@@ -165,8 +214,8 @@ def __init__(
                 FCLayer(
                     in_dim,
                     hidden_dim,
-                    activation=mid_activation,
-                    batch_norm=mid_batch_norm,
+                    activation=actibation,
+                    batch_norm=batch_norm,
                     dropout=dropout,
                 )
             )
@@ -175,8 +224,8 @@ def __init__(
                     FCLayer(
                         hidden_dim,
                         hidden_dim,
-                        activation=mid_activation,
-                        batch_norm=mid_batch_norm,
+                        activation=actibation,
+                        batch_norm=batch_norm,
                         dropout=dropout,
                     )
                 )
@@ -190,41 +239,72 @@ def __init__(
                 )
             )
 
-    def forward(self, x):
+    def forward(self, h: torch.Tensor) -> torch.Tensor:
+        r"""
+        Apply the MLP on the input features.
+
+        Parameters:
+
+            h: `torch.Tensor[..., Din]`:
+                Input feature tensor, before the MLP.
+                `Din` is the number of input features
+
+        Returns:
+
+            `torch.Tensor[..., Dout]`:
+                Output feature tensor, after the MLP.
+                `Dout` is the number of output features
+
+        """
         for fc in self.fully_connected:
-            x = fc(x)
-        return x
+            h = fc(h)
+        return h
 
     def __repr__(self):
+        r"""
+        Controls how the class is printed
+        """
         return self.__class__.__name__ + " (" + str(self.in_dim) + " -> " + str(self.out_dim) + ")"
 
 
 class GRU(nn.Module):
-    r"""
-    Wrapper class for the GRU used by the GNN framework, nn.GRU is used for the Gated Recurrent Unit itself
-    """
+    def __init__(self, in_dim: int, hidden_dim: int):
+        r"""
+        Wrapper class for the GRU used by the GNN framework, nn.GRU is used for the Gated Recurrent Unit itself
+
+        Parameters:
+            in_dim:
+                Input dimension of the GRU layer
+            hidden_dim:
+                Hidden dimension of the GRU layer.
+        """
 
-    def __init__(self, input_size, hidden_dim, device):
         super().__init__()
-        self.input_size = input_size
+        self.in_dim = in_dim
         self.hidden_dim = hidden_dim
-        self.gru = nn.GRU(input_size=input_size, hidden_dim=hidden_dim).to(device)
+        self.gru = nn.GRU(in_dim=in_dim, hidden_dim=hidden_dim)
 
     def forward(self, x, y):
         r"""
-        :param x:   shape: (B, N, Din) where Din <= input_size (difference is padded)
-        :param y:   shape: (B, N, Dh) where Dh <= hidden_dim (difference is padded)
-        :return:    shape: (B, N, Dh)
+        Parameters:
+            x:  `torch.Tensor[B, N, Din]`
+                where Din <= in_dim (difference is padded)
+            y:  `torch.Tensor[B, N, Dh]`
+                where Dh <= hidden_dim (difference is padded)
+
+        Returns:
+            torch.Tensor: `torch.Tensor[B, N, Dh]`
+
         """
-        assert x.shape[-1] <= self.input_size and y.shape[-1] <= self.hidden_dim
+        assert x.shape[-1] <= self.in_dim and y.shape[-1] <= self.hidden_dim
 
         (B, N, _) = x.shape
         x = x.reshape(1, B * N, -1).contiguous()
         y = y.reshape(1, B * N, -1).contiguous()
 
         # padding if necessary
-        if x.shape[-1] < self.input_size:
-            x = F.pad(input=x, pad=[0, self.input_size - x.shape[-1]], mode="constant", value=0)
+        if x.shape[-1] < self.in_dim:
+            x = F.pad(input=x, pad=[0, self.in_dim - x.shape[-1]], mode="constant", value=0)
         if y.shape[-1] < self.hidden_dim:
             y = F.pad(input=y, pad=[0, self.hidden_dim - y.shape[-1]], mode="constant", value=0)
 
diff --git a/goli/dgl/dgl_layers/__init__.py b/goli/dgl/dgl_layers/__init__.py
index e878218a0..5753272ea 100644
--- a/goli/dgl/dgl_layers/__init__.py
+++ b/goli/dgl/dgl_layers/__init__.py
@@ -1,3 +1,4 @@
+from goli.dgl.dgl_layers.base_dgl_layer import BaseDGLLayer
 from goli.dgl.dgl_layers.gat_layer import GATLayer
 from goli.dgl.dgl_layers.gcn_layer import GCNLayer
 from goli.dgl.dgl_layers.gin_layer import GINLayer
diff --git a/goli/dgl/dgl_layers/base_dgl_layer.py b/goli/dgl/dgl_layers/base_dgl_layer.py
index 1eb297de0..ccc6519d9 100644
--- a/goli/dgl/dgl_layers/base_dgl_layer.py
+++ b/goli/dgl/dgl_layers/base_dgl_layer.py
@@ -1,6 +1,7 @@
+import torch
 import torch.nn as nn
 import abc
-from typing import List, Dict, Tuple
+from typing import List, Dict, Tuple, Union, Callable
 
 from goli.dgl.base_layers import get_activation
 from goli.commons.decorators import classproperty
@@ -8,7 +9,12 @@
 
 class BaseDGLLayer(nn.Module):
     def __init__(
-        self, in_dim: int, out_dim: int, activation="relu", dropout: float = 0.0, batch_norm: bool = False
+        self,
+        in_dim: int,
+        out_dim: int,
+        activation: Union[str, Callable] = "relu",
+        dropout: float = 0.0,
+        batch_norm: bool = False,
     ):
         r"""
         Abstract class used to standardize the implementation of DGL layers
@@ -18,19 +24,19 @@ def __init__(
 
         Parameters:
 
-            in_dim: int
+            in_dim:
                 Input feature dimensions of the layer
 
-            out_dim: int
+            out_dim:
                 Output feature dimensions of the layer
 
-            activation: str, Callable
+            activation:
                 activation function to use in the layer
 
-            dropout: float
+            dropout:
                 The ratio of units to dropout. Must be between 0 and 1
 
-            batch_norm: bool
+            batch_norm:
                 Whether to use batch normalization
         """
 
@@ -52,10 +58,10 @@ def __init__(
 
         self.batch_norm_layer = None
         if batch_norm:
-            self.batch_norm_layer = nn.BatchNorm1d(out_dim)
+            self.batch_norm_layer = nn.BatchNorm1d(out_dim * self.out_dim_factor)
 
     def apply_norm_activation_dropout(
-        self, h, batch_norm: bool = True, activation: bool = True, dropout: bool = True
+        self, h: torch.Tensor, batch_norm: bool = True, activation: bool = True, dropout: bool = True
     ):
         r"""
         Apply the different normalization and the dropout to the
@@ -63,21 +69,21 @@ def apply_norm_activation_dropout(
 
         Parameters:
 
-            h: torch.Tensor()
+            h:
                 Feature tensor, to be normalized
 
-            batch_norm: bool
+            batch_norm:
                 Whether to apply the batch_norm layer
 
-            activation: bool
+            activation:
                 Whether to apply the activation layer
 
-            dropout: bool
+            dropout:
                 Whether to apply the dropout layer
 
         Returns:
 
-            h: torch.Tensor()
+            h:
                 Normalized and dropped-out features
 
         """
@@ -95,13 +101,13 @@ def apply_norm_activation_dropout(
 
     @classproperty
     def layer_supports_edges(cls) -> bool:
-        """
+        r"""
         Abstract method. Return a boolean specifying if the layer type
         supports output edges edges or not.
 
         Returns:
 
-            supports_edges: bool
+            bool:
                 Whether the layer supports the use of edges
         """
         ...
@@ -109,7 +115,7 @@ def layer_supports_edges(cls) -> bool:
     @property
     @abc.abstractmethod
     def layer_inputs_edges(self) -> bool:
-        """
+        r"""
         Abstract method. Return a boolean specifying if the layer type
         uses edges as input or not.
         It is different from ``layer_supports_input_edges`` since a layer that
@@ -117,7 +123,7 @@ def layer_inputs_edges(self) -> bool:
 
         Returns:
 
-            uses_edges: bool
+            bool:
                 Whether the layer uses input edges in the forward pass
         """
         ...
@@ -125,7 +131,7 @@ def layer_inputs_edges(self) -> bool:
     @property
     @abc.abstractmethod
     def layer_outputs_edges(self) -> bool:
-        """
+        r"""
         Abstract method. Return a boolean specifying if the layer type
         uses edges as input or not.
         It is different from ``layer_supports_output_edges`` since a layer that
@@ -133,7 +139,7 @@ def layer_outputs_edges(self) -> bool:
 
         Returns:
 
-            uses_edges: bool
+            bool:
                 Whether the layer outputs edges in the forward pass
         """
         ...
@@ -141,7 +147,7 @@ def layer_outputs_edges(self) -> bool:
     @property
     @abc.abstractmethod
     def out_dim_factor(self) -> int:
-        """
+        r"""
         Abstract method.
         Get the factor by which the output dimension is multiplied for
         the next layer.
@@ -155,8 +161,8 @@ def out_dim_factor(self) -> int:
 
         Returns:
 
-            dim_factor: int
-                The factor that multiplies the dimensions
+            int:
+                The factor that multiplies the output dimensions
         """
         ...
 
@@ -164,4 +170,6 @@ def __repr__(self):
         r"""
         Controls how the class is printed
         """
-        return f"{self.__class__.__name__}({self.in_dim} -> {self.out_dim}, activation={self.activation})"
+        f = self.out_dim_factor
+        out_dim_f_print = "" if f == 1 else f" * {f}"
+        return f"{self.__class__.__name__}({self.in_dim} -> {self.out_dim}{out_dim_f_print}, activation={self.activation})"
diff --git a/goli/dgl/dgl_layers/gat_layer.py b/goli/dgl/dgl_layers/gat_layer.py
index b3f4a2f47..b16b91456 100644
--- a/goli/dgl/dgl_layers/gat_layer.py
+++ b/goli/dgl/dgl_layers/gat_layer.py
@@ -53,6 +53,8 @@ def __init__(
                 Whether to use batch normalization
         """
 
+        self.num_heads = num_heads
+
         super().__init__(
             in_dim=in_dim,
             out_dim=out_dim,
@@ -61,8 +63,6 @@ def __init__(
             batch_norm=batch_norm,
         )
 
-        self.num_heads = num_heads
-
         self.gatconv = GATConv(
             in_feats=self.in_dim,
             out_feats=self.out_dim,
@@ -73,7 +73,7 @@ def __init__(
         )
 
     def forward(self, g: DGLGraph, h: torch.Tensor) -> torch.Tensor:
-        """
+        r"""
         Apply the graph convolutional layer, with the specified activations,
         normalizations and dropout.
 
@@ -82,13 +82,13 @@ def forward(self, g: DGLGraph, h: torch.Tensor) -> torch.Tensor:
             g: dgl.DGLGraph
                 graph on which the convolution is done
 
-            h: torch.Tensor(..., N, Din)
+            h: `torch.Tensor[..., N, Din]`
                 Node feature tensor, before convolution.
                 N is the number of nodes, Din is the input dimension ``self.in_dim``
 
         Returns:
 
-            h: torch.Tensor(..., N, Dout)
+            `torch.Tensor[..., N, Dout]`:
                 Node feature tensor, after convolution.
                 N is the number of nodes, Dout is the output dimension ``self.out_dim``
 
@@ -101,7 +101,7 @@ def forward(self, g: DGLGraph, h: torch.Tensor) -> torch.Tensor:
 
     @classproperty
     def layer_supports_edges(cls) -> bool:
-        """
+        r"""
         Return a boolean specifying if the layer type supports edges or not.
 
         Returns:
@@ -113,7 +113,7 @@ def layer_supports_edges(cls) -> bool:
 
     @property
     def layer_inputs_edges(self) -> bool:
-        """
+        r"""
         Return a boolean specifying if the layer type
         uses edges as input or not.
         It is different from ``layer_supports_edges`` since a layer that
@@ -128,7 +128,7 @@ def layer_inputs_edges(self) -> bool:
 
     @property
     def layer_outputs_edges(self) -> bool:
-        """
+        r"""
         Abstract method. Return a boolean specifying if the layer type
         uses edges as input or not.
         It is different from ``layer_supports_edges`` since a layer that
@@ -143,7 +143,7 @@ def layer_outputs_edges(self) -> bool:
 
     @property
     def out_dim_factor(self) -> int:
-        """
+        r"""
         Get the factor by which the output dimension is multiplied for
         the next layer.
 
diff --git a/goli/dgl/dgl_layers/gated_gcn_layer.py b/goli/dgl/dgl_layers/gated_gcn_layer.py
index b44ffccff..52ff53e8f 100644
--- a/goli/dgl/dgl_layers/gated_gcn_layer.py
+++ b/goli/dgl/dgl_layers/gated_gcn_layer.py
@@ -2,7 +2,7 @@
 import torch.nn as nn
 import torch.nn.functional as F
 from copy import deepcopy
-from typing import Tuple
+from typing import Tuple, Union, Callable
 from dgl import DGLGraph
 
 from goli.dgl.dgl_layers.base_dgl_layer import BaseDGLLayer
@@ -22,7 +22,7 @@ def __init__(
         out_dim: int,
         in_dim_edges: int,
         out_dim_edges: int,
-        activation="relu",
+        activation: Union[Callable, str] = "relu",
         dropout: float = 0.0,
         batch_norm: bool = False,
     ):
@@ -33,22 +33,22 @@ def __init__(
 
         Parameters:
 
-            in_dim: int
+            in_dim:
                 Input feature dimensions of the layer
 
-            out_dim: int
+            out_dim:
                 Output feature dimensions of the layer, and for the edges
 
-            in_dim_edges: int
+            in_dim_edges:
                 Input edge-feature dimensions of the layer
 
-            activation: str, Callable
+            activation:
                 activation function to use in the layer
 
-            dropout: float
+            dropout:
                 The ratio of units to dropout. Must be between 0 and 1
 
-            batch_norm: bool
+            batch_norm:
                 Whether to use batch normalization
 
         """
@@ -85,30 +85,29 @@ def reduce_func(self, nodes):
         return {"h": h}
 
     def forward(self, g: DGLGraph, h: torch.Tensor, e: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
-        """
+        r"""
         Apply the graph convolutional layer, with the specified activations,
         normalizations and dropout.
 
         Parameters:
 
-            g: dgl.DGLGraph
+            g:
                 graph on which the convolution is done
 
-            h: torch.Tensor(..., N, Din)
+            h: `torch.Tensor[..., N, Din]`
                 Node feature tensor, before convolution.
                 N is the number of nodes, Din is the input dimension ``self.in_dim``
 
-            e: torch.Tensor(..., N, Din_edges)
+            e: `torch.Tensor[..., N, Din_edges]`
                 Edge feature tensor, before convolution.
                 N is the number of nodes, Din is the input edge dimension  ``self.in_dim_edges``
 
         Returns:
-
-            h: torch.Tensor(..., N, Dout)
+            `torch.Tensor[..., N, Dout]`:
                 Node feature tensor, after convolution.
                 N is the number of nodes, Dout is the output dimension ``self.out_dim``
 
-            e: torch.Tensor(..., N, Dout)
+            `torch.Tensor[..., N, Dout]`:
                 Edge feature tensor, after convolution.
                 N is the number of nodes, Dout_edges is the output edge dimension ``self.out_dim``
 
@@ -132,19 +131,19 @@ def forward(self, g: DGLGraph, h: torch.Tensor, e: torch.Tensor) -> Tuple[torch.
 
     @classproperty
     def layer_supports_edges(cls) -> bool:
-        """
+        r"""
         Return a boolean specifying if the layer type supports edges or not.
 
         Returns:
 
-            supports_edges: bool
+            bool:
                 Always ``True`` for the current class
         """
         return True
 
     @property
     def layer_inputs_edges(self) -> bool:
-        """
+        r"""
         Return a boolean specifying if the layer type
         uses edges as input or not.
         It is different from ``layer_supports_edges`` since a layer that
@@ -152,14 +151,14 @@ def layer_inputs_edges(self) -> bool:
 
         Returns:
 
-            uses_edges: bool
+            bool:
                 Always ``True`` for the current class
         """
         return True
 
     @property
     def layer_outputs_edges(self) -> bool:
-        """
+        r"""
         Abstract method. Return a boolean specifying if the layer type
         uses edges as input or not.
         It is different from ``layer_supports_edges`` since a layer that
@@ -167,14 +166,14 @@ def layer_outputs_edges(self) -> bool:
 
         Returns:
 
-            uses_edges: bool
+            bool:
                 Always ``True`` for the current class
         """
         return True
 
     @property
     def out_dim_factor(self) -> int:
-        """
+        r"""
         Get the factor by which the output dimension is multiplied for
         the next layer.
 
@@ -187,7 +186,7 @@ def out_dim_factor(self) -> int:
 
         Returns:
 
-            dim_factor: int
+            int:
                 Always ``1`` for the current class
         """
         return 1
diff --git a/goli/dgl/dgl_layers/gcn_layer.py b/goli/dgl/dgl_layers/gcn_layer.py
index 11df1acfe..424ecb18d 100644
--- a/goli/dgl/dgl_layers/gcn_layer.py
+++ b/goli/dgl/dgl_layers/gcn_layer.py
@@ -1,7 +1,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from typing import List, Dict, Tuple
+from typing import List, Dict, Tuple, Union, Callable
 
 import dgl.function as fn
 from dgl.nn.pytorch import GraphConv
@@ -19,39 +19,44 @@
 
 class GCNLayer(BaseDGLLayer):
     def __init__(
-        self, in_dim: int, out_dim: int, activation="relu", dropout: float = 0.0, batch_norm: bool = False
+        self,
+        in_dim: int,
+        out_dim: int,
+        activation: Union[str, Callable] = "relu",
+        dropout: float = 0.0,
+        batch_norm: bool = False,
     ):
-
-        super().__init__(
-            in_dim=in_dim,
-            out_dim=out_dim,
-            activation=activation,
-            dropout=dropout,
-            batch_norm=batch_norm,
-        )
-        """
+        r"""
         Graph convolutional network (GCN) layer from
         Thomas N. Kipf, Max Welling, Semi-Supervised Classification with Graph Convolutional Networks (ICLR 2017)
         http://arxiv.org/abs/1609.02907
 
         Parameters:
 
-            in_dim: int
+            in_dim:
                 Input feature dimensions of the layer
 
-            out_dim: int
+            out_dim:
                 Output feature dimensions of the layer
 
-            activation: str, Callable
+            activation:
                 activation function to use in the layer
 
-            dropout: float
+            dropout:
                 The ratio of units to dropout. Must be between 0 and 1
 
-            batch_norm: bool
+            batch_norm:
                 Whether to use batch normalization
         """
 
+        super().__init__(
+            in_dim=in_dim,
+            out_dim=out_dim,
+            activation=activation,
+            dropout=dropout,
+            batch_norm=batch_norm,
+        )
+
         self.conv = GraphConv(
             in_feats=in_dim,
             out_feats=out_dim,
@@ -63,22 +68,22 @@ def __init__(
         )
 
     def forward(self, g: DGLGraph, h: torch.Tensor) -> torch.Tensor:
-        """
+        r"""
         Apply the graph convolutional layer, with the specified activations,
         normalizations and dropout.
 
         Parameters:
 
-            g: dgl.DGLGraph
+            g:
                 graph on which the convolution is done
 
-            h: torch.Tensor(..., N, Din)
+            h: `torch.Tensor[..., N, Din]`
                 Node feature tensor, before convolution.
                 N is the number of nodes, Din is the input dimension ``self.in_dim``
 
         Returns:
 
-            h: torch.Tensor(..., N, Dout)
+            `torch.Tensor[..., N, Dout]`:
                 Node feature tensor, after convolution.
                 N is the number of nodes, Dout is the output dimension ``self.out_dim``
 
@@ -91,19 +96,19 @@ def forward(self, g: DGLGraph, h: torch.Tensor) -> torch.Tensor:
 
     @classproperty
     def layer_supports_edges(cls) -> bool:
-        """
+        r"""
         Return a boolean specifying if the layer type supports edges or not.
 
         Returns:
 
-            supports_edges: bool
+            bool
                 Always ``False`` for the current class
         """
         return False
 
     @property
     def layer_inputs_edges(self) -> bool:
-        """
+        r"""
         Return a boolean specifying if the layer type
         uses edges as input or not.
         It is different from ``layer_supports_edges`` since a layer that
@@ -111,14 +116,14 @@ def layer_inputs_edges(self) -> bool:
 
         Returns:
 
-            uses_edges: bool
+            bool:
                 Always ``False`` for the current class
         """
         return False
 
     @property
     def layer_outputs_edges(self) -> bool:
-        """
+        r"""
         Abstract method. Return a boolean specifying if the layer type
         uses edges as input or not.
         It is different from ``layer_supports_edges`` since a layer that
@@ -126,14 +131,14 @@ def layer_outputs_edges(self) -> bool:
 
         Returns:
 
-            uses_edges: bool
+            bool:
                 Always ``False`` for the current class
         """
         return False
 
     @property
     def out_dim_factor(self) -> int:
-        """
+        r"""
         Get the factor by which the output dimension is multiplied for
         the next layer.
 
@@ -146,7 +151,7 @@ def out_dim_factor(self) -> int:
 
         Returns:
 
-            dim_factor: int
+            int:
                 Always ``1`` for the current class
         """
         return 1
diff --git a/goli/dgl/dgl_layers/gin_layer.py b/goli/dgl/dgl_layers/gin_layer.py
index 2203a7324..b69b3e847 100644
--- a/goli/dgl/dgl_layers/gin_layer.py
+++ b/goli/dgl/dgl_layers/gin_layer.py
@@ -3,6 +3,7 @@
 import torch.nn.functional as F
 import dgl.function as fn
 from dgl import DGLGraph
+from typing import Callable, Union
 
 from goli.dgl.dgl_layers.base_dgl_layer import BaseDGLLayer
 from goli.dgl.base_layers import MLP
@@ -20,11 +21,11 @@ def __init__(
         self,
         in_dim: int,
         out_dim: int,
-        activation="relu",
+        activation: Union[Callable, str] = "relu",
         dropout: float = 0.0,
         batch_norm: bool = False,
-        init_eps=0.0,
-        learn_eps=True,
+        init_eps: float = 0.0,
+        learn_eps: bool = True,
     ):
         r"""
         GIN: Graph Isomorphism Networks
@@ -35,25 +36,25 @@ def __init__(
 
         Parameters:
 
-            in_dim: int
+            in_dim:
                 Input feature dimensions of the layer
 
-            out_dim: int
+            out_dim:
                 Output feature dimensions of the layer
 
-            activation: str, Callable
+            activation:
                 activation function to use in the layer
 
-            dropout: float
+            dropout:
                 The ratio of units to dropout. Must be between 0 and 1
 
-            batch_norm: bool
+            batch_norm:
                 Whether to use batch normalization
 
-            init_eps : float
+            init_eps :
                 Initial :math:`\epsilon` value, default: ``0``.
 
-            learn_eps : bool
+            learn_eps :
                 If True, :math:`\epsilon` will be a learnable parameter.
 
         """
@@ -80,9 +81,9 @@ def __init__(
             hidden_dim=self.in_dim,
             out_dim=self.out_dim,
             layers=2,
-            mid_activation=self.activation_layer,
+            actibation=self.activation_layer,
             last_activation="none",
-            mid_batch_norm=self.batch_norm,
+            batch_norm=self.batch_norm,
             last_batch_norm=False,
         )
 
@@ -98,22 +99,22 @@ def message_func(self, g):
         return func
 
     def forward(self, g: DGLGraph, h: torch.Tensor) -> torch.Tensor:
-        """
+        r"""
         Apply the GIN convolutional layer, with the specified activations,
         normalizations and dropout.
 
         Parameters:
 
-            g: dgl.DGLGraph
+            g:
                 graph on which the convolution is done
 
-            h: torch.Tensor(..., N, Din)
+            h: `torch.Tensor[..., N, Din]`
                 Node feature tensor, before convolution.
                 N is the number of nodes, Din is the input dimension ``self.in_dim``
 
         Returns:
 
-            h: torch.Tensor(..., N, Dout)
+            `torch.Tensor[..., N, Dout]`:
                 Node feature tensor, after convolution.
                 N is the number of nodes, Dout is the output dimension ``self.out_dim``
 
@@ -134,7 +135,7 @@ def forward(self, g: DGLGraph, h: torch.Tensor) -> torch.Tensor:
 
     @classproperty
     def layer_supports_edges(cls) -> bool:
-        """
+        r"""
         Return a boolean specifying if the layer type supports edges or not.
 
         Returns:
@@ -146,7 +147,7 @@ def layer_supports_edges(cls) -> bool:
 
     @property
     def layer_inputs_edges(self) -> bool:
-        """
+        r"""
         Return a boolean specifying if the layer type
         uses edges as input or not.
         It is different from ``layer_supports_edges`` since a layer that
@@ -154,14 +155,14 @@ def layer_inputs_edges(self) -> bool:
 
         Returns:
 
-            uses_edges: bool
+            bool:
                 Always ``False`` for the current class
         """
         return False
 
     @property
     def layer_outputs_edges(self) -> bool:
-        """
+        r"""
         Abstract method. Return a boolean specifying if the layer type
         uses edges as input or not.
         It is different from ``layer_supports_edges`` since a layer that
@@ -169,14 +170,14 @@ def layer_outputs_edges(self) -> bool:
 
         Returns:
 
-            uses_edges: bool
+            bool:
                 Always ``False`` for the current class
         """
         return False
 
     @property
     def out_dim_factor(self) -> int:
-        """
+        r"""
         Get the factor by which the output dimension is multiplied for
         the next layer.
 
@@ -189,7 +190,7 @@ def out_dim_factor(self) -> int:
 
         Returns:
 
-            dim_factor: int
+            int:
                 Always ``1`` for the current class
         """
         return 1
diff --git a/goli/dgl/dgl_layers/pna_layer.py b/goli/dgl/dgl_layers/pna_layer.py
index 35bd4a041..e464ac746 100644
--- a/goli/dgl/dgl_layers/pna_layer.py
+++ b/goli/dgl/dgl_layers/pna_layer.py
@@ -5,7 +5,7 @@
 import torch.nn.functional as F
 import dgl.function as fn
 from dgl import DGLGraph
-from typing import Dict, List, Tuple
+from typing import Dict, List, Tuple, Union, Callable
 
 from goli.dgl.pna_operations import PNA_AGGREGATORS, PNA_SCALERS
 from goli.dgl.base_layers import MLP, FCLayer, get_activation
@@ -26,11 +26,11 @@ def __init__(
         out_dim: int,
         aggregators: List[str],
         scalers: List[str],
-        activation="relu",
+        activation: Union[Callable, str] = "relu",
         dropout: float = 0.0,
         batch_norm: bool = False,
         avg_d: float = 1.0,
-        last_activation="none",
+        last_activation: Union[Callable, str] = "none",
     ):
         r"""
         Abstract class used to standardize the implementation of PNA layers
@@ -44,35 +44,35 @@ def __init__(
 
         Parameters:
 
-            in_dim: int
+            in_dim:
                 Input feature dimensions of the layer
 
-            out_dim: int
+            out_dim:
                 Output feature dimensions of the layer
 
-            aggregators: list(str)
+            aggregators:
                 Set of aggregation function identifiers,
                 e.g. "mean", "max", "min", "std", "sum", "var", "moment3".
                 The results from all aggregators will be concatenated.
 
-            scalers: list(str)
+            scalers:
                 Set of scaling functions identifiers
                 e.g. "identidy", "amplification", "attenuation"
                 The results from all scalers will be concatenated
 
-            activation: str, Callable
+            activation:
                 activation function to use in the layer
 
-            dropout: float
+            dropout:
                 The ratio of units to dropout. Must be between 0 and 1
 
-            batch_norm: bool
+            batch_norm:
                 Whether to use batch normalization
 
-            avg_d: float
+            avg_d:
                 Average degree of nodes in the training set, used by scalers to normalize
 
-            last_activation: str, Callable
+            last_activation:
                 activation function to use in the last layer of the internal MLP
 
         """
@@ -94,13 +94,13 @@ def __init__(
         self.scalers = [PNA_SCALERS[scale] for scale in scalers]
 
     def message_func(self, edges) -> Dict[str, torch.Tensor]:
-        """
+        r"""
         The message function to generate messages along the edges.
         """
         return {"e": edges.data["e"]}
 
     def reduce_func(self, nodes) -> Dict[str, torch.Tensor]:
-        """
+        r"""
         The reduce function to aggregate the messages.
         Apply the aggregators and scalers, and concatenate the results.
         """
@@ -108,16 +108,16 @@ def reduce_func(self, nodes) -> Dict[str, torch.Tensor]:
         h = nodes.mailbox["e"]
         D = h.shape[-2]
         h_to_cat = [aggr(h=h, h_in=h_in) for aggr in self.aggregators]
-        h = torch.cat(h_to_cat, dim=1)
+        h = torch.cat(h_to_cat, dim=-1)
 
         if len(self.scalers) > 1:
-            h = torch.cat([scale(h, D=D, avg_d=self.avg_d) for scale in self.scalers], dim=1)
+            h = torch.cat([scale(h, D=D, avg_d=self.avg_d) for scale in self.scalers], dim=-1)
 
         return {"h": h}
 
     @property
     def layer_outputs_edges(self) -> bool:
-        """
+        r"""
         Abstract method. Return a boolean specifying if the layer type
         uses edges as input or not.
         It is different from ``layer_supports_edges`` since a layer that
@@ -125,14 +125,14 @@ def layer_outputs_edges(self) -> bool:
 
         Returns:
 
-            uses_edges: bool
+            bool:
                 Always ``False`` for the current class
         """
         return False
 
     @property
     def out_dim_factor(self) -> int:
-        """
+        r"""
         Get the factor by which the output dimension is multiplied for
         the next layer.
 
@@ -145,7 +145,7 @@ def out_dim_factor(self) -> int:
 
         Returns:
 
-            dim_factor: int
+            int:
                 Always ``1`` for the current class
         """
         return 1
@@ -158,11 +158,11 @@ def __init__(
         out_dim: int,
         aggregators: List[str],
         scalers: List[str],
-        activation="relu",
+        activation: Union[Callable, str] = "relu",
         dropout: float = 0.0,
         batch_norm: bool = False,
         avg_d: Dict[str, float] = {"log": 1.0},
-        last_activation="none",
+        last_activation: Union[Callable, str] = "none",
         posttrans_layers: int = 1,
     ):
         r"""
@@ -178,38 +178,38 @@ def __init__(
 
         Parameters:
 
-            in_dim: int
+            in_dim:
                 Input feature dimensions of the layer
 
-            out_dim: int
+            out_dim:
                 Output feature dimensions of the layer
 
-            aggregators: list(str)
+            aggregators:
                 Set of aggregation function identifiers,
                 e.g. "mean", "max", "min", "std", "sum", "var", "moment3".
                 The results from all aggregators will be concatenated.
 
-            scalers: list(str)
+            scalers:
                 Set of scaling functions identifiers
                 e.g. "identidy", "amplification", "attenuation"
                 The results from all scalers will be concatenated
 
-            activation: str, Callable
+            activation:
                 activation function to use in the layer
 
-            dropout: float
+            dropout:
                 The ratio of units to dropout. Must be between 0 and 1
 
-            batch_norm: bool
+            batch_norm:
                 Whether to use batch normalization
 
-            avg_d: Dict(str, float)
+            avg_d:
                 Average degree of nodes in the training set, used by scalers to normalize
 
-            last_activation: str, Callable
+            last_activation:
                 activation function to use in the last layer of the internal MLP
 
-            posttrans_layers: int
+            posttrans_layers:
                 number of layers in the MLP transformation after the aggregation
 
         """
@@ -232,35 +232,35 @@ def __init__(
             hidden_dim=self.out_dim,
             out_dim=self.out_dim,
             layers=posttrans_layers,
-            mid_activation=self.activation,
+            actibation=self.activation,
             last_activation=self.last_activation,
             dropout=dropout,
-            mid_batch_norm=batch_norm,
+            batch_norm=batch_norm,
             last_batch_norm=batch_norm,
         )
 
     def pretrans_edges(self, edges) -> Dict[str, torch.Tensor]:
-        """
+        r"""
         Return a mapping to the features of the source nodes.
         """
         return {"e": edges.src["h"]}
 
     def forward(self, g: DGLGraph, h: torch.Tensor) -> torch.Tensor:
-        """
+        r"""
         Apply the PNA convolutional layer, with the specified post transformation
 
         Parameters:
 
-            g: dgl.DGLGraph
+            g:
                 graph on which the convolution is done
 
-            h: torch.Tensor(..., N, Din)
+            h: `torch.Tensor[..., N, Din]`
                 Node feature tensor, before convolution.
                 N is the number of nodes, Din is the input dimension ``self.in_dim``
 
         Returns:
 
-            h: torch.Tensor(..., N, Dout)
+            `torch.Tensor[..., N, Dout]`:
                 Node feature tensor, after convolution.
                 N is the number of nodes, Dout is the output dimension ``self.out_dim``
 
@@ -280,7 +280,7 @@ def forward(self, g: DGLGraph, h: torch.Tensor) -> torch.Tensor:
 
     @property
     def layer_inputs_edges(self) -> bool:
-        """
+        r"""
         Return a boolean specifying if the layer type
         uses edges as input or not.
         It is different from ``layer_supports_edges`` since a layer that
@@ -288,19 +288,19 @@ def layer_inputs_edges(self) -> bool:
 
         Returns:
 
-            uses_edges: bool
+            bool:
                 Always ``False`` for the current class
         """
         return False
 
     @classproperty
     def layer_supports_edges(cls) -> bool:
-        """
+        r"""
         Return a boolean specifying if the layer type supports edges or not.
 
         Returns:
 
-            supports_edges: bool
+            bool:
                 Always ``False`` for the current class
         """
         return False
@@ -313,11 +313,11 @@ def __init__(
         out_dim: int,
         aggregators: List[str],
         scalers: List[str],
-        activation="relu",
+        activation: Union[Callable, str] = "relu",
         dropout: float = 0.0,
         batch_norm: bool = False,
         avg_d: Dict[str, float] = {"log": 1.0},
-        last_activation="none",
+        last_activation: Union[Callable, str] = "none",
         posttrans_layers: int = 1,
         pretrans_layers: int = 1,
         in_dim_edges: int = 0,
@@ -335,44 +335,44 @@ def __init__(
 
         Parameters:
 
-            in_dim: int
+            in_dim:
                 Input feature dimensions of the layer
 
-            out_dim: int
+            out_dim:
                 Output feature dimensions of the layer
 
-            aggregators: list(str)
+            aggregators:
                 Set of aggregation function identifiers,
                 e.g. "mean", "max", "min", "std", "sum", "var", "moment3".
                 The results from all aggregators will be concatenated.
 
-            scalers: list(str)
+            scalers:
                 Set of scaling functions identifiers
                 e.g. "identidy", "amplification", "attenuation"
                 The results from all scalers will be concatenated
 
-            activation: str, Callable
+            activation:
                 activation function to use in the layer
 
-            dropout: float
+            dropout:
                 The ratio of units to dropout. Must be between 0 and 1
 
-            batch_norm: bool
+            batch_norm:
                 Whether to use batch normalization
 
-            avg_d: Dict(str, float)
+            avg_d:
                 Average degree of nodes in the training set, used by scalers to normalize
 
-            last_activation: str, Callable
+            last_activation:
                 activation function to use in the last layer of the internal MLP
 
-            posttrans_layers: int
+            posttrans_layers:
                 number of layers in the MLP transformation after the aggregation
 
-            pretrans_layers: int
+            pretrans_layers:
                 number of layers in the transformation before the aggregation
 
-            in_dim_edges: int
+            in_dim_edges:
                 size of the edge features. If 0, edges are ignored
 
         """
@@ -398,10 +398,10 @@ def __init__(
             hidden_dim=in_dim,
             out_dim=in_dim,
             layers=pretrans_layers,
-            mid_activation=self.activation,
+            actibation=self.activation,
             last_activation=self.last_activation,
             dropout=dropout,
-            mid_batch_norm=batch_norm,
+            batch_norm=batch_norm,
             last_batch_norm=batch_norm,
         )
 
@@ -411,40 +411,38 @@ def __init__(
             hidden_dim=out_dim,
             out_dim=out_dim,
             layers=posttrans_layers,
-            mid_activation=self.activation,
+            actibation=self.activation,
             last_activation=self.last_activation,
             dropout=dropout,
-            mid_batch_norm=batch_norm,
+            batch_norm=batch_norm,
             last_batch_norm=batch_norm,
         )
 
     def pretrans_edges(self, edges) -> Dict[str, torch.Tensor]:
-        """
+        r"""
         Return a mapping to the concatenation of the features from
         the source node, the destination node, and the edge between them (if applicable).
         """
         if self.edge_features:
-            z2 = torch.cat([edges.src["h"], edges.dst["h"], edges.data["ef"]], dim=1)
+            z2 = torch.cat([edges.src["h"], edges.dst["h"], edges.data["ef"]], dim=-1)
         else:
-            z2 = torch.cat([edges.src["h"], edges.dst["h"]], dim=1)
+            z2 = torch.cat([edges.src["h"], edges.dst["h"]], dim=-1)
         return {"e": self.pretrans(z2)}
 
-    def forward(
-        self, g: DGLGraph, h: torch.Tensor, e: torch.Tensor = None
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        """
+    def forward(self, g: DGLGraph, h: torch.Tensor, e: torch.Tensor = None) -> torch.Tensor:
+        r"""
         Apply the PNA Message passing layer, with the specified pre/post transformations
 
         Parameters:
 
-            g: dgl.DGLGraph
+            g:
                 graph on which the convolution is done
 
-            h: torch.Tensor(..., N, Din)
+            h: `torch.Tensor[..., N, Din]`
                 Node feature tensor, before convolution.
                 N is the number of nodes, Din is the input dimension ``self.in_dim``
 
-            e: torch.Tensor(..., N, Din_edges) or None
+            e: `torch.Tensor(..., N, Din_edges)` or `None`
                 Edge feature tensor, before convolution.
                 N is the number of nodes, Din is the input edge dimension
 
@@ -453,7 +451,7 @@ def forward(
 
         Returns:
 
-            h: torch.Tensor(..., N, Dout)
+            `torch.Tensor[..., N, Dout]`:
                 Node feature tensor, after convolution.
                 N is the number of nodes, Dout is the output dimension ``self.out_dim``
 
@@ -468,7 +466,7 @@ def forward(
 
         # aggregation
         g.update_all(self.message_func, self.reduce_func)
-        h = torch.cat([h, g.ndata["h"]], dim=1)
+        h = torch.cat([h, g.ndata["h"]], dim=-1)
 
         # post-transformation
         h = self.posttrans(h)
@@ -477,7 +475,7 @@ def forward(
 
     @property
     def layer_inputs_edges(self) -> bool:
-        """
+        r"""
         Return a boolean specifying if the layer type
         uses edges as input or not.
         It is different from ``layer_supports_edges`` since a layer that
@@ -485,19 +483,19 @@ def layer_inputs_edges(self) -> bool:
 
         Returns:
 
-            uses_edges: bool
+            bool:
                 Returns ``self.edge_features``
         """
         return self.edge_features
 
     @classproperty
     def layer_supports_edges(cls) -> bool:
-        """
+        r"""
         Return a boolean specifying if the layer type supports edges or not.
 
         Returns:
 
-            supports_edges: bool
+            bool:
                 Always ``True`` for the current class
         """
         return True
diff --git a/goli/dgl/dgl_layers/pooling.py b/goli/dgl/dgl_layers/pooling.py
index 0aa0ae4d4..42a14ec39 100644
--- a/goli/dgl/dgl_layers/pooling.py
+++ b/goli/dgl/dgl_layers/pooling.py
@@ -1,7 +1,8 @@
 import torch
 import torch.nn as nn
-from typing import List
+from typing import List, Union, Callable, Tuple
 
+import dgl
 from dgl.nn.pytorch.glob import SumPooling, AvgPooling, MaxPooling, Set2Set, GlobalAttentionPooling
 from dgl import mean_nodes, sum_nodes, max_nodes
 
@@ -9,12 +10,23 @@
 from goli.commons.utils import ModuleListConcat
 
 
+EPS = 1e-6
+
+
 class S2SReadout(nn.Module):
     r"""
     Performs a Set2Set aggregation of all the graph nodes' features followed by a series of fully connected layers
     """
 
-    def __init__(self, in_dim, hidden_dim, out_dim, fc_layers=3, device="cpu", final_activation="relu"):
+    def __init__(
+        self,
+        in_dim,
+        hidden_dim,
+        out_dim,
+        fc_layers=3,
+        device="cpu",
+        final_activation: Union[str, Callable] = "relu",
+    ):
         super().__init__()
 
         # set2set aggregation
@@ -26,9 +38,9 @@ def __init__(self, in_dim, hidden_dim, out_dim, fc_layers=3, device="cpu", final
             hidden_dim=hidden_dim,
             out_dim=out_dim,
             layers=fc_layers,
-            mid_activation="relu",
+            actibation="relu",
             last_activation=final_activation,
-            mid_batch_norm=True,
+            batch_norm=True,
             last_batch_norm=False,
             device=device,
         )
@@ -109,12 +121,12 @@ def parse_pooling_layer(in_dim: int, pooling: List[str], n_iters: int = 2, n_lay
         pooling: list(str)
             The list of pooling layers to use. The accepted strings are:
 
-            - "sum": SumPooling
-            - "mean": MeanPooling
-            - "max": MaxPooling
-            - "min": MinPooling
-            - "std": StdPooling
-            - "s2s": Set2Set
+            - "sum": `SumPooling`
+            - "mean": `MeanPooling`
+            - "max": `MaxPooling`
+            - "min": `MinPooling`
+            - "std": `StdPooling`
+            - "s2s": `Set2Set`
 
         n_iters: int
             IGNORED FOR ALL POOLING LAYERS, EXCEPT "s2s".
@@ -128,7 +140,6 @@ def parse_pooling_layer(in_dim: int, pooling: List[str], n_iters: int = 2, n_lay
     # TODO: Add configuration for the pooling layer kwargs
 
     # Create the pooling layer
-    pooling = pooling.lower()
     pool_layer = ModuleListConcat()
     out_pool_dim = 0
 
@@ -145,6 +156,7 @@ def parse_pooling_layer(in_dim: int, pooling: List[str], n_iters: int = 2, n_lay
             pool_layer.append(MinPooling())
         elif this_pool == "std":
             pool_layer.append(StdPooling())
+        elif this_pool == "s2s":
             pool_layer.append(Set2Set(input_dim=in_dim, n_iters=n_iters, n_layers=n_layers))
             out_pool_dim += in_dim
         elif (this_pool == "none") or (this_pool is None):
@@ -157,8 +169,42 @@ def parse_pooling_layer(in_dim: int, pooling: List[str], n_iters: int = 2, n_lay
 
 class VirtualNode(nn.Module):
     def __init__(
-        self, dim, dropout, batch_norm=False, bias=True, activation="relu", residual=True, vn_type="sum"
+        self,
+        dim: int,
+        vn_type: Union[type(None), str] = "sum",
+        activation: Union[str, Callable] = "relu",
+        dropout: float = 0.0,
+        batch_norm: bool = False,
+        bias: bool = True,
+        residual: bool = True,
     ):
+        r"""
+        The VirtualNode is a layer that pool the features of the graph,
+        applies a neural network layer on the pooled features,
+        then add the result back to the node features of every node.
+
+        Parameters:
+
+            in_dim:
+                Input feature dimensions of the virtual node layer
+
+            activation:
+                activation function to use in the neural network layer.
+
+            dropout:
+                The ratio of units to dropout. Must be between 0 and 1
+
+            batch_norm:
+                Whether to use batch normalization
+
+            bias:
+                Whether to add a bias to the neural network
+
+            residual:
+                Whether all virtual nodes should be connected together
+                via a residual connection
+
+        """
         super().__init__()
         if (vn_type is None) or (vn_type.lower() == "none"):
             self.vn_type = None
@@ -177,13 +223,43 @@ def __init__(
             bias=bias,
         )
 
-    def forward(self, g, h, vn_h):
+    def forward(
+        self, g: dgl.DGLGraph, h: torch.Tensor, vn_h: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        r"""
+        Apply the virtual node layer.
+
+        Parameters:
+
+            g:
+                graph on which the convolution is done
+
+            h (torch.Tensor[..., N, Din]):
+                Node feature tensor, before convolution.
+                `N` is the number of nodes, `Din` is the input features
+
+            vn_h (torch.Tensor[..., M, Din]):
+                Graph feature of the previous virtual node, or `None`
+                `M` is the number of graphs, `Din` is the input features.
+                It is added to the result after the MLP, as a residual connection
+
+        Returns:
+
+            `h = torch.Tensor[..., N, Dout]`:
+                Node feature tensor, after convolution and residual.
+                `N` is the number of nodes, `Dout` is the output features of the layer and residual
+
+            `vn_h = torch.Tensor[..., M, Dout]`:
+                Graph feature tensor to be used at the next virtual node, or `None`
+                `M` is the number of graphs, `Dout` is the output features
+
+        """
 
         g.ndata["h"] = h
 
         # Pool the features
         if self.vn_type is None:
-            return vn_h, h
+            return h, vn_h
         elif self.vn_type == "mean":
             pool = mean_nodes(g, "h")
         elif self.vn_type == "max":
@@ -213,4 +289,4 @@ def forward(self, g, h, vn_h):
         )
         h = h + temp_h
 
-        return vn_h, h
+        return h, vn_h
diff --git a/goli/dgl/pna_operations.py b/goli/dgl/pna_operations.py
index da71ad434..a95b2ce6e 100644
--- a/goli/dgl/pna_operations.py
+++ b/goli/dgl/pna_operations.py
@@ -6,15 +6,15 @@
 
 
 def aggregate_mean(h, **kwargs):
-    return torch.mean(h, dim=1)
+    return torch.mean(h, dim=-2)
 
 
 def aggregate_max(h, **kwargs):
-    return torch.max(h, dim=1)[0]
+    return torch.max(h, dim=-2)[0]
 
 
 def aggregate_min(h, **kwargs):
-    return torch.min(h, dim=1)[0]
+    return torch.min(h, dim=-2)[0]
 
 
 def aggregate_std(h, **kwargs):
@@ -22,8 +22,8 @@ def aggregate_std(h, **kwargs):
 
 
 def aggregate_var(h, **kwargs):
-    h_mean_squares = torch.mean(h * h, dim=1)
-    h_mean = torch.mean(h, dim=1)
+    h_mean_squares = torch.mean(h * h, dim=-2)
+    h_mean = torch.mean(h, dim=-2)
     var = torch.relu(h_mean_squares - h_mean * h_mean)
     return var
 
@@ -31,14 +31,14 @@ def aggregate_var(h, **kwargs):
 def aggregate_moment(h, n=3, **kwargs):
     # for each node (E[(X-E[X])^n])^{1/n}
     # EPS is added to the absolute value of expectation before taking the nth root for stability
-    h_mean = torch.mean(h, dim=1, keepdim=True)
-    h_n = torch.mean(torch.pow(h - h_mean, n), dim=1)
+    h_mean = torch.mean(h, dim=-2, keepdim=True)
+    h_n = torch.mean(torch.pow(h - h_mean, n), dim=-2)
     rooted_h_n = torch.sign(h_n) * torch.pow(torch.abs(h_n) + EPS, 1.0 / n)
     return rooted_h_n
 
 
 def aggregate_sum(h, **kwargs):
-    return torch.sum(h, dim=1)
+    return torch.sum(h, dim=-2)
 
 
 # each scaler is a function that takes as input X (B x N x Din), adj (B x N x N) and
diff --git a/goli/dgl/residual_connections.py b/goli/dgl/residual_connections.py
index 754bfa2d1..fba21679b 100644
--- a/goli/dgl/residual_connections.py
+++ b/goli/dgl/residual_connections.py
@@ -6,7 +6,7 @@
 import abc
 import torch
 import torch.nn as nn
-from typing import List
+from typing import List, Union, Callable
 
 from goli.dgl.base_layers import FCLayer
 from goli.commons.decorators import classproperty
@@ -28,6 +28,8 @@ def __init__(self, skip_steps: int = 1):
 
             skip_steps: int
                 The number of steps to skip between the residual connections.
+                If `1`, all the layers are connected. If `2`, half of the
+                layers are connected.
         """
 
         super().__init__()
@@ -77,7 +79,7 @@ def get_true_out_dims(self, out_dims: List):
 
         true_out_dims = [out_dims[0]]
         out_dims_at_skip = [out_dims[0]]
-        for ii in range(1, len(out_dims)):
+        for ii in range(1, len(out_dims) - 1):
 
             # For the `None` type, don't change the output dims
             if self.h_dim_increase_type is None:
@@ -174,6 +176,8 @@ def __init__(self, skip_steps: int = 1):
 
             skip_steps: int
                 The number of steps to skip between the residual connections.
+                If `1`, all the layers are connected. If `2`, half of the
+                layers are connected.
         """
         super().__init__(skip_steps=skip_steps)
 
@@ -237,7 +241,13 @@ def forward(self, h: torch.Tensor, h_prev: torch.Tensor, step_idx: int):
 
 class ResidualConnectionWeighted(ResidualConnectionBase):
     def __init__(
-        self, out_dims, skip_steps: int = 1, dropout=0.0, activation="none", batch_norm=False, bias=False
+        self,
+        out_dims,
+        skip_steps: int = 1,
+        dropout=0.0,
+        activation: Union[str, Callable] = "none",
+        batch_norm=False,
+        bias=False,
     ):
         r"""
         Class for the simple residual connections proposed by ResNet,
@@ -249,6 +259,8 @@ def __init__(
 
             skip_steps: int
                 The number of steps to skip between the residual connections.
+                If `1`, all the layers are connected. If `2`, half of the
+                layers are connected.
 
             out_dims: list(int)
                 list of all output dimensions for the network
@@ -276,7 +288,7 @@ def __init__(
         self.skip_count = 0
         self.out_dims = out_dims
 
-        for ii in range(0, len(self.out_dims), self.skip_steps):
+        for ii in range(0, len(self.out_dims) - 1, self.skip_steps):
             this_dim = self.out_dims[ii]
             self.residual_list.append(
                 FCLayer(
@@ -359,6 +371,8 @@ def __init__(self, skip_steps: int = 1):
 
             skip_steps: int
                 The number of steps to skip between the residual connections.
+                If `1`, all the layers are connected. If `2`, half of the
+                layers are connected.
         """
 
         super().__init__(skip_steps=skip_steps)
@@ -434,6 +448,8 @@ def __init__(self, skip_steps: int = 1):
 
             skip_steps: int
                 The number of steps to skip between the residual connections.
+                If `1`, all the layers are connected. If `2`, half of the
+                layers are connected.
         """
 
         super().__init__(skip_steps=skip_steps)
diff --git a/goli/trainer/metrics.py b/goli/trainer/metrics.py
index a00275875..3dd57b595 100644
--- a/goli/trainer/metrics.py
+++ b/goli/trainer/metrics.py
@@ -90,7 +90,7 @@ def forward(
 
 
 def pearsonr(preds: torch.Tensor, target: torch.Tensor, reduction: str = "elementwise_mean") -> torch.Tensor:
-    """
+    r"""
     Computes the pearsonr correlation.
 
     Parameters:
@@ -125,7 +125,7 @@ def pearsonr(preds: torch.Tensor, target: torch.Tensor, reduction: str = "elemen
 
 
 def spearmanr(preds: torch.Tensor, target: torch.Tensor, reduction: str = "elementwise_mean") -> torch.Tensor:
-    """
+    r"""
     Computes the spearmanr correlation.
 
     Parameters:
diff --git a/goli/trainer/model_wrapper.py b/goli/trainer/model_wrapper.py
index f10e59850..78d3e003e 100644
--- a/goli/trainer/model_wrapper.py
+++ b/goli/trainer/model_wrapper.py
@@ -18,7 +18,7 @@
 
 
 class EpochSummary:
-    """Container for collecting epoch-wise results"""
+    r"""Container for collecting epoch-wise results"""
 
     class Results:
         def __init__(self, targets: torch.Tensor, predictions: torch.Tensor, loss: float, metrics: dict):
diff --git a/mkdocs.yml b/mkdocs.yml
index e7751c8a7..f1487d8ef 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -11,6 +11,10 @@ docs_dir: "docs"
 
 nav:
   - Overview: index.md
+  - Examples:
+      - Using GNN layers: examples/basics/using_gnn_layers.ipynb
+      - Implementing GNN layers: examples/basics/implementing_gnn_layers.ipynb
+      - Making GNN networks: examples/basics/making_gnn_networks.ipynb
   - API:
     # - goli.cli: api/goli.cli.md
     - goli.commons: api/goli.commons.md


Test patch:
diff --git a/tests/unit_tests/dgl/test_architectures.py b/tests/unit_tests/dgl/test_architectures.py
new file mode 100644
index 000000000..d01f29b8c
--- /dev/null
+++ b/tests/unit_tests/dgl/test_architectures.py
@@ -0,0 +1,634 @@
+"""
+Unit tests for the different layers of goli/dgl/dgl_layers/...
+
+The layers are not thoroughly tested due to the difficulty of testing them
+"""
+
+import numpy as np
+import torch
+import unittest as ut
+import dgl
+from copy import deepcopy
+
+from goli.dgl.architectures import FeedForwardNN, FeedForwardDGL, FullDGLNetwork, FullDGLSiameseNetwork
+from goli.dgl.base_layers import FCLayer
+from goli.dgl.residual_connections import (
+    ResidualConnectionConcat,
+    ResidualConnectionDenseNet,
+    ResidualConnectionNone,
+    ResidualConnectionSimple,
+    ResidualConnectionWeighted,
+)
+
+
+class test_FeedForwardNN(ut.TestCase):
+
+    kwargs = {
+        "activation": "relu",
+        "last_activation": "none",
+        "batch_norm": False,
+        "dropout": 0.2,
+        "name": "LNN",
+        "layer_type": FCLayer,
+    }
+
+    def test_forward_no_residual(self):
+        in_dim = 8
+        out_dim = 16
+        hidden_dims = [5, 6, 7]
+        batch = 2
+
+        lnn = FeedForwardNN(
+            in_dim=in_dim,
+            out_dim=out_dim,
+            hidden_dims=hidden_dims,
+            residual_type="none",
+            residual_skip_steps=1,
+            **self.kwargs,
+        )
+
+        self.assertEqual(len(lnn.layers), len(hidden_dims) + 1)
+        self.assertEqual(lnn.layers[0].in_dim, in_dim)
+        self.assertEqual(lnn.layers[1].in_dim, hidden_dims[0])
+        self.assertEqual(lnn.layers[2].in_dim, hidden_dims[1])
+        self.assertEqual(lnn.layers[3].in_dim, hidden_dims[2])
+
+        h = torch.FloatTensor(batch, in_dim)
+        h_out = lnn.forward(h)
+
+        self.assertListEqual(list(h_out.shape), [batch, out_dim])
+
+    def test_forward_simple_residual_1(self):
+        in_dim = 8
+        out_dim = 16
+        hidden_dims = [6, 6, 6, 6, 6]
+        batch = 2
+
+        lnn = FeedForwardNN(
+            in_dim=in_dim,
+            out_dim=out_dim,
+            hidden_dims=hidden_dims,
+            residual_type="simple",
+            residual_skip_steps=1,
+            **self.kwargs,
+        )
+
+        self.assertEqual(len(lnn.layers), len(hidden_dims) + 1)
+        self.assertEqual(lnn.layers[0].in_dim, in_dim)
+        self.assertEqual(lnn.layers[1].in_dim, hidden_dims[0])
+        self.assertEqual(lnn.layers[2].in_dim, hidden_dims[1])
+        self.assertEqual(lnn.layers[3].in_dim, hidden_dims[2])
+
+        h = torch.FloatTensor(batch, in_dim)
+        h_out = lnn.forward(h)
+
+        self.assertListEqual(list(h_out.shape), [batch, out_dim])
+
+    def test_forward_simple_residual_2(self):
+        in_dim = 8
+        out_dim = 16
+        hidden_dims = [6, 6, 6, 6, 6]
+        batch = 2
+
+        lnn = FeedForwardNN(
+            in_dim=in_dim,
+            out_dim=out_dim,
+            hidden_dims=hidden_dims,
+            residual_type="simple",
+            residual_skip_steps=2,
+            **self.kwargs,
+        )
+
+        self.assertEqual(len(lnn.layers), len(hidden_dims) + 1)
+        self.assertEqual(lnn.layers[0].in_dim, in_dim)
+        self.assertEqual(lnn.layers[1].in_dim, hidden_dims[0])
+        self.assertEqual(lnn.layers[2].in_dim, hidden_dims[1])
+        self.assertEqual(lnn.layers[3].in_dim, hidden_dims[2])
+        self.assertEqual(lnn.layers[4].in_dim, hidden_dims[3])
+        self.assertEqual(lnn.layers[5].in_dim, hidden_dims[4])
+
+        h = torch.FloatTensor(batch, in_dim)
+        h_out = lnn.forward(h)
+
+        self.assertListEqual(list(h_out.shape), [batch, out_dim])
+
+    def test_forward_concat_residual_1(self):
+        in_dim = 8
+        out_dim = 16
+        hidden_dims = [6, 6, 6, 6, 6]
+        batch = 2
+
+        lnn = FeedForwardNN(
+            in_dim=in_dim,
+            out_dim=out_dim,
+            hidden_dims=hidden_dims,
+            residual_type="concat",
+            residual_skip_steps=1,
+            **self.kwargs,
+        )
+
+        self.assertEqual(len(lnn.layers), len(hidden_dims) + 1)
+        self.assertEqual(lnn.layers[0].in_dim, in_dim)
+        self.assertEqual(lnn.layers[1].in_dim, hidden_dims[0])
+        self.assertEqual(lnn.layers[2].in_dim, 2 * hidden_dims[1])
+        self.assertEqual(lnn.layers[3].in_dim, 2 * hidden_dims[2])
+        self.assertEqual(lnn.layers[4].in_dim, 2 * hidden_dims[3])
+        self.assertEqual(lnn.layers[5].in_dim, 2 * hidden_dims[4])
+
+        h = torch.FloatTensor(batch, in_dim)
+        h_out = lnn.forward(h)
+
+        self.assertListEqual(list(h_out.shape), [batch, out_dim])
+
+    def test_forward_concat_residual_2(self):
+        in_dim = 8
+        out_dim = 16
+        hidden_dims = [6, 6, 6, 6, 6]
+        batch = 2
+
+        lnn = FeedForwardNN(
+            in_dim=in_dim,
+            out_dim=out_dim,
+            hidden_dims=hidden_dims,
+            residual_type="concat",
+            residual_skip_steps=2,
+            **self.kwargs,
+        )
+
+        self.assertEqual(len(lnn.layers), len(hidden_dims) + 1)
+        self.assertEqual(lnn.layers[0].in_dim, in_dim)
+        self.assertEqual(lnn.layers[1].in_dim, 1 * hidden_dims[0])
+        self.assertEqual(lnn.layers[2].in_dim, 1 * hidden_dims[1])
+        self.assertEqual(lnn.layers[3].in_dim, 2 * hidden_dims[2])
+        self.assertEqual(lnn.layers[4].in_dim, 1 * hidden_dims[3])
+        self.assertEqual(lnn.layers[5].in_dim, 2 * hidden_dims[4])
+
+        h = torch.FloatTensor(batch, in_dim)
+        h_out = lnn.forward(h)
+
+        self.assertListEqual(list(h_out.shape), [batch, out_dim])
+
+    def test_forward_densenet_residual_1(self):
+        in_dim = 8
+        out_dim = 16
+        hidden_dims = [6, 6, 6, 6, 6]
+        batch = 2
+
+        lnn = FeedForwardNN(
+            in_dim=in_dim,
+            out_dim=out_dim,
+            hidden_dims=hidden_dims,
+            residual_type="densenet",
+            residual_skip_steps=1,
+            **self.kwargs,
+        )
+
+        self.assertEqual(len(lnn.layers), len(hidden_dims) + 1)
+        self.assertEqual(lnn.layers[0].in_dim, in_dim)
+        self.assertEqual(lnn.layers[1].in_dim, hidden_dims[0])
+        self.assertEqual(lnn.layers[2].in_dim, 2 * hidden_dims[1])
+        self.assertEqual(lnn.layers[3].in_dim, 3 * hidden_dims[2])
+        self.assertEqual(lnn.layers[4].in_dim, 4 * hidden_dims[3])
+        self.assertEqual(lnn.layers[5].in_dim, 5 * hidden_dims[4])
+
+        h = torch.FloatTensor(batch, in_dim)
+        h_out = lnn.forward(h)
+
+        self.assertListEqual(list(h_out.shape), [batch, out_dim])
+
+    def test_forward_densenet_residual_2(self):
+        in_dim = 8
+        out_dim = 16
+        hidden_dims = [6, 6, 6, 6, 6]
+        batch = 2
+
+        lnn = FeedForwardNN(
+            in_dim=in_dim,
+            out_dim=out_dim,
+            hidden_dims=hidden_dims,
+            residual_type="densenet",
+            residual_skip_steps=2,
+            **self.kwargs,
+        )
+
+        self.assertEqual(len(lnn.layers), len(hidden_dims) + 1)
+        self.assertEqual(lnn.layers[0].in_dim, in_dim)
+        self.assertEqual(lnn.layers[1].in_dim, hidden_dims[0])
+        self.assertEqual(lnn.layers[2].in_dim, 1 * hidden_dims[1])
+        self.assertEqual(lnn.layers[3].in_dim, 2 * hidden_dims[2])
+        self.assertEqual(lnn.layers[4].in_dim, 1 * hidden_dims[3])
+        self.assertEqual(lnn.layers[5].in_dim, 3 * hidden_dims[4])
+
+        h = torch.FloatTensor(batch, in_dim)
+        h_out = lnn.forward(h)
+
+        self.assertListEqual(list(h_out.shape), [batch, out_dim])
+
+    def test_forward_weighted_residual_1(self):
+        in_dim = 8
+        out_dim = 16
+        hidden_dims = [6, 6, 6, 6, 6]
+        batch = 2
+
+        lnn = FeedForwardNN(
+            in_dim=in_dim,
+            out_dim=out_dim,
+            hidden_dims=hidden_dims,
+            residual_type="weighted",
+            residual_skip_steps=1,
+            **self.kwargs,
+        )
+
+        self.assertEqual(len(lnn.layers), len(hidden_dims) + 1)
+        self.assertEqual(lnn.layers[0].in_dim, in_dim)
+        self.assertEqual(lnn.layers[1].in_dim, hidden_dims[0])
+        self.assertEqual(lnn.layers[2].in_dim, hidden_dims[1])
+        self.assertEqual(lnn.layers[3].in_dim, hidden_dims[2])
+        self.assertEqual(lnn.layers[4].in_dim, hidden_dims[3])
+        self.assertEqual(lnn.layers[5].in_dim, hidden_dims[4])
+
+        self.assertEqual(len(lnn.residual_layer.residual_list), len(hidden_dims))
+
+        h = torch.FloatTensor(batch, in_dim)
+        h_out = lnn.forward(h)
+
+        self.assertListEqual(list(h_out.shape), [batch, out_dim])
+
+    def test_forward_weighted_residual_2(self):
+        in_dim = 8
+        out_dim = 16
+        hidden_dims = [6, 6, 6, 6, 6]
+        batch = 2
+
+        lnn = FeedForwardNN(
+            in_dim=in_dim,
+            out_dim=out_dim,
+            hidden_dims=hidden_dims,
+            residual_type="weighted",
+            residual_skip_steps=2,
+            **self.kwargs,
+        )
+
+        self.assertEqual(len(lnn.layers), len(hidden_dims) + 1)
+        self.assertEqual(lnn.layers[0].in_dim, in_dim)
+        self.assertEqual(lnn.layers[1].in_dim, hidden_dims[0])
+        self.assertEqual(lnn.layers[2].in_dim, hidden_dims[1])
+        self.assertEqual(lnn.layers[3].in_dim, hidden_dims[2])
+        self.assertEqual(lnn.layers[4].in_dim, hidden_dims[3])
+        self.assertEqual(lnn.layers[5].in_dim, hidden_dims[4])
+
+        self.assertEqual(len(lnn.residual_layer.residual_list), (len(hidden_dims) // 2 + 1))
+
+        h = torch.FloatTensor(batch, in_dim)
+        h_out = lnn.forward(h)
+
+        self.assertListEqual(list(h_out.shape), [batch, out_dim])
+
+
+class test_FeedForwardDGL(ut.TestCase):
+
+    kwargs = {
+        "activation": "relu",
+        "last_activation": "none",
+        "batch_norm": False,
+        "dropout": 0.2,
+        "name": "LNN",
+    }
+
+    in_dim = 7
+    out_dim = 11
+    in_dim_edges = 13
+    hidden_dims = [6, 6, 6, 6, 6]
+
+    g1 = dgl.graph((torch.tensor([0, 1, 2]), torch.tensor([1, 2, 3])))
+    g2 = dgl.graph((torch.tensor([0, 0, 0, 1]), torch.tensor([0, 1, 2, 0])))
+    g1.ndata["h"] = torch.zeros(g1.num_nodes(), in_dim, dtype=torch.float32)
+    g1.edata["e"] = torch.ones(g1.num_edges(), in_dim_edges, dtype=torch.float32)
+    g2.ndata["h"] = torch.ones(g2.num_nodes(), in_dim, dtype=torch.float32)
+    g2.edata["e"] = torch.zeros(g2.num_edges(), in_dim_edges, dtype=torch.float32)
+    bg = dgl.batch([g1, g2])
+    bg = dgl.add_self_loop(bg)
+
+    virtual_nodes = ["none", "mean", "sum"]
+    pna_kwargs = {"aggregators": ["mean", "max", "sum"], "scalers": ["identity", "amplification"]}
+
+    layers_kwargs = {
+        "gcn": {},
+        "gin": {},
+        "gat": {"num_heads": 3},
+        "gated-gcn": {"in_dim_edges": in_dim_edges, "hidden_dims_edges": hidden_dims},
+        "pna-conv": {**pna_kwargs},
+        "pna-msgpass#1": {**pna_kwargs, "in_dim_edges": 0},
+        "pna-msgpass#2": {**pna_kwargs, "in_dim_edges": in_dim_edges},
+    }
+
+    def test_dgl_forward_no_residual(self):
+
+        for pooling in [["none"], ["sum"], ["mean", "s2s", "max"]]:
+            for residual_skip_steps in [1, 2, 3]:
+                for virtual_node in self.virtual_nodes:
+                    for layer_name, this_kwargs in self.layers_kwargs.items():
+                        err_msg = f"pooling={pooling}, virtual_node={virtual_node}, layer_name={layer_name}, residual_skip_steps={residual_skip_steps}"
+                        layer_type = layer_name.split("#")[0]
+
+                        gnn = FeedForwardDGL(
+                            in_dim=self.in_dim,
+                            out_dim=self.out_dim,
+                            hidden_dims=self.hidden_dims,
+                            residual_type="none",
+                            residual_skip_steps=residual_skip_steps,
+                            layer_type=layer_type,
+                            pooling=pooling,
+                            **this_kwargs,
+                            **self.kwargs,
+                        )
+                        gnn.to(torch.float32)
+
+                        self.assertIsInstance(gnn.residual_layer, ResidualConnectionNone)
+                        self.assertEqual(len(gnn.layers), len(self.hidden_dims) + 1, msg=err_msg)
+                        self.assertEqual(gnn.layers[0].out_dim, self.hidden_dims[0], msg=err_msg)
+                        self.assertEqual(gnn.layers[1].out_dim, self.hidden_dims[1], msg=err_msg)
+                        self.assertEqual(gnn.layers[2].out_dim, self.hidden_dims[2], msg=err_msg)
+                        self.assertEqual(gnn.layers[3].out_dim, self.hidden_dims[3], msg=err_msg)
+                        self.assertEqual(gnn.layers[4].out_dim, self.hidden_dims[4], msg=err_msg)
+                        self.assertEqual(gnn.layers[5].out_dim, self.out_dim, msg=err_msg)
+
+                        f = gnn.layers[0].out_dim_factor
+                        self.assertEqual(gnn.layers[0].in_dim, self.in_dim, msg=err_msg)
+                        self.assertEqual(gnn.layers[1].in_dim, f * self.hidden_dims[0], msg=err_msg)
+                        self.assertEqual(gnn.layers[2].in_dim, f * self.hidden_dims[1], msg=err_msg)
+                        self.assertEqual(gnn.layers[3].in_dim, f * self.hidden_dims[2], msg=err_msg)
+                        self.assertEqual(gnn.layers[4].in_dim, f * self.hidden_dims[3], msg=err_msg)
+                        self.assertEqual(gnn.layers[5].in_dim, f * self.hidden_dims[4], msg=err_msg)
+
+                        bg = deepcopy(self.bg)
+                        h_out = gnn.forward(bg)
+
+                        dim_1 = bg.num_nodes() if pooling == ["none"] else 1
+                        self.assertListEqual(list(h_out.shape), [dim_1, self.out_dim], msg=err_msg)
+
+    def test_dgl_forward_simple_residual(self):
+
+        for pooling in [["none"], ["sum"], ["mean", "s2s", "max"]]:
+            for residual_skip_steps in [1, 2, 3]:
+                for virtual_node in self.virtual_nodes:
+                    for layer_name, this_kwargs in self.layers_kwargs.items():
+                        err_msg = f"pooling={pooling}, virtual_node={virtual_node}, layer_name={layer_name}, residual_skip_steps={residual_skip_steps}"
+                        layer_type = layer_name.split("#")[0]
+
+                        gnn = FeedForwardDGL(
+                            in_dim=self.in_dim,
+                            out_dim=self.out_dim,
+                            hidden_dims=self.hidden_dims,
+                            residual_type="simple",
+                            residual_skip_steps=1,
+                            layer_type=layer_type,
+                            pooling=pooling,
+                            **this_kwargs,
+                            **self.kwargs,
+                        )
+                        gnn.to(torch.float32)
+
+                        self.assertIsInstance(gnn.residual_layer, ResidualConnectionSimple)
+                        self.assertEqual(len(gnn.layers), len(self.hidden_dims) + 1, msg=err_msg)
+                        self.assertEqual(gnn.layers[0].out_dim, self.hidden_dims[0], msg=err_msg)
+                        self.assertEqual(gnn.layers[1].out_dim, self.hidden_dims[1], msg=err_msg)
+                        self.assertEqual(gnn.layers[2].out_dim, self.hidden_dims[2], msg=err_msg)
+                        self.assertEqual(gnn.layers[3].out_dim, self.hidden_dims[3], msg=err_msg)
+                        self.assertEqual(gnn.layers[4].out_dim, self.hidden_dims[4], msg=err_msg)
+                        self.assertEqual(gnn.layers[5].out_dim, self.out_dim, msg=err_msg)
+
+                        f = gnn.layers[0].out_dim_factor
+                        self.assertEqual(gnn.layers[0].in_dim, self.in_dim, msg=err_msg)
+                        self.assertEqual(gnn.layers[1].in_dim, f * self.hidden_dims[0], msg=err_msg)
+                        self.assertEqual(gnn.layers[2].in_dim, f * self.hidden_dims[1], msg=err_msg)
+                        self.assertEqual(gnn.layers[3].in_dim, f * self.hidden_dims[2], msg=err_msg)
+                        self.assertEqual(gnn.layers[4].in_dim, f * self.hidden_dims[3], msg=err_msg)
+                        self.assertEqual(gnn.layers[5].in_dim, f * self.hidden_dims[4], msg=err_msg)
+
+                        bg = deepcopy(self.bg)
+                        h_out = gnn.forward(bg)
+
+                        dim_1 = bg.num_nodes() if pooling == ["none"] else 1
+                        self.assertListEqual(list(h_out.shape), [dim_1, self.out_dim], msg=err_msg)
+
+    def test_dgl_forward_weighted_residual(self):
+
+        for pooling in [["none"], ["sum"], ["mean", "s2s", "max"]]:
+            for residual_skip_steps in [1, 2, 3]:
+                for virtual_node in self.virtual_nodes:
+                    for layer_name, this_kwargs in self.layers_kwargs.items():
+                        err_msg = f"pooling={pooling}, virtual_node={virtual_node}, layer_name={layer_name}, residual_skip_steps={residual_skip_steps}"
+                        layer_type = layer_name.split("#")[0]
+
+                        gnn = FeedForwardDGL(
+                            in_dim=self.in_dim,
+                            out_dim=self.out_dim,
+                            hidden_dims=self.hidden_dims,
+                            residual_type="weighted",
+                            residual_skip_steps=residual_skip_steps,
+                            layer_type=layer_type,
+                            pooling=pooling,
+                            **this_kwargs,
+                            **self.kwargs,
+                        )
+                        gnn.to(torch.float32)
+
+                        self.assertIsInstance(gnn.residual_layer, ResidualConnectionWeighted)
+                        self.assertEqual(len(gnn.layers), len(self.hidden_dims) + 1, msg=err_msg)
+                        self.assertEqual(gnn.layers[0].out_dim, self.hidden_dims[0], msg=err_msg)
+                        self.assertEqual(gnn.layers[1].out_dim, self.hidden_dims[1], msg=err_msg)
+                        self.assertEqual(gnn.layers[2].out_dim, self.hidden_dims[2], msg=err_msg)
+                        self.assertEqual(gnn.layers[3].out_dim, self.hidden_dims[3], msg=err_msg)
+                        self.assertEqual(gnn.layers[4].out_dim, self.hidden_dims[4], msg=err_msg)
+                        self.assertEqual(gnn.layers[5].out_dim, self.out_dim, msg=err_msg)
+
+                        f = gnn.layers[0].out_dim_factor
+                        self.assertEqual(gnn.layers[0].in_dim, self.in_dim, msg=err_msg)
+                        self.assertEqual(gnn.layers[1].in_dim, f * self.hidden_dims[0], msg=err_msg)
+                        self.assertEqual(gnn.layers[2].in_dim, f * self.hidden_dims[1], msg=err_msg)
+                        self.assertEqual(gnn.layers[3].in_dim, f * self.hidden_dims[2], msg=err_msg)
+                        self.assertEqual(gnn.layers[4].in_dim, f * self.hidden_dims[3], msg=err_msg)
+                        self.assertEqual(gnn.layers[5].in_dim, f * self.hidden_dims[4], msg=err_msg)
+
+                        bg = deepcopy(self.bg)
+                        h_out = gnn.forward(bg)
+
+                        dim_1 = bg.num_nodes() if pooling == ["none"] else 1
+                        self.assertListEqual(list(h_out.shape), [dim_1, self.out_dim], msg=err_msg)
+
+    def test_dgl_forward_concat_residual(self):
+
+        for pooling in [["none"], ["sum"], ["mean", "s2s", "max"]]:
+            for residual_skip_steps in [1, 2, 3]:
+                for virtual_node in self.virtual_nodes:
+                    for layer_name, this_kwargs in self.layers_kwargs.items():
+                        err_msg = f"pooling={pooling}, virtual_node={virtual_node}, layer_name={layer_name}, residual_skip_steps={residual_skip_steps}"
+                        layer_type = layer_name.split("#")[0]
+
+                        gnn = FeedForwardDGL(
+                            in_dim=self.in_dim,
+                            out_dim=self.out_dim,
+                            hidden_dims=self.hidden_dims,
+                            residual_type="concat",
+                            residual_skip_steps=residual_skip_steps,
+                            layer_type=layer_type,
+                            pooling=pooling,
+                            **this_kwargs,
+                            **self.kwargs,
+                        )
+                        gnn.to(torch.float32)
+
+                        self.assertIsInstance(gnn.residual_layer, ResidualConnectionConcat)
+                        self.assertEqual(len(gnn.layers), len(self.hidden_dims) + 1, msg=err_msg)
+                        self.assertEqual(gnn.layers[0].out_dim, self.hidden_dims[0], msg=err_msg)
+                        self.assertEqual(gnn.layers[1].out_dim, self.hidden_dims[1], msg=err_msg)
+                        self.assertEqual(gnn.layers[2].out_dim, self.hidden_dims[2], msg=err_msg)
+                        self.assertEqual(gnn.layers[3].out_dim, self.hidden_dims[3], msg=err_msg)
+                        self.assertEqual(gnn.layers[4].out_dim, self.hidden_dims[4], msg=err_msg)
+                        self.assertEqual(gnn.layers[5].out_dim, self.out_dim, msg=err_msg)
+
+                        f = gnn.layers[0].out_dim_factor
+                        f2 = [2 * f if ((ii % residual_skip_steps) == 0 and ii > 0) else f for ii in range(6)]
+                        self.assertEqual(gnn.layers[0].in_dim, self.in_dim, msg=err_msg)
+                        self.assertEqual(gnn.layers[1].in_dim, f2[0] * self.hidden_dims[0], msg=err_msg)
+                        self.assertEqual(gnn.layers[2].in_dim, f2[1] * self.hidden_dims[1], msg=err_msg)
+                        self.assertEqual(gnn.layers[3].in_dim, f2[2] * self.hidden_dims[2], msg=err_msg)
+                        self.assertEqual(gnn.layers[4].in_dim, f2[3] * self.hidden_dims[3], msg=err_msg)
+                        self.assertEqual(gnn.layers[5].in_dim, f2[4] * self.hidden_dims[4], msg=err_msg)
+
+                        bg = deepcopy(self.bg)
+                        h_out = gnn.forward(bg)
+
+                        dim_1 = bg.num_nodes() if pooling == ["none"] else 1
+                        self.assertListEqual(list(h_out.shape), [dim_1, self.out_dim], msg=err_msg)
+
+    def test_dgl_forward_densenet_residual(self):
+
+        for pooling in [["none"], ["sum"], ["mean", "s2s", "max"]]:
+            for residual_skip_steps in [1, 2, 3]:
+                for virtual_node in self.virtual_nodes:
+                    for layer_name, this_kwargs in self.layers_kwargs.items():
+                        err_msg = f"pooling={pooling}, virtual_node={virtual_node}, layer_name={layer_name}, residual_skip_steps={residual_skip_steps}"
+                        layer_type = layer_name.split("#")[0]
+
+                        gnn = FeedForwardDGL(
+                            in_dim=self.in_dim,
+                            out_dim=self.out_dim,
+                            hidden_dims=self.hidden_dims,
+                            residual_type="densenet",
+                            residual_skip_steps=residual_skip_steps,
+                            layer_type=layer_type,
+                            pooling=pooling,
+                            **this_kwargs,
+                            **self.kwargs,
+                        )
+                        gnn.to(torch.float32)
+
+                        self.assertIsInstance(gnn.residual_layer, ResidualConnectionDenseNet)
+                        self.assertEqual(len(gnn.layers), len(self.hidden_dims) + 1, msg=err_msg)
+                        self.assertEqual(gnn.layers[0].out_dim, self.hidden_dims[0], msg=err_msg)
+                        self.assertEqual(gnn.layers[1].out_dim, self.hidden_dims[1], msg=err_msg)
+                        self.assertEqual(gnn.layers[2].out_dim, self.hidden_dims[2], msg=err_msg)
+                        self.assertEqual(gnn.layers[3].out_dim, self.hidden_dims[3], msg=err_msg)
+                        self.assertEqual(gnn.layers[4].out_dim, self.hidden_dims[4], msg=err_msg)
+                        self.assertEqual(gnn.layers[5].out_dim, self.out_dim, msg=err_msg)
+
+                        f = gnn.layers[0].out_dim_factor
+                        f2 = [
+                            ((ii // residual_skip_steps) + 1) * f
+                            if ((ii % residual_skip_steps) == 0 and ii > 0)
+                            else f
+                            for ii in range(6)
+                        ]
+                        self.assertEqual(gnn.layers[0].in_dim, self.in_dim, msg=err_msg)
+                        self.assertEqual(gnn.layers[1].in_dim, f2[0] * self.hidden_dims[0], msg=err_msg)
+                        self.assertEqual(gnn.layers[2].in_dim, f2[1] * self.hidden_dims[1], msg=err_msg)
+                        self.assertEqual(gnn.layers[3].in_dim, f2[2] * self.hidden_dims[2], msg=err_msg)
+                        self.assertEqual(gnn.layers[4].in_dim, f2[3] * self.hidden_dims[3], msg=err_msg)
+                        self.assertEqual(gnn.layers[5].in_dim, f2[4] * self.hidden_dims[4], msg=err_msg)
+
+                        bg = deepcopy(self.bg)
+                        h_out = gnn.forward(bg)
+
+                        dim_1 = bg.num_nodes() if pooling == ["none"] else 1
+                        self.assertListEqual(list(h_out.shape), [dim_1, self.out_dim], msg=err_msg)
+
+
+class test_FullDGLNetwork(ut.TestCase):
+
+    kwargs = {
+        "activation": "relu",
+        "last_activation": "none",
+        "batch_norm": False,
+        "dropout": 0.2,
+        "name": "LNN",
+    }
+
+    in_dim = 7
+    out_dim = 11
+    in_dim_edges = 13
+    hidden_dims = [6, 6, 6, 6, 6]
+
+    g1 = dgl.graph((torch.tensor([0, 1, 2]), torch.tensor([1, 2, 3])))
+    g2 = dgl.graph((torch.tensor([0, 0, 0, 1]), torch.tensor([0, 1, 2, 0])))
+    g1.ndata["h"] = torch.zeros(g1.num_nodes(), in_dim, dtype=torch.float32)
+    g1.edata["e"] = torch.ones(g1.num_edges(), in_dim_edges, dtype=torch.float32)
+    g2.ndata["h"] = torch.ones(g2.num_nodes(), in_dim, dtype=torch.float32)
+    g2.edata["e"] = torch.zeros(g2.num_edges(), in_dim_edges, dtype=torch.float32)
+    bg = dgl.batch([g1, g2])
+    bg = dgl.add_self_loop(bg)
+
+    virtual_nodes = ["none", "mean", "sum"]
+    pna_kwargs = {"aggregators": ["mean", "max", "sum"], "scalers": ["identity", "amplification"]}
+
+    gnn_layers_kwargs = {
+        "gcn": {},
+        "gin": {},
+        "gat": {"num_heads": 3},
+        "gated-gcn": {"in_dim_edges": in_dim_edges, "hidden_dims_edges": hidden_dims},
+        "pna-conv": {**pna_kwargs},
+        "pna-msgpass#1": {**pna_kwargs, "in_dim_edges": 0},
+        "pna-msgpass#2": {**pna_kwargs, "in_dim_edges": in_dim_edges},
+    }
+
+    def test_full_network_densenet(self):
+
+        temp_dim_1 = 5
+        temp_dim_2 = 17
+
+        pre_nn_kwargs = dict(in_dim=self.in_dim, out_dim=temp_dim_1, hidden_dims=[4, 4, 4, 4, 4])
+
+        post_nn_kwargs = dict(in_dim=temp_dim_2, out_dim=self.out_dim, hidden_dims=[3, 3, 3, 3])
+
+        for pooling in [["none"], ["sum"], ["mean", "s2s", "max"]]:
+            for residual_skip_steps in [1, 2, 3]:
+                for virtual_node in self.virtual_nodes:
+                    for layer_name, this_kwargs in self.gnn_layers_kwargs.items():
+                        err_msg = f"pooling={pooling}, virtual_node={virtual_node}, layer_name={layer_name}, residual_skip_steps={residual_skip_steps}"
+                        layer_type = layer_name.split("#")[0]
+
+                        gnn_kwargs = dict(
+                            in_dim=temp_dim_1,
+                            out_dim=temp_dim_2,
+                            hidden_dims=self.hidden_dims,
+                            residual_type="densenet",
+                            residual_skip_steps=residual_skip_steps,
+                            layer_type=layer_type,
+                            pooling=pooling,
+                            **this_kwargs,
+                            **self.kwargs,
+                        )
+
+                        net = FullDGLNetwork(
+                            pre_nn_kwargs=pre_nn_kwargs, gnn_kwargs=gnn_kwargs, post_nn_kwargs=post_nn_kwargs
+                        )
+                        bg = deepcopy(self.bg)
+                        h_out = net.forward(bg)
+
+                        dim_1 = bg.num_nodes() if pooling == ["none"] else 1
+                        self.assertListEqual(list(h_out.shape), [dim_1, self.out_dim], msg=err_msg)
+
+
+if __name__ == "__main__":
+    ut.main()
diff --git a/tests/unit_tests/dgl/test_gnn_layers.py b/tests/unit_tests/dgl/test_gnn_layers.py
index ba987893c..be2d18223 100644
--- a/tests/unit_tests/dgl/test_gnn_layers.py
+++ b/tests/unit_tests/dgl/test_gnn_layers.py
@@ -36,12 +36,17 @@ class test_DGL_Layers(ut.TestCase):
     bg = dgl.batch([g1, g2])
     bg = dgl.add_self_loop(bg)
 
+    kwargs = {
+        "activation": "relu",
+        "dropout": 0.1,
+        "batch_norm": True,
+    }
+
     def test_gcnlayer(self):
 
         bg = deepcopy(self.bg)
         h_in = bg.ndata["h"]
-        e_in = bg.edata["e"]
-        layer = GCNLayer(in_dim=self.in_dim, out_dim=self.out_dim).to(float)
+        layer = GCNLayer(in_dim=self.in_dim, out_dim=self.out_dim, **self.kwargs).to(float)
 
         # Check the re-implementation of abstract methods
         self.assertFalse(layer.layer_supports_edges)
@@ -58,8 +63,7 @@ def test_ginlayer(self):
 
         bg = deepcopy(self.bg)
         h_in = bg.ndata["h"]
-        e_in = bg.edata["e"]
-        layer = GINLayer(in_dim=self.in_dim, out_dim=self.out_dim).to(float)
+        layer = GINLayer(in_dim=self.in_dim, out_dim=self.out_dim, **self.kwargs).to(float)
 
         # Check the re-implementation of abstract methods
         self.assertFalse(layer.layer_supports_edges)
@@ -77,8 +81,9 @@ def test_gatlayer(self):
         num_heads = 3
         bg = deepcopy(self.bg)
         h_in = bg.ndata["h"]
-        e_in = bg.edata["e"]
-        layer = GATLayer(in_dim=self.in_dim, out_dim=self.out_dim, num_heads=num_heads).to(float)
+        layer = GATLayer(in_dim=self.in_dim, out_dim=self.out_dim, num_heads=num_heads, **self.kwargs).to(
+            float
+        )
 
         # Check the re-implementation of abstract methods
         self.assertFalse(layer.layer_supports_edges)
@@ -101,6 +106,7 @@ def test_gatedgcnlayer(self):
             out_dim=self.out_dim,
             in_dim_edges=self.in_dim_edges,
             out_dim_edges=self.out_dim_edges,
+            **self.kwargs,
         ).to(float)
 
         # Check the re-implementation of abstract methods
@@ -129,7 +135,7 @@ def test_pnaconvolutionallayer(self):
         scalers = ["identity", "amplification", "attenuation"]
 
         layer = PNAConvolutionalLayer(
-            in_dim=self.in_dim, out_dim=self.out_dim, aggregators=aggregators, scalers=scalers
+            in_dim=self.in_dim, out_dim=self.out_dim, aggregators=aggregators, scalers=scalers, **self.kwargs
         ).to(float)
 
         # Check the re-implementation of abstract methods
@@ -152,7 +158,7 @@ def test_pnamessagepassinglayer(self):
         scalers = ["identity", "amplification", "attenuation"]
 
         layer = PNAMessagePassingLayer(
-            in_dim=self.in_dim, out_dim=self.out_dim, aggregators=aggregators, scalers=scalers
+            in_dim=self.in_dim, out_dim=self.out_dim, aggregators=aggregators, scalers=scalers, **self.kwargs
         ).to(float)
 
         # Check the re-implementation of abstract methods
@@ -174,6 +180,7 @@ def test_pnamessagepassinglayer(self):
             aggregators=aggregators,
             scalers=scalers,
             in_dim_edges=self.in_dim_edges,
+            **self.kwargs,
         ).to(float)
 
         # Check the re-implementation of abstract methods
diff --git a/tests/unit_tests/dgl/test_residual_connections.py b/tests/unit_tests/dgl/test_residual_connections.py
index a05a84982..a67ac076e 100644
--- a/tests/unit_tests/dgl/test_residual_connections.py
+++ b/tests/unit_tests/dgl/test_residual_connections.py
@@ -21,7 +21,7 @@ def test_get_true_out_dims_none(self):
         in_dims, out_dims = full_dims[:-1], full_dims[1:]
         rc = ResidualConnectionNone(skip_steps=1)
         true_out_dims = rc.get_true_out_dims(out_dims)
-        expected_out_dims = out_dims
+        expected_out_dims = out_dims[:-1]
 
         self.assertListEqual(expected_out_dims, true_out_dims)
 
@@ -44,7 +44,7 @@ def test_get_true_out_dims_simple(self):
         in_dims, out_dims = full_dims[:-1], full_dims[1:]
         rc = ResidualConnectionSimple(skip_steps=1)
         true_out_dims = rc.get_true_out_dims(out_dims)
-        expected_out_dims = out_dims
+        expected_out_dims = out_dims[:-1]
 
         self.assertListEqual(expected_out_dims, true_out_dims)
 
@@ -84,7 +84,7 @@ def test_get_true_out_dims_weighted(self):
         in_dims, out_dims = full_dims[:-1], full_dims[1:]
         rc = ResidualConnectionWeighted(skip_steps=1, out_dims=full_dims[1:])
         true_out_dims = rc.get_true_out_dims(out_dims)
-        expected_out_dims = out_dims
+        expected_out_dims = out_dims[:-1]
 
         self.assertListEqual(expected_out_dims, true_out_dims)
 
@@ -136,19 +136,19 @@ def test_get_true_out_dims_concat(self):
         # skip_steps=1
         rc = ResidualConnectionConcat(skip_steps=1)
         true_out_dims = rc.get_true_out_dims(out_dims)
-        expected_out_dims = [6, 14, 18, 22, 26, 30, 34, 38]
+        expected_out_dims = [6, 14, 18, 22, 26, 30, 34]
         self.assertListEqual(expected_out_dims, true_out_dims)
 
         # skip_steps=2
         rc = ResidualConnectionConcat(skip_steps=2)
         true_out_dims = rc.get_true_out_dims(out_dims)
-        expected_out_dims = [6, 8, 16, 12, 24, 16, 32, 20]
+        expected_out_dims = [6, 8, 16, 12, 24, 16, 32]
         self.assertListEqual(expected_out_dims, true_out_dims)
 
         # skip_steps=3
         rc = ResidualConnectionConcat(skip_steps=3)
         true_out_dims = rc.get_true_out_dims(out_dims)
-        expected_out_dims = [6, 8, 10, 18, 14, 16, 30, 20]
+        expected_out_dims = [6, 8, 10, 18, 14, 16, 30]
         self.assertListEqual(expected_out_dims, true_out_dims)
 
     def test_forward_concat(self):
@@ -189,19 +189,19 @@ def test_get_true_out_dims_densenet(self):
         # skip_steps=1
         rc = ResidualConnectionDenseNet(skip_steps=1)
         true_out_dims = rc.get_true_out_dims(out_dims)
-        expected_out_dims = np.cumsum(out_dims).tolist()
+        expected_out_dims = np.cumsum(out_dims).tolist()[:-1]
         self.assertListEqual(expected_out_dims, true_out_dims)
 
         # skip_steps=2
         rc = ResidualConnectionDenseNet(skip_steps=2)
         true_out_dims = rc.get_true_out_dims(out_dims)
-        expected_out_dims = [6, 8, 10 + 6, 12, 14 + 10 + 6, 16, 18 + 14 + 10 + 6, 20]
+        expected_out_dims = [6, 8, 10 + 6, 12, 14 + 10 + 6, 16, 18 + 14 + 10 + 6]
         self.assertListEqual(expected_out_dims, true_out_dims)
 
         # skip_steps=3
         rc = ResidualConnectionDenseNet(skip_steps=3)
         true_out_dims = rc.get_true_out_dims(out_dims)
-        expected_out_dims = [6, 8, 10, 12 + 6, 14, 16, 18 + 12 + 6, 20]
+        expected_out_dims = [6, 8, 10, 12 + 6, 14, 16, 18 + 12 + 6]
         self.assertListEqual(expected_out_dims, true_out_dims)
 
     def test_forward_densenet(self):