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

Issue:
Fix errors in the documentation
## Some block functions do not work with mkdocs, but used to work with Sphinx
- ..note::
- ..math::
- ..code-block::
- ..warning::

## Some sections also do not work

The first quote doesn't work, but the second seems to work. The same is true for `Math` and `See Also`

```
Example:
    This is an example that doesn't work.
```

```
Example
-----------
   This is an example that works, but is the older syntax.
```

## @classproperty does not appear in the docs

The methods that are defined as class properties with the decorator `@classproperty` does not appear in the rendered docs.





Discussion / hints:
Will look into that.
You will need to use the extended markdown syntax in general and not rst.  (quote for code, $$ for latex/mathjax)

Also `..note::`, etc should use admonitions (https://squidfunk.github.io/mkdocs-material/reference/admonitions/). 
Admonitions seem to be the solution for most of the problems. I will implement them in the next PR.
For code-blocks, we should use the notation with the triple `

The only problem that will remain is the methods with @classproperty decorators not showing in the docs
> For code-blocks, we should use the notation with the triple `

Yup, tripled ` as in markdown will work.


> The only problem that will remain is the methods with @classproperty decorators not showing in the docs

Since `classproperty ` is custom, I am not sure but I will have a look.
There's also the minor issue of defining the types in the docstring such that it appears in the parameters tables, since sometimes multiple types are accepted and cannot be added to the function definition.
![image](https://user-images.githubusercontent.com/47570400/107290583-cc198400-6a34-11eb-9295-78c5dd514d59.png)

I spotted this too. I'll fix it.

About multiple types you can use `Union` from the `typing` module.
Continuing the discussion in https://github.com/invivoai/goli/pull/17
I'm gonna close this issue since most of it is solved, and I opened this other issue #20 regarding the problem with `@classproperty` methods not appearing in the docs

Fix patch:
diff --git a/.gitignore b/.gitignore
index 7b8c7590c..1bbc1327d 100644
--- a/.gitignore
+++ b/.gitignore
@@ -29,6 +29,7 @@ goli/data/mega-pubchem/
 goli/data/cache/
 goli/data/b3lyp/
 goli/data/PCQM4Mv2/
+goli/data/PCQM4M/
 predictions/
 
 # Others
diff --git a/docs/images/datamodule.png b/docs/images/datamodule.png
index 73153fb54..d45d665c7 100644
Binary files a/docs/images/datamodule.png and b/docs/images/datamodule.png differ
diff --git a/docs/images/full_graph_network.png b/docs/images/full_graph_network.png
index e80ede227..3576f41ac 100644
Binary files a/docs/images/full_graph_network.png and b/docs/images/full_graph_network.png differ
diff --git a/docs/tutorials/basics/implementing_gnn_layers.ipynb b/docs/tutorials/basics/implementing_gnn_layers.ipynb
deleted file mode 100644
index d5420757c..000000000
--- a/docs/tutorials/basics/implementing_gnn_layers.ipynb
+++ /dev/null
@@ -1,325 +0,0 @@
-{
- "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.8.8"
-  },
-  "orig_nbformat": 2,
-  "kernelspec": {
-   "name": "python3",
-   "display_name": "Python 3.8.8 64-bit ('goli': conda)"
-  },
-  "interpreter": {
-   "hash": "f4a99d018a205fcbcc0480c84566beaebcb91b08d0414b39a842df533e2a1d25"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 2,
- "cells": [
-  {
-   "source": [
-    "# Creating GNN layers\n",
-    "\n",
-    "One of the primary advantage of the current library is the fact that GNN layers are independant from model architecture, thus allowing more flexibility with the code by easily swapping different layer types as hyper-parameters. However, this requires that the layers be implemented using the DGL library, and must be inherited from the class `BaseDGLLayer`, which standardizes the inputs, outputs and properties of the layers. Thus, the architecture can be handled independantly using the class `FeedForwardDGL`, or any similar custom class.\n",
-    "\n",
-    "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.nn.dgl_layers import BaseDGLLayer\n",
-    "from goli.nn.base_layers import FCLayer\n",
-    "from goli.utils.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 `FeedForwardDGL`, it needs to inherit from the class `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, normalization):\n",
-    "        # Initialize the parent class\n",
-    "        super().__init__(   in_dim=in_dim, out_dim=out_dim, activation=activation,\n",
-    "                            dropout=dropout, normalization=normalization)\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, normalization=\"batch_norm\").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, normalization):\n",
-    "        # Initialize the parent class\n",
-    "        super().__init__(   in_dim=in_dim, out_dim=out_dim, activation=activation,\n",
-    "                            dropout=dropout, normalization=normalization)\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, normalization=\"batch_norm\").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/tutorials/basics/parquet_reader.ipynb b/docs/tutorials/basics/parquet_reader.ipynb
deleted file mode 100644
index cc3f7706a..000000000
--- a/docs/tutorials/basics/parquet_reader.ipynb
+++ /dev/null
@@ -1,488 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "code",
-   "execution_count": 4,
-   "metadata": {},
-   "outputs": [
-    {
-     "ename": "NotFound",
-     "evalue": "404 GET https://storage.googleapis.com/download/storage/v1/b/gs://goli-private/datasets/pm6/o/Compound_000025001_000050000.parquet?alt=media: Not Found: ('Request failed with status code', 404, 'Expected one of', <HTTPStatus.OK: 200>, <HTTPStatus.PARTIAL_CONTENT: 206>)",
-     "output_type": "error",
-     "traceback": [
-      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
-      "\u001b[0;31mInvalidResponse\u001b[0m                           Traceback (most recent call last)",
-      "File \u001b[0;32m~/.venv/goli_ipu/lib/python3.8/site-packages/google/cloud/storage/client.py:1151\u001b[0m, in \u001b[0;36mClient.download_blob_to_file\u001b[0;34m(self, blob_or_uri, file_obj, start, end, raw_download, if_etag_match, if_etag_not_match, if_generation_match, if_generation_not_match, if_metageneration_match, if_metageneration_not_match, timeout, checksum, retry)\u001b[0m\n\u001b[1;32m   1150\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[0;32m-> 1151\u001b[0m     blob_or_uri\u001b[39m.\u001b[39;49m_do_download(\n\u001b[1;32m   1152\u001b[0m         transport,\n\u001b[1;32m   1153\u001b[0m         file_obj,\n\u001b[1;32m   1154\u001b[0m         download_url,\n\u001b[1;32m   1155\u001b[0m         headers,\n\u001b[1;32m   1156\u001b[0m         start,\n\u001b[1;32m   1157\u001b[0m         end,\n\u001b[1;32m   1158\u001b[0m         raw_download,\n\u001b[1;32m   1159\u001b[0m         timeout\u001b[39m=\u001b[39;49mtimeout,\n\u001b[1;32m   1160\u001b[0m         checksum\u001b[39m=\u001b[39;49mchecksum,\n\u001b[1;32m   1161\u001b[0m         retry\u001b[39m=\u001b[39;49mretry,\n\u001b[1;32m   1162\u001b[0m     )\n\u001b[1;32m   1163\u001b[0m \u001b[39mexcept\u001b[39;00m resumable_media\u001b[39m.\u001b[39mInvalidResponse \u001b[39mas\u001b[39;00m exc:\n",
-      "File \u001b[0;32m~/.venv/goli_ipu/lib/python3.8/site-packages/google/cloud/storage/blob.py:989\u001b[0m, in \u001b[0;36mBlob._do_download\u001b[0;34m(self, transport, file_obj, download_url, headers, start, end, raw_download, timeout, checksum, retry)\u001b[0m\n\u001b[1;32m    988\u001b[0m download\u001b[39m.\u001b[39m_retry_strategy \u001b[39m=\u001b[39m retry_strategy\n\u001b[0;32m--> 989\u001b[0m response \u001b[39m=\u001b[39m download\u001b[39m.\u001b[39;49mconsume(transport, timeout\u001b[39m=\u001b[39;49mtimeout)\n\u001b[1;32m    990\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_extract_headers_from_download(response)\n",
-      "File \u001b[0;32m~/.venv/goli_ipu/lib/python3.8/site-packages/google/resumable_media/requests/download.py:214\u001b[0m, in \u001b[0;36mDownload.consume\u001b[0;34m(self, transport, timeout)\u001b[0m\n\u001b[1;32m    212\u001b[0m     \u001b[39mreturn\u001b[39;00m result\n\u001b[0;32m--> 214\u001b[0m \u001b[39mreturn\u001b[39;00m _request_helpers\u001b[39m.\u001b[39;49mwait_and_retry(\n\u001b[1;32m    215\u001b[0m     retriable_request, \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_get_status_code, \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_retry_strategy\n\u001b[1;32m    216\u001b[0m )\n",
-      "File \u001b[0;32m~/.venv/goli_ipu/lib/python3.8/site-packages/google/resumable_media/requests/_request_helpers.py:148\u001b[0m, in \u001b[0;36mwait_and_retry\u001b[0;34m(func, get_status_code, retry_strategy)\u001b[0m\n\u001b[1;32m    147\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[0;32m--> 148\u001b[0m     response \u001b[39m=\u001b[39m func()\n\u001b[1;32m    149\u001b[0m \u001b[39mexcept\u001b[39;00m _CONNECTION_ERROR_CLASSES \u001b[39mas\u001b[39;00m e:\n",
-      "File \u001b[0;32m~/.venv/goli_ipu/lib/python3.8/site-packages/google/resumable_media/requests/download.py:207\u001b[0m, in \u001b[0;36mDownload.consume.<locals>.retriable_request\u001b[0;34m()\u001b[0m\n\u001b[1;32m    203\u001b[0m     \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_object_generation \u001b[39m=\u001b[39m _helpers\u001b[39m.\u001b[39m_parse_generation_header(\n\u001b[1;32m    204\u001b[0m         result, \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_get_headers\n\u001b[1;32m    205\u001b[0m     )\n\u001b[0;32m--> 207\u001b[0m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_process_response(result)\n\u001b[1;32m    209\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_stream \u001b[39mis\u001b[39;00m \u001b[39mnot\u001b[39;00m \u001b[39mNone\u001b[39;00m:\n",
-      "File \u001b[0;32m~/.venv/goli_ipu/lib/python3.8/site-packages/google/resumable_media/_download.py:188\u001b[0m, in \u001b[0;36mDownload._process_response\u001b[0;34m(self, response)\u001b[0m\n\u001b[1;32m    187\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_finished \u001b[39m=\u001b[39m \u001b[39mTrue\u001b[39;00m\n\u001b[0;32m--> 188\u001b[0m _helpers\u001b[39m.\u001b[39;49mrequire_status_code(\n\u001b[1;32m    189\u001b[0m     response, _ACCEPTABLE_STATUS_CODES, \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_get_status_code\n\u001b[1;32m    190\u001b[0m )\n",
-      "File \u001b[0;32m~/.venv/goli_ipu/lib/python3.8/site-packages/google/resumable_media/_helpers.py:105\u001b[0m, in \u001b[0;36mrequire_status_code\u001b[0;34m(response, status_codes, get_status_code, callback)\u001b[0m\n\u001b[1;32m    104\u001b[0m         callback()\n\u001b[0;32m--> 105\u001b[0m     \u001b[39mraise\u001b[39;00m common\u001b[39m.\u001b[39mInvalidResponse(\n\u001b[1;32m    106\u001b[0m         response,\n\u001b[1;32m    107\u001b[0m         \u001b[39m\"\u001b[39m\u001b[39mRequest failed with status code\u001b[39m\u001b[39m\"\u001b[39m,\n\u001b[1;32m    108\u001b[0m         status_code,\n\u001b[1;32m    109\u001b[0m         \u001b[39m\"\u001b[39m\u001b[39mExpected one of\u001b[39m\u001b[39m\"\u001b[39m,\n\u001b[1;32m    110\u001b[0m         \u001b[39m*\u001b[39mstatus_codes\n\u001b[1;32m    111\u001b[0m     )\n\u001b[1;32m    112\u001b[0m \u001b[39mreturn\u001b[39;00m status_code\n",
-      "\u001b[0;31mInvalidResponse\u001b[0m: ('Request failed with status code', 404, 'Expected one of', <HTTPStatus.OK: 200>, <HTTPStatus.PARTIAL_CONTENT: 206>)",
-      "\nDuring handling of the above exception, another exception occurred:\n",
-      "\u001b[0;31mNotFound\u001b[0m                                  Traceback (most recent call last)",
-      "\u001b[1;32m/home/oleksandr/code/goli/tests/parquet_reader.ipynb Cell 1\u001b[0m in \u001b[0;36m<cell line: 39>\u001b[0;34m()\u001b[0m\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2B193.57.88.53/home/oleksandr/code/goli/tests/parquet_reader.ipynb#ch0000000vscode-remote?line=23'>24</a>\u001b[0m     \u001b[39mprint\u001b[39m(\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2B193.57.88.53/home/oleksandr/code/goli/tests/parquet_reader.ipynb#ch0000000vscode-remote?line=24'>25</a>\u001b[0m         \u001b[39m\"\u001b[39m\u001b[39mFile \u001b[39m\u001b[39m{}\u001b[39;00m\u001b[39m read successfully  from Bucket  \u001b[39m\u001b[39m{}\u001b[39;00m\u001b[39m.\u001b[39m\u001b[39m\"\u001b[39m\u001b[39m.\u001b[39mformat(\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2B193.57.88.53/home/oleksandr/code/goli/tests/parquet_reader.ipynb#ch0000000vscode-remote?line=25'>26</a>\u001b[0m             destination_blob_name, bucket_name\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2B193.57.88.53/home/oleksandr/code/goli/tests/parquet_reader.ipynb#ch0000000vscode-remote?line=26'>27</a>\u001b[0m         )\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2B193.57.88.53/home/oleksandr/code/goli/tests/parquet_reader.ipynb#ch0000000vscode-remote?line=27'>28</a>\u001b[0m     )\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2B193.57.88.53/home/oleksandr/code/goli/tests/parquet_reader.ipynb#ch0000000vscode-remote?line=29'>30</a>\u001b[0m     \u001b[39m# files = os.listdir(path)\u001b[39;00m\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2B193.57.88.53/home/oleksandr/code/goli/tests/parquet_reader.ipynb#ch0000000vscode-remote?line=30'>31</a>\u001b[0m     \u001b[39m# print('files is', files)\u001b[39;00m\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2B193.57.88.53/home/oleksandr/code/goli/tests/parquet_reader.ipynb#ch0000000vscode-remote?line=31'>32</a>\u001b[0m \n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2B193.57.88.53/home/oleksandr/code/goli/tests/parquet_reader.ipynb#ch0000000vscode-remote?line=35'>36</a>\u001b[0m     \u001b[39m# # df = pd.read_csv(path, **kwargs)\u001b[39;00m\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2B193.57.88.53/home/oleksandr/code/goli/tests/parquet_reader.ipynb#ch0000000vscode-remote?line=36'>37</a>\u001b[0m     \u001b[39m# return df\u001b[39;00m\n\u001b[0;32m---> <a href='vscode-notebook-cell://ssh-remote%2B193.57.88.53/home/oleksandr/code/goli/tests/parquet_reader.ipynb#ch0000000vscode-remote?line=38'>39</a>\u001b[0m _read_parquet(\u001b[39m'\u001b[39;49m\u001b[39mgs://goli-private/datasets/pm6\u001b[39;49m\u001b[39m'\u001b[39;49m)\n",
-      "\u001b[1;32m/home/oleksandr/code/goli/tests/parquet_reader.ipynb Cell 1\u001b[0m in \u001b[0;36m_read_parquet\u001b[0;34m(path, **kwargs)\u001b[0m\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2B193.57.88.53/home/oleksandr/code/goli/tests/parquet_reader.ipynb#ch0000000vscode-remote?line=15'>16</a>\u001b[0m blob \u001b[39m=\u001b[39m bucket\u001b[39m.\u001b[39mblob(destination_blob_name)\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2B193.57.88.53/home/oleksandr/code/goli/tests/parquet_reader.ipynb#ch0000000vscode-remote?line=17'>18</a>\u001b[0m \u001b[39m# read as string\u001b[39;00m\n\u001b[0;32m---> <a href='vscode-notebook-cell://ssh-remote%2B193.57.88.53/home/oleksandr/code/goli/tests/parquet_reader.ipynb#ch0000000vscode-remote?line=18'>19</a>\u001b[0m read_output \u001b[39m=\u001b[39m blob\u001b[39m.\u001b[39;49mdownload_as_string()\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2B193.57.88.53/home/oleksandr/code/goli/tests/parquet_reader.ipynb#ch0000000vscode-remote?line=20'>21</a>\u001b[0m \u001b[39m# blobstring = blob_service.get_blob_to_text(CONTAINERNAME,BLOBNAME)\u001b[39;00m\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2B193.57.88.53/home/oleksandr/code/goli/tests/parquet_reader.ipynb#ch0000000vscode-remote?line=21'>22</a>\u001b[0m df \u001b[39m=\u001b[39m pd\u001b[39m.\u001b[39mread_csv(StringIO(read_output))\n",
-      "File \u001b[0;32m~/.venv/goli_ipu/lib/python3.8/site-packages/google/cloud/storage/blob.py:1516\u001b[0m, in \u001b[0;36mBlob.download_as_string\u001b[0;34m(self, client, start, end, raw_download, if_etag_match, if_etag_not_match, if_generation_match, if_generation_not_match, if_metageneration_match, if_metageneration_not_match, timeout, retry)\u001b[0m\n\u001b[1;32m   1434\u001b[0m \u001b[39m\"\"\"(Deprecated) Download the contents of this blob as a bytes object.\u001b[39;00m\n\u001b[1;32m   1435\u001b[0m \n\u001b[1;32m   1436\u001b[0m \u001b[39mIf :attr:`user_project` is set on the bucket, bills the API request\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m   1511\u001b[0m \u001b[39m:raises: :class:`google.cloud.exceptions.NotFound`\u001b[39;00m\n\u001b[1;32m   1512\u001b[0m \u001b[39m\"\"\"\u001b[39;00m\n\u001b[1;32m   1513\u001b[0m warnings\u001b[39m.\u001b[39mwarn(\n\u001b[1;32m   1514\u001b[0m     _DOWNLOAD_AS_STRING_DEPRECATED, \u001b[39mPendingDeprecationWarning\u001b[39;00m, stacklevel\u001b[39m=\u001b[39m\u001b[39m2\u001b[39m\n\u001b[1;32m   1515\u001b[0m )\n\u001b[0;32m-> 1516\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mdownload_as_bytes(\n\u001b[1;32m   1517\u001b[0m     client\u001b[39m=\u001b[39;49mclient,\n\u001b[1;32m   1518\u001b[0m     start\u001b[39m=\u001b[39;49mstart,\n\u001b[1;32m   1519\u001b[0m     end\u001b[39m=\u001b[39;49mend,\n\u001b[1;32m   1520\u001b[0m     raw_download\u001b[39m=\u001b[39;49mraw_download,\n\u001b[1;32m   1521\u001b[0m     if_etag_match\u001b[39m=\u001b[39;49mif_etag_match,\n\u001b[1;32m   1522\u001b[0m     if_etag_not_match\u001b[39m=\u001b[39;49mif_etag_not_match,\n\u001b[1;32m   1523\u001b[0m     if_generation_match\u001b[39m=\u001b[39;49mif_generation_match,\n\u001b[1;32m   1524\u001b[0m     if_generation_not_match\u001b[39m=\u001b[39;49mif_generation_not_match,\n\u001b[1;32m   1525\u001b[0m     if_metageneration_match\u001b[39m=\u001b[39;49mif_metageneration_match,\n\u001b[1;32m   1526\u001b[0m     if_metageneration_not_match\u001b[39m=\u001b[39;49mif_metageneration_not_match,\n\u001b[1;32m   1527\u001b[0m     timeout\u001b[39m=\u001b[39;49mtimeout,\n\u001b[1;32m   1528\u001b[0m     retry\u001b[39m=\u001b[39;49mretry,\n\u001b[1;32m   1529\u001b[0m )\n",
-      "File \u001b[0;32m~/.venv/goli_ipu/lib/python3.8/site-packages/google/cloud/storage/blob.py:1401\u001b[0m, in \u001b[0;36mBlob.download_as_bytes\u001b[0;34m(self, client, start, end, raw_download, if_etag_match, if_etag_not_match, if_generation_match, if_generation_not_match, if_metageneration_match, if_metageneration_not_match, timeout, checksum, retry)\u001b[0m\n\u001b[1;32m   1399\u001b[0m client \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_require_client(client)\n\u001b[1;32m   1400\u001b[0m string_buffer \u001b[39m=\u001b[39m BytesIO()\n\u001b[0;32m-> 1401\u001b[0m client\u001b[39m.\u001b[39;49mdownload_blob_to_file(\n\u001b[1;32m   1402\u001b[0m     \u001b[39mself\u001b[39;49m,\n\u001b[1;32m   1403\u001b[0m     string_buffer,\n\u001b[1;32m   1404\u001b[0m     start\u001b[39m=\u001b[39;49mstart,\n\u001b[1;32m   1405\u001b[0m     end\u001b[39m=\u001b[39;49mend,\n\u001b[1;32m   1406\u001b[0m     raw_download\u001b[39m=\u001b[39;49mraw_download,\n\u001b[1;32m   1407\u001b[0m     if_etag_match\u001b[39m=\u001b[39;49mif_etag_match,\n\u001b[1;32m   1408\u001b[0m     if_etag_not_match\u001b[39m=\u001b[39;49mif_etag_not_match,\n\u001b[1;32m   1409\u001b[0m     if_generation_match\u001b[39m=\u001b[39;49mif_generation_match,\n\u001b[1;32m   1410\u001b[0m     if_generation_not_match\u001b[39m=\u001b[39;49mif_generation_not_match,\n\u001b[1;32m   1411\u001b[0m     if_metageneration_match\u001b[39m=\u001b[39;49mif_metageneration_match,\n\u001b[1;32m   1412\u001b[0m     if_metageneration_not_match\u001b[39m=\u001b[39;49mif_metageneration_not_match,\n\u001b[1;32m   1413\u001b[0m     timeout\u001b[39m=\u001b[39;49mtimeout,\n\u001b[1;32m   1414\u001b[0m     checksum\u001b[39m=\u001b[39;49mchecksum,\n\u001b[1;32m   1415\u001b[0m     retry\u001b[39m=\u001b[39;49mretry,\n\u001b[1;32m   1416\u001b[0m )\n\u001b[1;32m   1417\u001b[0m \u001b[39mreturn\u001b[39;00m string_buffer\u001b[39m.\u001b[39mgetvalue()\n",
-      "File \u001b[0;32m~/.venv/goli_ipu/lib/python3.8/site-packages/google/cloud/storage/client.py:1164\u001b[0m, in \u001b[0;36mClient.download_blob_to_file\u001b[0;34m(self, blob_or_uri, file_obj, start, end, raw_download, if_etag_match, if_etag_not_match, if_generation_match, if_generation_not_match, if_metageneration_match, if_metageneration_not_match, timeout, checksum, retry)\u001b[0m\n\u001b[1;32m   1151\u001b[0m     blob_or_uri\u001b[39m.\u001b[39m_do_download(\n\u001b[1;32m   1152\u001b[0m         transport,\n\u001b[1;32m   1153\u001b[0m         file_obj,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m   1161\u001b[0m         retry\u001b[39m=\u001b[39mretry,\n\u001b[1;32m   1162\u001b[0m     )\n\u001b[1;32m   1163\u001b[0m \u001b[39mexcept\u001b[39;00m resumable_media\u001b[39m.\u001b[39mInvalidResponse \u001b[39mas\u001b[39;00m exc:\n\u001b[0;32m-> 1164\u001b[0m     _raise_from_invalid_response(exc)\n",
-      "File \u001b[0;32m~/.venv/goli_ipu/lib/python3.8/site-packages/google/cloud/storage/blob.py:4457\u001b[0m, in \u001b[0;36m_raise_from_invalid_response\u001b[0;34m(error)\u001b[0m\n\u001b[1;32m   4453\u001b[0m     error_message \u001b[39m=\u001b[39m \u001b[39mstr\u001b[39m(error)\n\u001b[1;32m   4455\u001b[0m message \u001b[39m=\u001b[39m \u001b[39mf\u001b[39m\u001b[39m\"\u001b[39m\u001b[39m{\u001b[39;00mresponse\u001b[39m.\u001b[39mrequest\u001b[39m.\u001b[39mmethod\u001b[39m}\u001b[39;00m\u001b[39m \u001b[39m\u001b[39m{\u001b[39;00mresponse\u001b[39m.\u001b[39mrequest\u001b[39m.\u001b[39murl\u001b[39m}\u001b[39;00m\u001b[39m: \u001b[39m\u001b[39m{\u001b[39;00merror_message\u001b[39m}\u001b[39;00m\u001b[39m\"\u001b[39m\n\u001b[0;32m-> 4457\u001b[0m \u001b[39mraise\u001b[39;00m exceptions\u001b[39m.\u001b[39mfrom_http_status(response\u001b[39m.\u001b[39mstatus_code, message, response\u001b[39m=\u001b[39mresponse)\n",
-      "\u001b[0;31mNotFound\u001b[0m: 404 GET https://storage.googleapis.com/download/storage/v1/b/gs://goli-private/datasets/pm6/o/Compound_000025001_000050000.parquet?alt=media: Not Found: ('Request failed with status code', 404, 'Expected one of', <HTTPStatus.OK: 200>, <HTTPStatus.PARTIAL_CONTENT: 206>)"
-     ]
-    }
-   ],
-   "source": [
-    "import os\n",
-    "import pyarrow as pa\n",
-    "import pyarrow.parquet as pq\n",
-    "import pandas as pd\n",
-    "from google.cloud import storage\n",
-    "from io import StringIO\n",
-    "\n",
-    "# TODO create funciton to read parquet, test from GCP storage (put it in path, should support gs and path, explore function from Pandas instead of parquet \"pq\")\n",
-    "def _read_parquet(path, **kwargs):\n",
-    "\n",
-    "    bucket_name = path\n",
-    "    destination_blob_name = 'Compound_000025001_000050000.parquet'\n",
-    "    \n",
-    "    storage_client = storage.Client()\n",
-    "    bucket = storage_client.bucket(bucket_name)\n",
-    "    blob = bucket.blob(destination_blob_name)\n",
-    " \n",
-    "    # read as string\n",
-    "    read_output = blob.download_as_string()\n",
-    "    df = pd.read_csv(StringIO(read_output))\n",
-    " \n",
-    "    print(\n",
-    "        \"File {} read successfully  from Bucket  {}.\".format(\n",
-    "            destination_blob_name, bucket_name\n",
-    "        )\n",
-    "    )\n",
-    "\n",
-    "_read_parquet('gs://goli-private/datasets/pm6')"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 1,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/html": [
-       "<div>\n",
-       "<style scoped>\n",
-       "    .dataframe tbody tr th:only-of-type {\n",
-       "        vertical-align: middle;\n",
-       "    }\n",
-       "\n",
-       "    .dataframe tbody tr th {\n",
-       "        vertical-align: top;\n",
-       "    }\n",
-       "\n",
-       "    .dataframe thead th {\n",
-       "        text-align: right;\n",
-       "    }\n",
-       "</style>\n",
-       "<table border=\"1\" class=\"dataframe\">\n",
-       "  <thead>\n",
-       "    <tr style=\"text-align: right;\">\n",
-       "      <th></th>\n",
-       "      <th>mol_id</th>\n",
-       "      <th>smiles</th>\n",
-       "      <th>A</th>\n",
-       "      <th>B</th>\n",
-       "      <th>C</th>\n",
-       "      <th>mu</th>\n",
-       "      <th>alpha</th>\n",
-       "      <th>homo</th>\n",
-       "      <th>lumo</th>\n",
-       "      <th>gap</th>\n",
-       "      <th>...</th>\n",
-       "      <th>zpve</th>\n",
-       "      <th>u0</th>\n",
-       "      <th>u298</th>\n",
-       "      <th>h298</th>\n",
-       "      <th>g298</th>\n",
-       "      <th>cv</th>\n",
-       "      <th>u0_atom</th>\n",
-       "      <th>u298_atom</th>\n",
-       "      <th>h298_atom</th>\n",
-       "      <th>g298_atom</th>\n",
-       "    </tr>\n",
-       "  </thead>\n",
-       "  <tbody>\n",
-       "    <tr>\n",
-       "      <th>0</th>\n",
-       "      <td>gdb_1</td>\n",
-       "      <td>C</td>\n",
-       "      <td>157.71180</td>\n",
-       "      <td>157.709970</td>\n",
-       "      <td>157.706990</td>\n",
-       "      <td>0.0000</td>\n",
-       "      <td>13.21</td>\n",
-       "      <td>-0.3877</td>\n",
-       "      <td>0.1171</td>\n",
-       "      <td>0.5048</td>\n",
-       "      <td>...</td>\n",
-       "      <td>0.044749</td>\n",
-       "      <td>-40.478930</td>\n",
-       "      <td>-40.476062</td>\n",
-       "      <td>-40.475117</td>\n",
-       "      <td>-40.498597</td>\n",
-       "      <td>6.469</td>\n",
-       "      <td>-395.999595</td>\n",
-       "      <td>-398.643290</td>\n",
-       "      <td>-401.014647</td>\n",
-       "      <td>-372.471772</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>1</th>\n",
-       "      <td>gdb_2</td>\n",
-       "      <td>N</td>\n",
-       "      <td>293.60975</td>\n",
-       "      <td>293.541110</td>\n",
-       "      <td>191.393970</td>\n",
-       "      <td>1.6256</td>\n",
-       "      <td>9.46</td>\n",
-       "      <td>-0.2570</td>\n",
-       "      <td>0.0829</td>\n",
-       "      <td>0.3399</td>\n",
-       "      <td>...</td>\n",
-       "      <td>0.034358</td>\n",
-       "      <td>-56.525887</td>\n",
-       "      <td>-56.523026</td>\n",
-       "      <td>-56.522082</td>\n",
-       "      <td>-56.544961</td>\n",
-       "      <td>6.316</td>\n",
-       "      <td>-276.861363</td>\n",
-       "      <td>-278.620271</td>\n",
-       "      <td>-280.399259</td>\n",
-       "      <td>-259.338802</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>2</th>\n",
-       "      <td>gdb_3</td>\n",
-       "      <td>O</td>\n",
-       "      <td>799.58812</td>\n",
-       "      <td>437.903860</td>\n",
-       "      <td>282.945450</td>\n",
-       "      <td>1.8511</td>\n",
-       "      <td>6.31</td>\n",
-       "      <td>-0.2928</td>\n",
-       "      <td>0.0687</td>\n",
-       "      <td>0.3615</td>\n",
-       "      <td>...</td>\n",
-       "      <td>0.021375</td>\n",
-       "      <td>-76.404702</td>\n",
-       "      <td>-76.401867</td>\n",
-       "      <td>-76.400922</td>\n",
-       "      <td>-76.422349</td>\n",
-       "      <td>6.002</td>\n",
-       "      <td>-213.087624</td>\n",
-       "      <td>-213.974294</td>\n",
-       "      <td>-215.159658</td>\n",
-       "      <td>-201.407171</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>3</th>\n",
-       "      <td>gdb_4</td>\n",
-       "      <td>C#C</td>\n",
-       "      <td>0.00000</td>\n",
-       "      <td>35.610036</td>\n",
-       "      <td>35.610036</td>\n",
-       "      <td>0.0000</td>\n",
-       "      <td>16.28</td>\n",
-       "      <td>-0.2845</td>\n",
-       "      <td>0.0506</td>\n",
-       "      <td>0.3351</td>\n",
-       "      <td>...</td>\n",
-       "      <td>0.026841</td>\n",
-       "      <td>-77.308427</td>\n",
-       "      <td>-77.305527</td>\n",
-       "      <td>-77.304583</td>\n",
-       "      <td>-77.327429</td>\n",
-       "      <td>8.574</td>\n",
-       "      <td>-385.501997</td>\n",
-       "      <td>-387.237686</td>\n",
-       "      <td>-389.016047</td>\n",
-       "      <td>-365.800724</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>4</th>\n",
-       "      <td>gdb_5</td>\n",
-       "      <td>C#N</td>\n",
-       "      <td>0.00000</td>\n",
-       "      <td>44.593883</td>\n",
-       "      <td>44.593883</td>\n",
-       "      <td>2.8937</td>\n",
-       "      <td>12.99</td>\n",
-       "      <td>-0.3604</td>\n",
-       "      <td>0.0191</td>\n",
-       "      <td>0.3796</td>\n",
-       "      <td>...</td>\n",
-       "      <td>0.016601</td>\n",
-       "      <td>-93.411888</td>\n",
-       "      <td>-93.409370</td>\n",
-       "      <td>-93.408425</td>\n",
-       "      <td>-93.431246</td>\n",
-       "      <td>6.278</td>\n",
-       "      <td>-301.820534</td>\n",
-       "      <td>-302.906752</td>\n",
-       "      <td>-304.091489</td>\n",
-       "      <td>-288.720028</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>...</th>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>1001</th>\n",
-       "      <td>gdb_94823</td>\n",
-       "      <td>CC1OCC2(CC2)C=C1</td>\n",
-       "      <td>3.73934</td>\n",
-       "      <td>1.156580</td>\n",
-       "      <td>1.019880</td>\n",
-       "      <td>1.0639</td>\n",
-       "      <td>86.86</td>\n",
-       "      <td>-0.2291</td>\n",
-       "      <td>0.0210</td>\n",
-       "      <td>0.2501</td>\n",
-       "      <td>...</td>\n",
-       "      <td>0.183181</td>\n",
-       "      <td>-387.089386</td>\n",
-       "      <td>-387.080974</td>\n",
-       "      <td>-387.080030</td>\n",
-       "      <td>-387.122163</td>\n",
-       "      <td>33.410</td>\n",
-       "      <td>-2037.744489</td>\n",
-       "      <td>-2051.130511</td>\n",
-       "      <td>-2062.985411</td>\n",
-       "      <td>-1895.686491</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>1002</th>\n",
-       "      <td>gdb_39986</td>\n",
-       "      <td>C1CC11CC2C3OC2C13</td>\n",
-       "      <td>4.01995</td>\n",
-       "      <td>1.375230</td>\n",
-       "      <td>1.276030</td>\n",
-       "      <td>2.0583</td>\n",
-       "      <td>78.65</td>\n",
-       "      <td>-0.2263</td>\n",
-       "      <td>0.0855</td>\n",
-       "      <td>0.3118</td>\n",
-       "      <td>...</td>\n",
-       "      <td>0.160873</td>\n",
-       "      <td>-385.816258</td>\n",
-       "      <td>-385.809490</td>\n",
-       "      <td>-385.808546</td>\n",
-       "      <td>-385.847235</td>\n",
-       "      <td>28.507</td>\n",
-       "      <td>-1866.696830</td>\n",
-       "      <td>-1879.337372</td>\n",
-       "      <td>-1890.006280</td>\n",
-       "      <td>-1736.880279</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>1003</th>\n",
-       "      <td>gdb_129889</td>\n",
-       "      <td>c1c(nc2n1cno2)N</td>\n",
-       "      <td>5.52693</td>\n",
-       "      <td>1.335230</td>\n",
-       "      <td>1.076560</td>\n",
-       "      <td>4.4714</td>\n",
-       "      <td>67.81</td>\n",
-       "      <td>-0.2034</td>\n",
-       "      <td>-0.0173</td>\n",
-       "      <td>0.1861</td>\n",
-       "      <td>...</td>\n",
-       "      <td>0.092058</td>\n",
-       "      <td>-448.922474</td>\n",
-       "      <td>-448.915855</td>\n",
-       "      <td>-448.914911</td>\n",
-       "      <td>-448.953258</td>\n",
-       "      <td>25.857</td>\n",
-       "      <td>-1339.274261</td>\n",
-       "      <td>-1346.674475</td>\n",
-       "      <td>-1353.785407</td>\n",
-       "      <td>-1248.217057</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>1004</th>\n",
-       "      <td>gdb_47001</td>\n",
-       "      <td>O=C1NCC2COC1C2</td>\n",
-       "      <td>2.67679</td>\n",
-       "      <td>1.787180</td>\n",
-       "      <td>1.428170</td>\n",
-       "      <td>4.8266</td>\n",
-       "      <td>70.94</td>\n",
-       "      <td>-0.2282</td>\n",
-       "      <td>0.0241</td>\n",
-       "      <td>0.2522</td>\n",
-       "      <td>...</td>\n",
-       "      <td>0.151811</td>\n",
-       "      <td>-439.131708</td>\n",
-       "      <td>-439.124600</td>\n",
-       "      <td>-439.123656</td>\n",
-       "      <td>-439.163299</td>\n",
-       "      <td>28.066</td>\n",
-       "      <td>-1779.364520</td>\n",
-       "      <td>-1790.901901</td>\n",
-       "      <td>-1800.977813</td>\n",
-       "      <td>-1656.054609</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>1005</th>\n",
-       "      <td>gdb_53123</td>\n",
-       "      <td>CC#CC1C2C3CN1C23</td>\n",
-       "      <td>4.54414</td>\n",
-       "      <td>0.981090</td>\n",
-       "      <td>0.944760</td>\n",
-       "      <td>1.7549</td>\n",
-       "      <td>84.57</td>\n",
-       "      <td>-0.2224</td>\n",
-       "      <td>0.0357</td>\n",
-       "      <td>0.2581</td>\n",
-       "      <td>...</td>\n",
-       "      <td>0.147781</td>\n",
-       "      <td>-364.706764</td>\n",
-       "      <td>-364.698604</td>\n",
-       "      <td>-364.697660</td>\n",
-       "      <td>-364.740861</td>\n",
-       "      <td>30.373</td>\n",
-       "      <td>-1786.060041</td>\n",
-       "      <td>-1796.938537</td>\n",
-       "      <td>-1807.014449</td>\n",
-       "      <td>-1664.834715</td>\n",
-       "    </tr>\n",
-       "  </tbody>\n",
-       "</table>\n",
-       "<p>1006 rows × 21 columns</p>\n",
-       "</div>"
-      ],
-      "text/plain": [
-       "          mol_id             smiles          A           B           C  \\\n",
-       "0          gdb_1                  C  157.71180  157.709970  157.706990   \n",
-       "1          gdb_2                  N  293.60975  293.541110  191.393970   \n",
-       "2          gdb_3                  O  799.58812  437.903860  282.945450   \n",
-       "3          gdb_4                C#C    0.00000   35.610036   35.610036   \n",
-       "4          gdb_5                C#N    0.00000   44.593883   44.593883   \n",
-       "...          ...                ...        ...         ...         ...   \n",
-       "1001   gdb_94823   CC1OCC2(CC2)C=C1    3.73934    1.156580    1.019880   \n",
-       "1002   gdb_39986  C1CC11CC2C3OC2C13    4.01995    1.375230    1.276030   \n",
-       "1003  gdb_129889    c1c(nc2n1cno2)N    5.52693    1.335230    1.076560   \n",
-       "1004   gdb_47001     O=C1NCC2COC1C2    2.67679    1.787180    1.428170   \n",
-       "1005   gdb_53123   CC#CC1C2C3CN1C23    4.54414    0.981090    0.944760   \n",
-       "\n",
-       "          mu  alpha    homo    lumo     gap  ...      zpve          u0  \\\n",
-       "0     0.0000  13.21 -0.3877  0.1171  0.5048  ...  0.044749  -40.478930   \n",
-       "1     1.6256   9.46 -0.2570  0.0829  0.3399  ...  0.034358  -56.525887   \n",
-       "2     1.8511   6.31 -0.2928  0.0687  0.3615  ...  0.021375  -76.404702   \n",
-       "3     0.0000  16.28 -0.2845  0.0506  0.3351  ...  0.026841  -77.308427   \n",
-       "4     2.8937  12.99 -0.3604  0.0191  0.3796  ...  0.016601  -93.411888   \n",
-       "...      ...    ...     ...     ...     ...  ...       ...         ...   \n",
-       "1001  1.0639  86.86 -0.2291  0.0210  0.2501  ...  0.183181 -387.089386   \n",
-       "1002  2.0583  78.65 -0.2263  0.0855  0.3118  ...  0.160873 -385.816258   \n",
-       "1003  4.4714  67.81 -0.2034 -0.0173  0.1861  ...  0.092058 -448.922474   \n",
-       "1004  4.8266  70.94 -0.2282  0.0241  0.2522  ...  0.151811 -439.131708   \n",
-       "1005  1.7549  84.57 -0.2224  0.0357  0.2581  ...  0.147781 -364.706764   \n",
-       "\n",
-       "            u298        h298        g298      cv      u0_atom    u298_atom  \\\n",
-       "0     -40.476062  -40.475117  -40.498597   6.469  -395.999595  -398.643290   \n",
-       "1     -56.523026  -56.522082  -56.544961   6.316  -276.861363  -278.620271   \n",
-       "2     -76.401867  -76.400922  -76.422349   6.002  -213.087624  -213.974294   \n",
-       "3     -77.305527  -77.304583  -77.327429   8.574  -385.501997  -387.237686   \n",
-       "4     -93.409370  -93.408425  -93.431246   6.278  -301.820534  -302.906752   \n",
-       "...          ...         ...         ...     ...          ...          ...   \n",
-       "1001 -387.080974 -387.080030 -387.122163  33.410 -2037.744489 -2051.130511   \n",
-       "1002 -385.809490 -385.808546 -385.847235  28.507 -1866.696830 -1879.337372   \n",
-       "1003 -448.915855 -448.914911 -448.953258  25.857 -1339.274261 -1346.674475   \n",
-       "1004 -439.124600 -439.123656 -439.163299  28.066 -1779.364520 -1790.901901   \n",
-       "1005 -364.698604 -364.697660 -364.740861  30.373 -1786.060041 -1796.938537   \n",
-       "\n",
-       "        h298_atom    g298_atom  \n",
-       "0     -401.014647  -372.471772  \n",
-       "1     -280.399259  -259.338802  \n",
-       "2     -215.159658  -201.407171  \n",
-       "3     -389.016047  -365.800724  \n",
-       "4     -304.091489  -288.720028  \n",
-       "...           ...          ...  \n",
-       "1001 -2062.985411 -1895.686491  \n",
-       "1002 -1890.006280 -1736.880279  \n",
-       "1003 -1353.785407 -1248.217057  \n",
-       "1004 -1800.977813 -1656.054609  \n",
-       "1005 -1807.014449 -1664.834715  \n",
-       "\n",
-       "[1006 rows x 21 columns]"
-      ]
-     },
-     "execution_count": 1,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "import pandas as pd\n",
-    "import goli\n",
-    "import os\n",
-    "from os.path import dirname, abspath\n",
-    "\n",
-    "MAIN_DIR = dirname(dirname(abspath(goli.__file__)))\n",
-    "\n",
-    "filename = MAIN_DIR + '/goli/data/QM9/micro_qm9'\n",
-    "df = pd.read_parquet(f'{filename}.parquet')\n",
-    "\n",
-    "df"
-   ]
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "Python 3.8.10 ('goli_ipu')",
-   "language": "python",
-   "name": "python3"
-  },
-  "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.8.10"
-  },
-  "orig_nbformat": 4,
-  "vscode": {
-   "interpreter": {
-    "hash": "b813b16c721ca8d9e65e6e6945a38a6ae48015ae799c455bab639da90d3bb278"
-   }
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 2
-}
diff --git a/docs/tutorials/feature_processing/add_new_positional_encoding.ipynb b/docs/tutorials/feature_processing/add_new_positional_encoding.ipynb
new file mode 100644
index 000000000..d4b6d2f56
--- /dev/null
+++ b/docs/tutorials/feature_processing/add_new_positional_encoding.ipynb
@@ -0,0 +1,296 @@
+{
+ "cells": [
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Add new positional encoding\n",
+    "\n",
+    "One of the main advantage of this library is the ability to easily incorporate novel positional encodings on the node, edge and graph level. The positional encodings are computed and feed into respective encoders and then the hidden embeddings from all pe encoders are pooled (according to if they are node, edge, or graph level) and then feed into the GNN layers as features. The designs allow any combination of positional encodings to be used by modifying the configuration file. For more details on the data processing part, please visit the [design page of the doc](https://valence-discovery.github.io/goli/design.html).\n",
+    "\n",
+    "Here is the workflow for computing and processing positional encoding in the library:\n",
+    "1. edit related parts in the yaml configuration file\n",
+    "\n",
+    "2. compute the raw positional encoding from the graph in [`goli/features/positional_encoding.py`](https://valence-discovery.github.io/goli/api/goli.features.html#goli.features.positional_encoding) (from the [`graph positional encoder`](https://valence-discovery.github.io/goli/api/goli.features.html#goli.features.positional_encoding.graph_positional_encoder))\n",
+    "\n",
+    "3. feed the raw positional encoding into the respective (specialized) encoders in [`goli/nn/encoders`](https://valence-discovery.github.io/goli/api/goli.nn/encoders.html). For example, a simple [`MLP positional encoder`](https://valence-discovery.github.io/goli/api/goli.nn/encoders.html#goli.nn.encoders.mlp_encoder) can be found. \n",
+    "\n",
+    "4. Output the hidden embeddings of pe from the encoders in their respective output keys: `feat`(node feature), `edge_feat`(edge feature), `graph_feat`(graph feature) and potentially other keys if needed such as `nodepair_feat` \n",
+    "\n",
+    "5. pool the hidden embeddings with same keys together: for example, all output with `feat` key will be pooled together\n",
+    "\n",
+    "6. Construct the [`PyG Batch`](https://pytorch-geometric.readthedocs.io/en/latest/generated/torch_geometric.data.Batch.html#torch_geometric.data.Batch), batch of graphs, each contain the output keys seen above, ready for use in the [GNN layers](https://valence-discovery.github.io/goli/api/goli.nn/pyg_layers.html) \n",
+    "\n",
+    "Since this library is built using PyG, we recommend looking at their [Docs](https://pytorch-geometric.readthedocs.io/en/latest/) and [Tutorials](https://pytorch-geometric.readthedocs.io/en/latest/get_started/introduction.html) for more info. \n",
+    "\n",
+    "We start by editing the configuration file first.\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Table of content\n",
+    "1. [edit the config file](#Edit-the-yaml-Configuration-File)\n",
+    "2. [compute the pe from mol](#Compute-the-Positional-Encoding)\n",
+    "3. [add existing encoder](#Add-Existing-Encoder)\n",
+    "4. [add specialized encoder](#Add-Specialized-Encoder)\n",
+    "5. [add the keys to spaces](#Add-the-Keys-to-Spaces)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Edit the yaml Configuration File\n",
+    "\n",
+    "### Computing Raw PE\n",
+    "We will use the degree of each node as a positional encoding in this tutorial. \n",
+    "First start with an existing yaml configuration file, you can find them in `expts/configs`\n",
+    "\n",
+    "We first look at where in the yaml file is the raw positional encodings computed. `deg_pos` is added as an example below. You can add relevant arguments for computing the positional encoding here as well such as `normalize` in the example."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "```\n",
+    "pos_encoding_as_features:\n",
+    "    pos_types:\n",
+    "      deg_pos: #example, degree centrality\n",
+    "        pos_type: degree\n",
+    "        normalize: False\n",
+    "```"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Specifying Encoders for the PE\n",
+    "Now we want to specify arguments for the encoders associated with the pe"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "```\n",
+    "pe_encoders:\n",
+    "    out_dim: 64\n",
+    "    pool: \"sum\" #choice of pooling across multiple pe encoders\n",
+    "    last_norm: None #\"batch_norm\", \"layer_norm\"\n",
+    "    encoders: \n",
+    "      deg_pos: #same name from the previous cell\n",
+    "        encoder_type: \"mlp\" #or you can specify your own specialized encoder\n",
+    "        input_keys: [\"degree\"] #same as the pos_type configured before\n",
+    "        output_keys: [\"feat\"] #node feature\n",
+    "        hidden_dim: 64\n",
+    "        num_layers: 1\n",
+    "        dropout: 0.1\n",
+    "        normalization: \"none\"   #\"batch_norm\" or \"layer_norm\"\n",
+    "        first_normalization: \"layer_norm\"   #\"batch_norm\" or \"layer_norm\"\n",
+    "```"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Compute the Positional Encoding\n",
+    "Next, we want to compute the raw degree of each node from the molecule graph.\n",
+    "\n",
+    "### add function to compute the pe\n",
+    "Go to [goli/features](https://valence-discovery.github.io/goli/api/goli.features.html) and add a new file `deg.py` to add the function to compute the pe. "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from typing import Tuple, Union, Optional\n",
+    "\n",
+    "from scipy import sparse\n",
+    "from scipy.sparse import spmatrix\n",
+    "import numpy as np\n",
+    "\n",
+    "def compute_deg(adj: Union[np.ndarray, spmatrix], normalize: bool) -> np.ndarray:\n",
+    "    \"\"\"\n",
+    "    Compute the node degree positional encoding \n",
+    "\n",
+    "    Parameters:\n",
+    "        adj: Adjacency matrix\n",
+    "        normalize: indicate if the degree across all nodes are normalized to [0,1] or not\n",
+    "    Returns:\n",
+    "        2D array with shape (num_nodes, 1) specifying (outgoing) degree for each node\n",
+    "    \"\"\"\n",
+    "    \n",
+    "    #first adj convert to scipy sparse matrix if not already\n",
+    "    if type(adj) is np.ndarray:\n",
+    "        adj = sparse.csr_matrix(adj)\n",
+    "    \n",
+    "    #https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.csr_matrix.sum.html\n",
+    "    degs = adj.sum(axis=0) #sum over each row\n",
+    "    \n",
+    "    if (normalize): #normalize the degree sequence to [0,1]\n",
+    "        degs = degs / np.max(degs)\n",
+    "    return degs"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Test with toy matrix\n",
+    "\n",
+    "here we will test if our code compute the degrees of each node correctly"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "matrix([[1., 1., 1., 1., 1.]])"
+      ]
+     },
+     "execution_count": 2,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "adj = np.identity(5) #make an identity matrix\n",
+    "normalize = True\n",
+    "\n",
+    "degs = compute_deg(adj, normalize=normalize)\n",
+    "\n",
+    "degs"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### add to positional_encoding.py\n",
+    "\n",
+    "To compute the new pe along with all existing pe, we need to add the function we wrote to [`goli/feature/positional_encoding.py`](https://valence-discovery.github.io/goli/api/goli.features.html#goli.features.positional_encoding). Modify the [`graph_positional_encoder`](https://valence-discovery.github.io/goli/api/goli.features.html#goli.features.positional_encoding.graph_positional_encoder) function by adding `pos_type == \"degree\"` logic "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Add Existing Encoder\n",
+    "\n",
+    "In order to pool over all the positional encodings, we need to add encoder to process the raw computed positional encoding and ensure the output dimension from all pe encoders are the same. When designing the encoder, you can either use an existing encoder or write a specialized encoder you made\n",
+    "\n",
+    "here we can simply specify [`MLPEncoder`](https://valence-discovery.github.io/goli/api/goli.nn/encoders.html#goli.nn.encoders.mlp_encoder.MLPEncoder) in the yaml file and the library will automatically feed the raw positional encoding to a mlp encoder based on the input arguments. Note that in this example, the encoder takes in the pe stored at the input key `degree` and then outputs to the output key `feat`"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "```\n",
+    "encoders: \n",
+    "  deg_pos: \n",
+    "    encoder_type: \"mlp\" \n",
+    "    input_keys: [\"degree\"] \n",
+    "    output_keys: [\"feat\"] # node feature\n",
+    "    hidden_dim: 64\n",
+    "    num_layers: 1\n",
+    "    dropout: 0.1\n",
+    "    normalization: \"none\"   #\"batch_norm\" or \"layer_norm\"\n",
+    "    first_normalization: \"layer_norm\"   #\"batch_norm\" or \"layer_norm\"\n",
+    "```"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Add Specialized Encoder\n",
+    "\n",
+    "You can also add specialized encoder, such as `laplacian_pe` for the laplacian eigenvectors and eigenvalues. Here, we can add a new `deg_pos_encoder.py` in [`goli/nn/encoders`](https://valence-discovery.github.io/goli/api/goli.nn/encoders.html). As an example and template, please see the [`MLPEncoder`](https://valence-discovery.github.io/goli/api/goli.nn/encoders.html#goli.nn.encoders.mlp_encoder.MLPEncoder)\n",
+    "\n",
+    "Note that all new encoders must inherent from [`BaseEncoder`](https://valence-discovery.github.io/goli/api/goli.nn/encoders.html#goli.nn.encoders.base_encoder.BaseEncoder) class and implement the following abstract methods\n",
+    "\n",
+    "- [`forward`](https://valence-discovery.github.io/goli/api/goli.nn/encoders.html#goli.nn.encoders.base_encoder.BaseEncoder.forward): the forward function of the encoder, how to process the input\n",
+    "\n",
+    "- [`parse_input_keys`](https://valence-discovery.github.io/goli/api/goli.nn/encoders.html#goli.nn.encoders.base_encoder.BaseEncoder.parse_input_keys): how to parse the input keys\n",
+    "\n",
+    "- [`parse_output_keys`](https://valence-discovery.github.io/goli/api/goli.nn/encoders.html#goli.nn.encoders.laplace_pos_encoder.LapPENodeEncoder.parse_output_keys): how to parse the output keys\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Add the Keys to Spaces\n",
+    "\n",
+    "In order to directly find the correct encoders from the yaml file, we need to specify which key corresponding to what class. \n",
+    "\n",
+    "- add our new `deg_pos_encoder` to `goli/utils/spaces.py` in the `PE_ENCODERS_DICT`\n",
+    "- add our new `deg_pos_encoder` to [`goli/nn/architectures/encoder_manager.py`](https://valence-discovery.github.io/goli/api/goli.nn/architectures.html#goli.nn.architectures.encoder_manager.EncoderManager) in the `PE_ENCODERS_DICT`\n",
+    "- add the import of our encoder to  `goli/nn/encoders/__init__.py`\n",
+    "\n",
+    "Now we can modify the yaml file to use our new encoder"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "```\n",
+    "encoders: \n",
+    "  deg_pos: \n",
+    "    encoder_type: \"deg_pos_encoder\" \n",
+    "    input_keys: [\"degree\"] \n",
+    "    output_keys: [\"feat\"] # node feature\n",
+    "    hidden_dim: 64\n",
+    "    #any other keys that might be used for initialization\n",
+    "```"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "interpreter": {
+   "hash": "f4a99d018a205fcbcc0480c84566beaebcb91b08d0414b39a842df533e2a1d25"
+  },
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "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.8.10"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/docs/tutorials/basics/csv_to_parquet.ipynb b/docs/tutorials/feature_processing/csv_to_parquet.ipynb
similarity index 100%
rename from docs/tutorials/basics/csv_to_parquet.ipynb
rename to docs/tutorials/feature_processing/csv_to_parquet.ipynb
diff --git a/docs/tutorials/basics/timing_parallel.ipynb b/docs/tutorials/feature_processing/timing_parallel.ipynb
similarity index 96%
rename from docs/tutorials/basics/timing_parallel.ipynb
rename to docs/tutorials/feature_processing/timing_parallel.ipynb
index f4931874c..c0c6af763 100644
--- a/docs/tutorials/basics/timing_parallel.ipynb
+++ b/docs/tutorials/feature_processing/timing_parallel.ipynb
@@ -1,5 +1,18 @@
 {
  "cells": [
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "id": "0693c5c9",
+   "metadata": {},
+   "source": [
+    "# Timing parallel processing\n",
+    "\n",
+    "This tutorial is meant to help you benchmark different parallel processing methods for the processing of molecules into graphs. This will allow you to chose the one most suitable for your machine, since the benchmarks vary per machine.\n",
+    "\n",
+    "In general, we find that using `joblib` with the `loky` parallel processing and a batch size of `1000` is most beneficial. The logic is abstracted into `datamol.parallelized_with_batches`"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 1,
@@ -262,16 +275,6 @@
     "        benchmark.append(datum)"
    ]
   },
-  {
-   "cell_type": "markdown",
-   "id": "cc3f7256-3f92-4bf7-ae96-010efe5f060c",
-   "metadata": {
-    "tags": []
-   },
-   "source": [
-    "# pandarallel"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": 7,
diff --git a/docs/tutorials/gnn/add_new_gnn_layers.ipynb b/docs/tutorials/gnn/add_new_gnn_layers.ipynb
new file mode 100644
index 000000000..b58eef584
--- /dev/null
+++ b/docs/tutorials/gnn/add_new_gnn_layers.ipynb
@@ -0,0 +1,575 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Creating GNN layers\n",
+    "\n",
+    "One of the primary advantage of the current library is the fact that GNN layers are independant from model architecture, thus allowing more flexibility with the code by easily swapping different layer types as hyper-parameters. However, this requires that the layers be implemented using the PyG library, and must be inherited from the class [`BaseGraphStructure`](https://valence-discovery.github.io/goli/api/goli.nn/goli.nn.html#goli.nn.base_graph_layer.BaseGraphStructure), which standardizes the inputs, outputs and properties of the layers. Thus, the architecture can be handled independantly using the class [`FeedForwardPyg`](https://valence-discovery.github.io/goli/api/goli.nn/architectures.html#goli.nn.architectures.pyg_architectures.FeedForwardPyg) or other custom classes.\n",
+    "\n",
+    "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 PyG, we recommend looking at their [Docs](https://pytorch-geometric.readthedocs.io/en/latest/) and [Tutorials](https://pytorch-geometric.readthedocs.io/en/latest/get_started/introduction.html) for more info. "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Table of content\n",
+    "1. [define test graph](#Define-synthetic-test-graph)\n",
+    "2. [create simple layer](#Creating-a-simple-layer)\n",
+    "3. [test simple layer](#Test-our-simple-layer)\n",
+    "4. [create complex layer](#Creating-a-complex-layer-with-edges)\n",
+    "5. [test complex layer](#Test-our-complex-layer) "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%load_ext autoreload\n",
+    "%autoreload 2\n",
+    "\n",
+    "import torch\n",
+    "from torch import Tensor\n",
+    "import torch_geometric as pyg\n",
+    "from torch_geometric.data import Data, Batch\n",
+    "from torch_geometric.nn.conv import MessagePassing\n",
+    "from torch_scatter import scatter\n",
+    "from torch_geometric.typing import (\n",
+    "    Adj,\n",
+    "    OptPairTensor,\n",
+    "    OptTensor,\n",
+    "    Size,\n",
+    "    SparseTensor,\n",
+    ")\n",
+    "\n",
+    "\n",
+    "\n",
+    "from copy import deepcopy\n",
+    "from typing import Callable, Union, Optional, List\n",
+    "from goli.nn.base_graph_layer import BaseGraphStructure\n",
+    "from goli.nn.base_layers import FCLayer\n",
+    "from goli.utils.decorators import classproperty\n",
+    "\n",
+    "\n",
+    "_ = torch.manual_seed(42)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Define synthetic test graph\n",
+    "\n",
+    "We define below a small batched graph on which we can test the created layers. You can also use synthetic graphs to define unit tests "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "DataBatch(edge_index=[2, 7], feat=[7, 5], edge_feat=[7, 13], batch=[7], ptr=[3])\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",
+    "\n",
+    "# Let's create 2 simple pyg graphs. \n",
+    "# start by specifying the edges with edge index\n",
+    "edge_idx1 = torch.stack([torch.tensor([0, 1, 2]), torch.tensor([1, 2, 3])])\n",
+    "edge_idx2 = torch.stack([torch.tensor([0, 0, 0, 1]), torch.tensor([0, 1, 2, 0])])\n",
+    "\n",
+    "# specify the node features, convention with variable x\n",
+    "x1 = torch.randn(edge_idx1.max() + 1, in_dim, dtype=torch.float32)\n",
+    "x2 = torch.randn(edge_idx2.max() + 1, in_dim, dtype=torch.float32)\n",
+    "\n",
+    "# specify the edge features in e\n",
+    "e1 = torch.randn(edge_idx1.shape[-1], in_dim_edges, dtype=torch.float32)\n",
+    "e2 = torch.randn(edge_idx2.shape[-1], in_dim_edges, dtype=torch.float32)\n",
+    "\n",
+    "\n",
+    "# make the pyg graph objects with our constructed features\n",
+    "g1 = Data(feat=x1, edge_index=edge_idx1, edge_feat=e1)\n",
+    "g2 = Data(feat=x2, edge_index=edge_idx2, edge_feat=e2)\n",
+    "\n",
+    "# put the two graphs into a Batch graph\n",
+    "bg = Batch.from_data_list([g1, g2])\n",
+    "\n",
+    "# The batched graph will show as a single graph with 7 nodes\n",
+    "print(bg)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Creating a simple layer\n",
+    "\n",
+    "Here, we will show how to create a GNN layer that simples does a mean aggregation on the neighbouring features.\n",
+    "\n",
+    "First, for the layer to be fully compatible with the flexible architecture provided by [`FeedForwardPyg`](https://valence-discovery.github.io/goli/api/goli.nn/architectures.html#goli.nn.architectures.pyg_architectures.FeedForwardPyg), it needs to inherit from the class [`BaseGraphStructure`](https://valence-discovery.github.io/goli/api/goli.nn/goli.nn.html#goli.nn.base_graph_layer.BaseGraphStructure). 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`](https://valence-discovery.github.io/goli/api/goli.nn/goli.nn.html#goli.nn.base_graph_layer.BaseGraphStructure): We want to return `False` since our layer doesn't support edges\n",
+    "- [`layer_inputs_edges`](https://valence-discovery.github.io/goli/api/goli.nn/goli.nn.html#goli.nn.base_graph_layer.BaseGraphStructure.layer_inputs_edges): We want to return `False` since our layer doesn't input edges\n",
+    "- [`layer_outputs_edges`](https://valence-discovery.github.io/goli/api/goli.nn/goli.nn.html#goli.nn.base_graph_layer.BaseGraphStructure.layer_outputs_edges): We want to return `False` since our layer doesn't output edges\n",
+    "- [`out_dim_factor`](https://valence-discovery.github.io/goli/api/goli.nn/goli.nn.html#goli.nn.base_graph_layer.BaseGraphStructure.out_dim_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": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#inherent from message passing class from pyg \n",
+    "#inherent from BaseGraphStructure\n",
+    "# this example is also based of : https://pytorch-geometric.readthedocs.io/en/latest/_modules/torch_geometric/nn/conv/simple_conv.html#SimpleConv.forward\n",
+    "\n",
+    "class SimpleMeanLayer(MessagePassing, BaseGraphStructure):\n",
+    "    def __init__(self, \n",
+    "                 in_dim: int, \n",
+    "                 out_dim: int, \n",
+    "                 activation: Union[Callable, str] = \"relu\", \n",
+    "                 dropout: float = 0.0, \n",
+    "                 normalization: Union[str, Callable] = \"none\",\n",
+    "                 aggr: str = \"sum\",\n",
+    "                ):\n",
+    "        r\"\"\"\n",
+    "        add documentation in the format shown here for this layer to be automatically added to the goli api reference\n",
+    "        the type information will also be automatically shown in the doc\n",
+    "        \n",
+    "        Parameters:\n",
+    "\n",
+    "            in_dim:\n",
+    "                Input feature dimensions of the layer\n",
+    "\n",
+    "            out_dim:\n",
+    "                Output feature dimensions of the layer\n",
+    "\n",
+    "            activation:\n",
+    "                activation function to use in the layer\n",
+    "\n",
+    "            dropout:\n",
+    "                The ratio of units to dropout. Must be between 0 and 1\n",
+    "\n",
+    "            normalization:\n",
+    "                Normalization to use. Choices:\n",
+    "\n",
+    "                - \"none\" or `None`: No normalization\n",
+    "                - \"batch_norm\": Batch normalization\n",
+    "                - \"layer_norm\": Layer normalization\n",
+    "                - `Callable`: Any callable function\n",
+    "            aggr:\n",
+    "                what aggregation to use\n",
+    "        \"\"\"\n",
+    "        # Initialize the parent class\n",
+    "        MessagePassing.__init__(self, node_dim=0)\n",
+    "        BaseGraphStructure.__init__(self,\n",
+    "                                    in_dim=in_dim, \n",
+    "                                    out_dim=out_dim, \n",
+    "                                    activation=activation,\n",
+    "                                    dropout=dropout, \n",
+    "                                    normalization=normalization)\n",
+    "        \n",
+    "        self.aggr = aggr\n",
+    "        self._initialize_activation_dropout_norm()\n",
+    "        # Create the mlp layer \n",
+    "        # https://valence-discovery.github.io/goli/api/goli.nn/goli.nn.html#goli.nn.base_layers.MLP\n",
+    "        self.transform = FCLayer(in_dim=in_dim, out_dim=out_dim)\n",
+    "            \n",
+    "    # define the forward function \n",
+    "    def forward(self, \n",
+    "                batch: Union[Data, Batch],\n",
+    "               ) -> Union[Data, Batch]:\n",
+    "        r\"\"\"\n",
+    "        similarly add documentation to the functions in the class\n",
+    "        \n",
+    "        Parameters:\n",
+    "            batch: pyg Batch graphs to pass through the layer\n",
+    "        Returns:\n",
+    "            batch: pyg Batch graphs\n",
+    "        \"\"\"\n",
+    "        # We first apply the mean aggregation\n",
+    "        \n",
+    "        x = batch.feat\n",
+    "        edge_index = batch.edge_index\n",
+    "        \n",
+    "        if isinstance(x, Tensor):\n",
+    "            x: OptPairTensor = (x, x)\n",
+    "\n",
+    "        # propagate_type: (x: OptPairTensor, edge_weight: OptTensor)\n",
+    "        h = self.propagate(edge_index, x=x)\n",
+    "        out = self.transform(h)\n",
+    "        out = self.apply_norm_activation_dropout(out,batch_idx=batch.batch)\n",
+    "        batch.feat = out\n",
+    "        return batch\n",
+    "    \n",
+    "    def message_and_aggregate(self, \n",
+    "                              adj_t: SparseTensor,\n",
+    "                              x: OptPairTensor) -> Tensor:\n",
+    "        return spmm(adj_t, x[0], reduce=self.aggr)\n",
+    "\n",
+    "    # Finally, we define all the virtual properties according to how the class works with proper documentation of course\n",
+    "    @classproperty\n",
+    "    def layer_supports_edges(cls) -> bool:\n",
+    "        r\"\"\"\n",
+    "        Return a boolean specifying if the layer type supports edges or not.\n",
+    "\n",
+    "        Returns:\n",
+    "\n",
+    "            bool:\n",
+    "                False for the current class\n",
+    "        \"\"\"\n",
+    "        return False\n",
+    "\n",
+    "    @property\n",
+    "    def layer_inputs_edges(self) -> bool:\n",
+    "        r\"\"\"\n",
+    "        Returns:\n",
+    "\n",
+    "            bool:\n",
+    "                Returns False\n",
+    "        \"\"\"\n",
+    "        return False\n",
+    "            \n",
+    "    @property\n",
+    "    def layer_outputs_edges(self) -> bool:\n",
+    "        r\"\"\"\n",
+    "        Returns:\n",
+    "\n",
+    "            bool:\n",
+    "                Always ``False`` for the current class\n",
+    "        \"\"\"\n",
+    "        return False\n",
+    "\n",
+    "    @property\n",
+    "    def out_dim_factor(self) -> int:\n",
+    "        r\"\"\"\n",
+    "        Get the factor by which the output dimension is multiplied for\n",
+    "        the next layer.\n",
+    "\n",
+    "        For standard layers, this will return ``1``.\n",
+    "        Returns:\n",
+    "\n",
+    "            int:\n",
+    "                Always ``1`` for the current class\n",
+    "        \"\"\"\n",
+    "        return 1"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Test our simple layer\n",
+    "\n",
+    "Now, we are ready to test the `SimpleMeanLayer` on our constructed graphs. Note that in this example, we **ignore** the edge features since they are not supported."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([7, 5])\n",
+      "torch.Size([7, 11])\n"
+     ]
+    }
+   ],
+   "source": [
+    "graph = deepcopy(bg)\n",
+    "print(graph.feat.shape)\n",
+    "\n",
+    "layer = SimpleMeanLayer(\n",
+    "            in_dim=in_dim, \n",
+    "            out_dim=out_dim, \n",
+    "            activation=\"relu\", \n",
+    "            dropout=.3, \n",
+    "            normalization=\"batch_norm\",\n",
+    "            aggr=\"sum\")\n",
+    "\n",
+    "graph = layer(graph)\n",
+    "print(graph.feat.shape)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "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`](https://valence-discovery.github.io/goli/api/goli.nn/goli.nn.html#goli.nn.base_graph_layer.BaseGraphStructure): We want to return `True` since our layer does support edges\n",
+    "- [`layer_inputs_edges`](https://valence-discovery.github.io/goli/api/goli.nn/goli.nn.html#goli.nn.base_graph_layer.BaseGraphStructure.layer_inputs_edges): We want to return `True` since our layer does input edges\n",
+    "- [`layer_outputs_edges`](https://valence-discovery.github.io/goli/api/goli.nn/goli.nn.html#goli.nn.base_graph_layer.BaseGraphStructure.layer_outputs_edges): We want to return `False` since our layer will not output new edges\n",
+    "- [`out_dim_factor`](https://valence-discovery.github.io/goli/api/goli.nn/goli.nn.html#goli.nn.base_graph_layer.BaseGraphStructure.out_dim_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": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#inherent from message passing class from pyg \n",
+    "#inherent from BaseGraphStructure\n",
+    "# adapting example from goli/nn/pyg_layers/png_pyg.py\n",
+    "\n",
+    "class ComplexMeanLayer(MessagePassing, BaseGraphStructure):\n",
+    "    def __init__(self, \n",
+    "                 in_dim: int, \n",
+    "                 out_dim: int, \n",
+    "                 in_dim_edges: int, \n",
+    "                 activation: Union[Callable, str] = \"relu\", \n",
+    "                 dropout: float = 0.0, \n",
+    "                 normalization: Union[str, Callable] = \"none\",\n",
+    "                 aggr: str = \"sum\",\n",
+    "                ):\n",
+    "        r\"\"\"\n",
+    "        add documentation in the format shown here for this layer to be automatically added to the goli api reference\n",
+    "        the type information will also be automatically shown in the doc\n",
+    "        \n",
+    "        Parameters:\n",
+    "\n",
+    "            in_dim:\n",
+    "                Input feature dimensions of the layer\n",
+    "\n",
+    "            out_dim:\n",
+    "                Output feature dimensions of the layer\n",
+    "            \n",
+    "            in_dim_edges:\n",
+    "                Input edge feature dimension of the layer\n",
+    "\n",
+    "            activation:\n",
+    "                activation function to use in the layer\n",
+    "\n",
+    "            dropout:\n",
+    "                The ratio of units to dropout. Must be between 0 and 1\n",
+    "\n",
+    "            normalization:\n",
+    "                Normalization to use. Choices:\n",
+    "\n",
+    "                - \"none\" or `None`: No normalization\n",
+    "                - \"batch_norm\": Batch normalization\n",
+    "                - \"layer_norm\": Layer normalization\n",
+    "                - `Callable`: Any callable function\n",
+    "            aggr:\n",
+    "                what aggregation to use\n",
+    "        \"\"\"\n",
+    "        # Initialize the parent class\n",
+    "        MessagePassing.__init__(self, node_dim=0)\n",
+    "        BaseGraphStructure.__init__(self,\n",
+    "                                    in_dim=in_dim, \n",
+    "                                    out_dim=out_dim, \n",
+    "                                    activation=activation,\n",
+    "                                    dropout=dropout, \n",
+    "                                    normalization=normalization)\n",
+    "        \n",
+    "        self.aggr = aggr\n",
+    "        self._initialize_activation_dropout_norm()\n",
+    "        # Create the mlp layer \n",
+    "        # https://valence-discovery.github.io/goli/api/goli.nn/goli.nn.html#goli.nn.base_layers.MLP\n",
+    "        self.transform = FCLayer(in_dim=(in_dim * 2 +in_dim_edges), out_dim=out_dim)\n",
+    "            \n",
+    "    # define the forward function \n",
+    "    def forward(self, \n",
+    "                batch: Union[Data, Batch],\n",
+    "               ) -> Union[Data, Batch]:\n",
+    "        r\"\"\"\n",
+    "        similarly add documentation to the functions in the class\n",
+    "        \n",
+    "        Parameters:\n",
+    "            batch: pyg Batch graphs to pass through the layer\n",
+    "        Returns:\n",
+    "            batch: pyg Batch graphs\n",
+    "        \"\"\"        \n",
+    "        x = batch.feat\n",
+    "        edge_index = batch.edge_index\n",
+    "        edge_feat = batch.edge_feat\n",
+    "        \n",
+    "        if isinstance(x, Tensor):\n",
+    "            x: OptPairTensor = (x, x)\n",
+    "\n",
+    "        # propagate_type: (x: OptPairTensor, edge_weight: OptTensor)\n",
+    "        h = self.propagate(edge_index, x=x, edge_feat=edge_feat, size=None)\n",
+    "        out = self.transform(h)\n",
+    "        out = self.apply_norm_activation_dropout(out,batch_idx=batch.batch)\n",
+    "        batch.feat = out\n",
+    "        return batch\n",
+    "    \n",
+    "    def message(self, x_i: Tensor, x_j: Tensor, edge_feat: OptTensor) -> Tensor:\n",
+    "        r\"\"\"\n",
+    "        message function\n",
+    "\n",
+    "        Parameters:\n",
+    "            x_i: node features\n",
+    "            x_j: neighbour node features\n",
+    "            edge_feat: edge features\n",
+    "        Returns:\n",
+    "            feat: the message\n",
+    "        \"\"\"\n",
+    "        feat = torch.cat([x_i, x_j, edge_feat], dim=-1)\n",
+    "        return feat\n",
+    "    \n",
+    "    def aggregate(\n",
+    "        self,\n",
+    "        inputs: Tensor,\n",
+    "        index: Tensor,\n",
+    "        edge_index: Tensor,\n",
+    "        dim_size: Optional[int] = None,\n",
+    "    ) -> Tensor:\n",
+    "        r\"\"\"\n",
+    "        aggregate function\n",
+    "        Parameters:\n",
+    "            inputs: input features\n",
+    "            index: index of the nodes\n",
+    "            edge_index: edge index\n",
+    "            dim_size: dimension size\n",
+    "        Returns:\n",
+    "            out: aggregated features\n",
+    "        \"\"\"\n",
+    "        out = scatter(inputs, index, 0, None, dim_size, reduce=self.aggr)\n",
+    "        return out\n",
+    "    \n",
+    "\n",
+    "    # Finally, we define all the virtual properties according to how the class works with proper documentation of course\n",
+    "    @classproperty\n",
+    "    def layer_supports_edges(cls) -> bool:\n",
+    "        r\"\"\"\n",
+    "        Return a boolean specifying if the layer type supports edges or not.\n",
+    "\n",
+    "        Returns:\n",
+    "\n",
+    "            bool:\n",
+    "                False for the current class\n",
+    "        \"\"\"\n",
+    "        return True\n",
+    "\n",
+    "    @property\n",
+    "    def layer_inputs_edges(self) -> bool:\n",
+    "        r\"\"\"\n",
+    "        Returns:\n",
+    "\n",
+    "            bool:\n",
+    "                Returns True\n",
+    "        \"\"\"\n",
+    "        return True\n",
+    "            \n",
+    "    @property\n",
+    "    def layer_outputs_edges(self) -> bool:\n",
+    "        r\"\"\"\n",
+    "        Returns:\n",
+    "\n",
+    "            bool:\n",
+    "                Always ``True`` for the current class\n",
+    "        \"\"\"\n",
+    "        return True\n",
+    "\n",
+    "    @property\n",
+    "    def out_dim_factor(self) -> int:\n",
+    "        r\"\"\"\n",
+    "        Get the factor by which the output dimension is multiplied for\n",
+    "        the next layer.\n",
+    "\n",
+    "        For standard layers, this will return ``1``.\n",
+    "        Returns:\n",
+    "\n",
+    "            int:\n",
+    "                Always ``1`` for the current class\n",
+    "        \"\"\"\n",
+    "        return 1"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Test our complex layer \n",
+    "\n",
+    "Now, we are ready to test the `ComplexMeanLayer` on our constructed graphs. Note that in this example, we utilize the edge features and node features"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([7, 5])\n",
+      "torch.Size([7, 11])\n"
+     ]
+    }
+   ],
+   "source": [
+    "graph = deepcopy(bg)\n",
+    "print(graph.feat.shape)\n",
+    "\n",
+    "layer = ComplexMeanLayer(\n",
+    "            in_dim=in_dim, \n",
+    "            out_dim=out_dim, \n",
+    "            in_dim_edges=in_dim_edges,\n",
+    "            activation=\"relu\", \n",
+    "            dropout=.3, \n",
+    "            normalization=\"batch_norm\",\n",
+    "            aggr=\"sum\")\n",
+    "\n",
+    "graph = layer(graph)\n",
+    "print(graph.feat.shape)"
+   ]
+  }
+ ],
+ "metadata": {
+  "interpreter": {
+   "hash": "f4a99d018a205fcbcc0480c84566beaebcb91b08d0414b39a842df533e2a1d25"
+  },
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "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.8.10"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/docs/tutorials/basics/making_gnn_networks.ipynb b/docs/tutorials/gnn/making_gnn_networks.ipynb
similarity index 100%
rename from docs/tutorials/basics/making_gnn_networks.ipynb
rename to docs/tutorials/gnn/making_gnn_networks.ipynb
diff --git a/docs/tutorials/basics/using_gnn_layers.ipynb b/docs/tutorials/gnn/using_gnn_layers.ipynb
similarity index 100%
rename from docs/tutorials/basics/using_gnn_layers.ipynb
rename to docs/tutorials/gnn/using_gnn_layers.ipynb
diff --git a/docs/tutorials/model_training/ipu_training.ipynb b/docs/tutorials/model_training/ipu_training.ipynb
index cc32c7c97..85b4d9097 100644
--- a/docs/tutorials/model_training/ipu_training.ipynb
+++ b/docs/tutorials/model_training/ipu_training.ipynb
@@ -1435,7 +1435,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "3.1.0+1181_poptorch",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
    "name": "python3"
   },
@@ -1449,7 +1449,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.10 (default, Jun 22 2022, 20:18:18) \n[GCC 9.4.0]"
+   "version": "3.8.10"
   },
   "vscode": {
    "interpreter": {
diff --git a/env.yml b/env.yml
index 89984c363..df1bbae7a 100644
--- a/env.yml
+++ b/env.yml
@@ -15,6 +15,7 @@ dependencies:
   - scipy >=1.4
   - pandas >=1.0
   - scikit-learn
+  - fastparquet
 
   # viz
   - matplotlib >=3.0.1
diff --git a/expts/configs/config_gpu_qm9.yaml b/expts/configs/config_gpu_qm9.yaml
deleted file mode 100644
index 00a46da2d..000000000
--- a/expts/configs/config_gpu_qm9.yaml
+++ /dev/null
@@ -1,324 +0,0 @@
-# Testing the multitask pipeline with the QM9 dataset on IPU, by splitting it up into three tasks: homo, alpha and cv.
-constants:
-  name: &name qm9_micro
-  seed: &seed 42
-  raise_train_error: true   # Whether the code should raise an error if it crashes during training
-  accelerator:
-<<<<<<< HEAD:expts/configs/config_gpu_qm9.yaml
-    type: gpu  #cpu or ipu or gpu
-=======
-    type: cpu  #cpu or ipu or gpu
->>>>>>> minor update, all tested fine:expts/configs/config_qm9_gpspp.yaml
-
-
-datamodule:
-  module_type: "MultitaskFromSmilesDataModule"
-  args: # Matches that in the test_multitask_datamodule.py case.
-    task_specific_args:   # To be replaced by a new class "DatasetParams"
-      homo:
-        df: null
-        # df_path: &df_path goli/data/QM9/norm_micro_qm9.csv #goli/data/QM9/qm9.csv #goli/data/QM9/norm_micro_qm9.csv #https://storage.googleapis.com/goli-public/datasets/QM9/norm_mini_qm9.csv #goli/data/QM9/micro_qm9.csv #goli/data/QM9/norm_micro_qm9.csv   #https://storage.googleapis.com/goli-public/datasets/QM9/norm_mini_qm9.csv
-        df_path: &df_path https://storage.googleapis.com/goli-public/datasets/QM9/norm_qm9.csv
-        smiles_col: "smiles"
-        label_cols: ["homo", "lumo"]
-        split_val: 0.2
-        split_test: 0.2
-        split_seed: *seed
-        splits_path: null                 # This may not always be provided
-        sample_size: null                 # This may not always be provided
-        idx_col: null                     # This may not always be provided
-        weights_col: null                 # This may not always be provided
-        weights_type: null                # This may not always be provided
-      alpha:
-        df: null
-        df_path: *df_path
-        smiles_col: "smiles"
-        label_cols: ["alpha"]
-        split_val: 0.2
-        split_test: 0.2
-        split_seed: *seed
-        splits_path: null                 # This may not always be provided
-        sample_size: null                 # This may not always be provided
-        idx_col: null                     # This may not always be provided
-        weights_col: null                 # This may not always be provided
-        weights_type: null                # This may not always be provided
-      cv:
-        df: null
-        df_path: *df_path
-        smiles_col: "smiles"
-        label_cols: ["cv"]
-        split_val: 0.2
-        split_test: 0.2
-        split_seed: *seed
-        splits_path: null                 # This may not always be provided
-        sample_size: null                 # This may not always be provided
-        idx_col: null                     # This may not always be provided
-        weights_col: null                 # This may not always be provided
-        weights_type: null                # This may not always be provided
-
-    # Featurization
-    prepare_dict_or_graph: pyg:graph
-    featurization_n_jobs: 32
-    featurization_progress: True
-    featurization:
-      atom_property_list_onehot: [atomic-number, valence]
-      atom_property_list_float: [mass, electronegativity, in-ring]
-      edge_property_list: [bond-type-onehot, stereo, in-ring]
-      add_self_loop: False
-      explicit_H: False
-      use_bonds_weights: False
-      pos_encoding_as_features:
-        pos_types:
-          la_pos: &pos_enc  #use same name as pe_encoder
-            pos_type: laplacian_eigvec_eigval #laplacian_eigvec
-            num_pos: 3
-            normalization: "none"
-            disconnected_comp: True
-          rw_pos: #use same name as pe_encoder
-            pos_type: rwse
-            ksteps: 16
-      # pos_encoding_as_directions: *pos_enc # Only for DGN or directional pooling
-
-    # Data handling-related
-    batch_size_training: 6000
-    batch_size_inference: 6000
-    # cache_data_path: null
-
-    num_workers: -1
-
-    ipu_dataloader_training_opts:
-      max_num_nodes_per_graph: 12
-      max_num_edges_per_graph: 24
-
-    ipu_dataloader_inference_opts:
-      max_num_nodes_per_graph: 12
-      max_num_edges_per_graph: 24
-
-architecture:
-  model_type: FullGraphMultiTaskNetwork
-  mup_base_path: null
-  pre_nn:   # Set as null to avoid a pre-nn network
-    out_dim: 32
-    hidden_dims: 32
-    depth: 1
-    activation: relu
-    last_activation: none
-    dropout: &dropout 0.1
-    normalization: &normalization batch_norm
-    last_normalization: *normalization
-    residual_type: none
-
-  pre_nn_edges:   # Set as null to avoid a pre-nn network
-    out_dim: 16
-    hidden_dims: 16
-    depth: 1
-    activation: relu
-    last_activation: none
-    dropout: *dropout
-    normalization: *normalization
-    last_normalization: *normalization
-    residual_type: none
-
-  pe_encoders:
-    out_dim: 64
-    pool: "sum" #"mean" "max"
-    last_norm: None #"batch_norm", "layer_norm"
-    encoders: #la_pos |  rw_pos
-      la_pos:  # Set as null to avoid a pre-nn network
-        encoder_type: "laplacian_pe"
-        input_keys: ["eigvecs", "eigvals"]
-        hidden_dim: 64
-        model_type: 'DeepSet' #'Transformer' or 'DeepSet'
-        num_layers: 1
-        num_layers_post: 0 # Num. layers to apply after pooling
-        dropout: 0.1
-        first_normalization: "none" #"batch_norm" or "layer_norm"
-      rw_pos:
-        encoder_type: "mlp"
-        input_keys: ["rwse"]
-        output_keys: "node"
-        hidden_dim: 64
-        num_layers: 1
-        dropout: 0.1
-        normalization: "none" #"batch_norm" or "layer_norm"
-        first_normalization: "batch_norm" #"batch_norm" or "layer_norm"
-
-
-  gnn:  # Set as null to avoid a post-nn network
-    out_dim: 500
-    hidden_dims: 500
-    depth: 16
-    activation: relu
-    last_activation: none
-    dropout: *dropout
-    normalization: *normalization
-    last_normalization: *normalization
-    residual_type: simple
-    pooling: [sum, mean, max]
-    virtual_node: 'none'
-    layer_type: 'pyg:gps' #pyg:gine #'pyg:gps' # pyg:gated-gcn, pyg:gine,pyg:gps
-    layer_kwargs:  # Parameters for the model itself. You could define dropout_attn: 0.1
-      mpnn_type: 'pyg:gine'
-      mpnn_kwargs: null
-        #out_dim_edges: 10
-      attn_type: "full-attention" # "full-attention", "none"
-      attn_kwargs: null
-
-
-      # out_dim_edges: 16
-      # aggregators: [mean, max]
-      # scalers: [identity, amplification, attenuation]
-      # num_heads: 3
-
-  post_nn:
-    out_dim: 32
-    hidden_dims: 400
-    depth: 2
-    activation: relu
-    last_activation: none
-    dropout: *dropout
-    normalization: *normalization
-    last_normalization: "none"
-    residual_type: none
-
-  task_heads:
-    homo:
-      out_dim: 2
-      hidden_dims: 32
-      depth: 2                          # Not needed if we have hidden_dims
-      activation: relu
-      last_activation: none
-      dropout: *dropout
-      normalization: *normalization
-      last_normalization: "none"
-      residual_type: none
-    alpha:
-      out_dim: 1
-      hidden_dims: 32
-      depth: 2                          # Not needed if we have hidden_dims
-      activation: relu
-      last_activation: none
-      dropout: *dropout
-      normalization: *normalization
-      last_normalization: "none"
-      residual_type: none
-    cv:
-      out_dim: 1
-      hidden_dims: 32
-      depth: 2                          # Not needed if we have hidden_dims
-      activation: relu
-      last_activation: none
-      dropout: *dropout
-      normalization: *normalization
-      last_normalization: "none"
-      residual_type: none
-
-#Task-specific
-predictor:
-  metrics_on_progress_bar:
-    homo: ["mae", "pearsonr", "spearmanr", "f1 > 0.5", "precision > 0.5"]
-    alpha: ["mae"]
-    cv: ["mae"]
-  loss_fun:
-    homo: mse_ipu
-    alpha: mse_ipu
-    cv: mse_ipu
-  random_seed: *seed
-  optim_kwargs:
-    lr: 1.e-4 #1.e-3
-    # weight_decay: 1.e-7
-  torch_scheduler_kwargs:
-    #module_type: ReduceLROnPlateau
-    #factor: 0.5
-    #patience: 7
-  scheduler_kwargs: null
-  #  monitor: &monitor loss/val
-  #  mode: min
-  #  frequency: 1
-  target_nan_mask: null # null: no mask, 0: 0 mask, ignore: ignore nan values from loss
-  flag_kwargs:
-    n_steps: 0 #1
-    alpha: 0.0 #0.01
-
-# Task-specific
-metrics:
-  homo:
-    - name: mae
-      metric: mae_ipu
-      target_nan_mask: null
-      multitask_handling: flatten
-      threshold_kwargs: null
-    - name: pearsonr
-      metric: pearsonr_ipu
-      threshold_kwargs: null
-      target_nan_mask: null
-      multitask_handling: mean-per-label
-    - name: spearmanr
-      metric: spearmanr_ipu
-      threshold_kwargs: null
-      target_nan_mask: null
-      multitask_handling: mean-per-label
-    - name: f1 > 0.5
-      metric: f1_ipu
-      num_classes: 2
-      average: micro
-      target_nan_mask: null
-      multitask_handling: mean-per-label
-      threshold_kwargs: &threshold_05
-        operator: greater
-        threshold: 0.5
-        th_on_preds: True
-        th_on_target: True
-        target_to_int: False
-    - name: precision > 0.5
-      metric: precision_ipu
-      average: micro
-      target_nan_mask: null
-      multitask_handling: mean-per-label
-      threshold_kwargs: *threshold_05
-
-  alpha:
-    - name: mae
-      metric: mae_ipu
-      target_nan_mask: null
-      multitask_handling: flatten
-      threshold_kwargs: null
-    - name: pearsonr
-      metric: pearsonr_ipu
-      target_nan_mask: null
-      multitask_handling: mean-per-label
-      threshold_kwargs: null
-
-  cv:
-    - name: mae
-      metric: mae_ipu
-      target_nan_mask: null
-      multitask_handling: flatten
-      threshold_kwargs: null
-    - name: pearsonr
-      metric: pearsonr_ipu
-      target_nan_mask: null
-      multitask_handling: mean-per-label
-      threshold_kwargs: null
-
-trainer:
-  logger:
-    save_dir: logs/QM9
-    name: *name
-  #early_stopping:
-  #  monitor: *monitor
-  #  min_delta: 0
-  #  patience: 10
-  #  mode: &mode min
-  model_checkpoint:
-    dirpath: models_checkpoints/QM9/
-    filename: *name
-    #monitor: *monitor
-    #mode: *mode
-    save_top_k: 1
-    every_n_epochs: 1
-  trainer:
-    precision: 32
-    max_epochs: 50 #50
-    min_epochs: 1 #5
-    accumulate_grad_batches: 1
diff --git a/expts/configs/config_ipu_qm9.yaml b/expts/configs/config_ipu_qm9.yaml
deleted file mode 100644
index 940e3d1b2..000000000
--- a/expts/configs/config_ipu_qm9.yaml
+++ /dev/null
@@ -1,322 +0,0 @@
-# Testing the multitask pipeline with the QM9 dataset on IPU, by splitting it up into three tasks: homo, alpha and cv.
-constants:
-  name: &name qm9_full_ipu
-  seed: &seed 42
-  raise_train_error: true   # Whether the code should raise an error if it crashes during training
-  accelerator:
-    type: ipu  #cpu or ipu or gpu
-
-
-datamodule:
-  module_type: "MultitaskFromSmilesDataModule"
-  args: # Matches that in the test_multitask_datamodule.py case.
-    task_specific_args:   # To be replaced by a new class "DatasetParams"
-      homo:
-        df: null
-        # df_path: &df_path goli/data/QM9/norm_micro_qm9.csv #goli/data/QM9/qm9.csv #goli/data/QM9/norm_micro_qm9.csv #https://storage.googleapis.com/goli-public/datasets/QM9/norm_mini_qm9.csv #goli/data/QM9/micro_qm9.csv #goli/data/QM9/norm_micro_qm9.csv   #https://storage.googleapis.com/goli-public/datasets/QM9/norm_mini_qm9.csv
-        df_path: &df_path https://storage.googleapis.com/goli-public/datasets/QM9/norm_micro_qm9.csv
-        smiles_col: "smiles"
-        label_cols: ["homo", "lumo"]
-        split_val: 0.2
-        split_test: 0.2
-        split_seed: *seed
-        splits_path: null                 # This may not always be provided
-        sample_size: null                 # This may not always be provided
-        idx_col: null                     # This may not always be provided
-        weights_col: null                 # This may not always be provided
-        weights_type: null                # This may not always be provided
-      alpha:
-        df: null
-        df_path: *df_path
-        smiles_col: "smiles"
-        label_cols: ["alpha"]
-        split_val: 0.2
-        split_test: 0.2
-        split_seed: *seed
-        splits_path: null                 # This may not always be provided
-        sample_size: null                 # This may not always be provided
-        idx_col: null                     # This may not always be provided
-        weights_col: null                 # This may not always be provided
-        weights_type: null                # This may not always be provided
-      cv:
-        df: null
-        df_path: *df_path
-        smiles_col: "smiles"
-        label_cols: ["cv"]
-        split_val: 0.2
-        split_test: 0.2
-        split_seed: *seed
-        splits_path: null                 # This may not always be provided
-        sample_size: null                 # This may not always be provided
-        idx_col: null                     # This may not always be provided
-        weights_col: null                 # This may not always be provided
-        weights_type: null                # This may not always be provided
-
-    # Featurization
-    prepare_dict_or_graph: pyg:graph
-    featurization_n_jobs: 32
-    featurization_progress: True
-    featurization:
-      atom_property_list_onehot: [atomic-number, valence]
-      atom_property_list_float: [mass, electronegativity, in-ring]
-      edge_property_list: [bond-type-onehot, stereo, in-ring]
-      add_self_loop: False
-      explicit_H: False
-      use_bonds_weights: False
-      pos_encoding_as_features:
-        pos_types:
-          la_pos: &pos_enc  #use same name as pe_encoder
-            pos_type: laplacian_eigvec_eigval #laplacian_eigvec
-            num_pos: 3
-            normalization: "none"
-            disconnected_comp: True
-          rw_pos: #use same name as pe_encoder
-            pos_type: rwse
-            ksteps: 16
-      # pos_encoding_as_directions: *pos_enc # Only for DGN or directional pooling
-
-    # Data handling-related
-    batch_size_training: 9
-    batch_size_inference: 9
-    batch_size_per_pack: 3
-    cache_data_path: "datacache"
-
-    num_workers: -1
-
-    ipu_dataloader_training_opts:
-      max_num_nodes_per_graph: 9
-      max_num_edges_per_graph: 24
-
-    ipu_dataloader_inference_opts:
-      max_num_nodes_per_graph: 9
-      max_num_edges_per_graph: 24
-
-architecture:
-  model_type: FullGraphMultiTaskNetwork
-  mup_base_path: null
-  pre_nn:   # Set as null to avoid a pre-nn network
-    out_dim: 32
-    hidden_dims: 32
-    depth: 1
-    activation: relu
-    last_activation: none
-    dropout: &dropout 0.1
-    normalization: &normalization batch_norm
-    last_normalization: *normalization
-    residual_type: none
-
-  pre_nn_edges:   # Set as null to avoid a pre-nn network
-    out_dim: 16
-    hidden_dims: 16
-    depth: 1
-    activation: relu
-    last_activation: none
-    dropout: *dropout
-    normalization: *normalization
-    last_normalization: *normalization
-    residual_type: none
-
-  pe_encoders:
-    out_dim: 64
-    pool: "sum" #"mean" "max"
-    last_norm: None #"batch_norm", "layer_norm"
-    encoders: #la_pos |  rw_pos
-      la_pos:  # Set as null to avoid a pre-nn network
-        encoder_type: "laplacian_pe"
-        input_keys: ["eigvecs", "eigvals"]
-        output_keys: ["feat"]
-        hidden_dim: 64
-        model_type: 'DeepSet' #'Transformer' or 'DeepSet'
-        num_layers: 1
-        num_layers_post: 0 # Num. layers to apply after pooling
-        dropout: 0.1
-        first_normalization: "none" #"batch_norm" or "layer_norm"
-      rw_pos:
-        encoder_type: "mlp"
-        input_keys: ["rwse"]
-        output_keys: ["feat"]
-        hidden_dim: 64
-        num_layers: 1
-        dropout: 0.1
-        normalization: "none" #"batch_norm" or "layer_norm"
-        first_normalization: "batch_norm" #"batch_norm" or "layer_norm"
-
-
-  gnn:  # Set as null to avoid a post-nn network
-    out_dim: 400
-    hidden_dims: 400
-    depth: 16
-    activation: relu
-    last_activation: none
-    dropout: *dropout
-    normalization: *normalization
-    last_normalization: *normalization
-    residual_type: simple
-    pooling: [sum, mean, max]
-    virtual_node: 'none'
-    layer_type: 'pyg:gps' #pyg:gine #'pyg:gps' # pyg:gated-gcn, pyg:gine,pyg:gps
-    layer_kwargs:  # Parameters for the model itself. You could define dropout_attn: 0.1
-      mpnn_type: 'pyg:gine'
-      mpnn_kwargs: null
-        #out_dim_edges: 10
-      attn_type: "full-attention" # "full-attention", "none"
-      attn_kwargs: null
-
-
-      # out_dim_edges: 16
-      # aggregators: [mean, max]
-      # scalers: [identity, amplification, attenuation]
-      # num_heads: 3
-
-  post_nn:
-    out_dim: 32
-    hidden_dims: 400
-    depth: 2
-    activation: relu
-    last_activation: none
-    dropout: *dropout
-    normalization: *normalization
-    last_normalization: "none"
-    residual_type: none
-
-  task_heads:
-    homo:
-      out_dim: 2
-      hidden_dims: 32
-      depth: 2                          # Not needed if we have hidden_dims
-      activation: relu
-      last_activation: none
-      dropout: *dropout
-      normalization: *normalization
-      last_normalization: "none"
-      residual_type: none
-    alpha:
-      out_dim: 1
-      hidden_dims: 32
-      depth: 2                          # Not needed if we have hidden_dims
-      activation: relu
-      last_activation: none
-      dropout: *dropout
-      normalization: *normalization
-      last_normalization: "none"
-      residual_type: none
-    cv:
-      out_dim: 1
-      hidden_dims: 32
-      depth: 2                          # Not needed if we have hidden_dims
-      activation: relu
-      last_activation: none
-      dropout: *dropout
-      normalization: *normalization
-      last_normalization: "none"
-      residual_type: none
-
-#Task-specific
-predictor:
-  metrics_on_progress_bar:
-    homo: ["mae", "pearsonr", "spearmanr", "f1 > 0.5", "precision > 0.5"]
-    alpha: ["mae"]
-    cv: ["mae"]
-  loss_fun:
-    homo: mse_ipu
-    alpha: mse_ipu
-    cv: mse_ipu
-  random_seed: *seed
-  optim_kwargs:
-    lr: 1.e-3
-    # weight_decay: 1.e-7
-  torch_scheduler_kwargs:
-    #module_type: ReduceLROnPlateau
-    #factor: 0.5
-    #patience: 7
-  scheduler_kwargs: null
-  #  monitor: &monitor loss/val
-  #  mode: min
-  #  frequency: 1
-  target_nan_mask: null # null: no mask, 0: 0 mask, ignore: ignore nan values from loss
-  flag_kwargs:
-    n_steps: 0 #1
-    alpha: 0.0 #0.01
-
-# Task-specific
-metrics:
-  homo:
-    - name: mae
-      metric: mae_ipu
-      target_nan_mask: null
-      multitask_handling: flatten
-      threshold_kwargs: null
-    - name: pearsonr
-      metric: pearsonr_ipu
-      threshold_kwargs: null
-      target_nan_mask: null
-      multitask_handling: mean-per-label
-    - name: spearmanr
-      metric: spearmanr_ipu
-      threshold_kwargs: null
-      target_nan_mask: null
-      multitask_handling: mean-per-label
-    - name: f1 > 0.5
-      metric: f1_ipu
-      num_classes: 2
-      average: micro
-      target_nan_mask: null
-      multitask_handling: mean-per-label
-      threshold_kwargs: &threshold_05
-        operator: greater
-        threshold: 0.5
-        th_on_preds: True
-        th_on_target: True
-        target_to_int: False
-    - name: precision > 0.5
-      metric: precision_ipu
-      average: micro
-      target_nan_mask: null
-      multitask_handling: mean-per-label
-      threshold_kwargs: *threshold_05
-
-  alpha:
-    - name: mae
-      metric: mae_ipu
-      target_nan_mask: null
-      multitask_handling: flatten
-      threshold_kwargs: null
-    - name: pearsonr
-      metric: pearsonr_ipu
-      target_nan_mask: null
-      multitask_handling: mean-per-label
-      threshold_kwargs: null
-
-  cv:
-    - name: mae
-      metric: mae_ipu
-      target_nan_mask: null
-      multitask_handling: flatten
-      threshold_kwargs: null
-    - name: pearsonr
-      metric: pearsonr_ipu
-      target_nan_mask: null
-      multitask_handling: mean-per-label
-      threshold_kwargs: null
-
-trainer:
-  logger:
-    save_dir: logs/QM9
-    name: *name
-  #early_stopping:
-  #  monitor: *monitor
-  #  min_delta: 0
-  #  patience: 10
-  #  mode: &mode min
-  model_checkpoint:
-    dirpath: models_checkpoints/QM9/
-    filename: *name
-    #monitor: *monitor
-    #mode: *mode
-    save_top_k: 1
-    every_n_epochs: 1
-  trainer:
-    precision: 32
-    max_epochs: 50 #50
-    min_epochs: 1 #5
-    accumulate_grad_batches: 2
diff --git a/expts/configs/config_ipu_reproduce.yaml b/expts/configs/config_ipu_reproduce.yaml
deleted file mode 100644
index cc756b428..000000000
--- a/expts/configs/config_ipu_reproduce.yaml
+++ /dev/null
@@ -1,301 +0,0 @@
-constants:
-  name: &name ipu-reproduce
-  seed: &seed 42
-  raise_train_error: true   # Whether the code should raise an error if it crashes during training
-  accelerator:
-    type: ipu
-
-
-datamodule:
-  module_type: "MultitaskFromSmilesDataModule"
-  args: # Matches that in the test_multitask_datamodule.py case.
-    task_specific_args:   # To be replaced by a new class "DatasetParams"
-      pc-assays:
-        df_path:  goli/data/mega-pubchem/mini.csv.gz
-        # df_path: goli/data/mega-pubchem/mega-pubchem-L1000-VCAP.csv.gz
-        # df_path: gs://goli-private/datasets/goli-pubchem-L1000/mega-pubchem-L1000-VCAP.csv.gz
-        smiles_col: "SMILES"
-        label_cols: "assayID-*"
-        split_val: 0.08
-        split_test: 0.01
-        split_seed: *seed
-        splits_path: null                 # This may not always be provided
-        sample_size: 715233 # 715233       # This may not always be provided
-        idx_col: null                     # This may not always be provided
-        weights_col: null                 # This may not always be provided
-        weights_type: null                # This may not always be provided
-
-    # Featurization
-    prepare_dict_or_graph: pyg:graph
-    featurization_n_jobs: 32
-    featurization_progress: True
-    featurization:
-      atom_property_list_onehot: [atomic-number, valence]
-      atom_property_list_float: [mass, electronegativity, in-ring, hybridization, chirality, aromatic, degree, formal-charge, single-bond, double-bond, radical-electron, vdw-radius, covalent-radius, metal]
-      edge_property_list: [bond-type-onehot, bond-type-float, stereo, in-ring, conjugated, estimated-bond-length]
-      add_self_loop: False
-      explicit_H: False
-      use_bonds_weights: False
-      max_num_atoms: 60
-      pos_encoding_as_features:
-        pos_types:
-          la_pos: &pos_enc  #use same name as pe_encoder
-            pos_type: laplacian_eigvec_eigval #laplacian_eigvec
-            num_pos: 3
-            normalization: "none"
-            disconnected_comp: True
-          rw_pos: #use same name as pe_encoder
-            pos_type: rwse
-            ksteps: 16
-      # pos_encoding_as_directions: *pos_enc # Only for DGN or directional pooling
-
-    # Data handling-related
-    batch_size_training: 9
-    batch_size_inference: 9
-    # cache_data_path: null
-
-    num_workers: 32
-    persistent_workers: False
-
-    ipu_dataloader_training_opts:
-      mode: async
-      max_num_nodes: &max_num_nodes 260
-      max_num_edges: &max_num_edges 600
-
-    ipu_dataloader_inference_opts:
-      mode: async
-      max_num_nodes: *max_num_nodes
-      max_num_edges: *max_num_edges
-
-architecture:
-  model_type: FullGraphMultiTaskNetwork
-  pre_nn:   # Set as null to avoid a pre-nn network
-    out_dim: &hidden_dim 600
-    hidden_dims: *hidden_dim
-    depth: 3
-    activation: relu
-    last_activation: none
-    dropout: &dropout_mlp 0.1
-    last_dropout: *dropout_mlp
-    normalization: &normalization batch_norm
-    last_normalization: *normalization
-    residual_type: simple
-
-  pre_nn_edges:   # Set as null to avoid a pre-nn network
-    out_dim: 32
-    hidden_dims: 32
-    depth: 2
-    activation: relu
-    last_activation: none
-    dropout: *dropout_mlp
-    last_dropout: *dropout_mlp
-    normalization: *normalization
-    last_normalization: *normalization
-    residual_type: simple
-
-  pe_encoders:
-    out_dim: 32
-    pool: "sum" #"mean" "max"
-    last_norm: None #"batch_norm", "layer_norm"
-    encoders: #la_pos |  rw_pos
-      la_pos:  # Set as null to avoid a pre-nn network
-        encoder_type: "laplacian_pe"
-        input_keys: ["eigvecs", "eigvals"]
-        hidden_dim: 32
-        model_type: 'DeepSet' #'Transformer' or 'DeepSet'
-        num_layers: 1
-        num_layers_post: 0 # Num. layers to apply after pooling
-        dropout: 0.1
-        first_normalization: "none" #"batch_norm" or "layer_norm"
-      rw_pos:
-        encoder_type: "mlp"
-        input_keys: ["rwse"]
-        output_keys: "node"
-        hidden_dim: 32
-        num_layers: 1
-        dropout: 0.1
-        normalization: "none" #"batch_norm" or "layer_norm"
-        first_normalization: "none" #"batch_norm" or "layer_norm"
-
-  gnn:  # Set as null to avoid a post-nn network
-    out_dim: &out_mlp_dim 1150
-    hidden_dims: *hidden_dim
-    depth: 8
-    activation: relu
-    last_activation: none
-    dropout: &dropout_gnn 0.07
-    last_dropout: *dropout_gnn
-    normalization: *normalization
-    last_normalization: *normalization
-    residual_type: simple
-    pooling: ['logsum', 'max', 'mean']
-    # pooling: ['mean']
-    virtual_node: 'none'
-    layer_type: "pyg:pna-msgpass" # 'pyg:gine'
-    layer_kwargs:
-      aggregators: [mean, max, sum]
-      scalers: [identity]
-
-  post_nn:
-    out_dim: *out_mlp_dim
-    hidden_dims: *out_mlp_dim
-    depth: 2
-    activation: relu
-    last_activation: relu
-    dropout: *dropout_mlp
-    last_dropout: 0.
-    normalization: *normalization
-    last_normalization: "none"
-    residual_type: simple
-
-
-  task_heads:
-    - task_name: "pc-assays"
-      out_dim: 561 # 2306
-      hidden_dims: *out_mlp_dim
-      depth: 1                          # Not needed if we have hidden_dims
-      activation: relu
-      last_activation: sigmoid
-      dropout: *dropout_mlp
-      normalization: *normalization
-      last_normalization: "none"
-      residual_type: none
-
-#Task-specific
-predictor:
-  metrics_on_progress_bar:
-    pc-assays: [] # ["mae", "averageprecision", "auroc"]
-  metrics_on_training_set:
-    pc-assays: [] # ["mae", "averageprecision", "auroc"]
-  loss_fun:
-    pc-assays: bce_ipu
-  random_seed: *seed
-  optim_kwargs:
-    lr: 2.e-4
-    # weight_decay: 1.e-7
-  torch_scheduler_kwargs:
-    #module_type: ReduceLROnPlateau
-    #factor: 0.5
-    #patience: 7
-  scheduler_kwargs: null
-  #  monitor: &monitor loss/val
-  #  mode: min
-  #  frequency: 1
-  target_nan_mask: null # null: no mask, 0: 0 mask, ignore-flatten, ignore-mean-label
-  flag_kwargs:
-    n_steps: 0 #1
-    alpha: 0.0 #0.01
-
-
-metrics:
-  pc-assays:
-    # - name: mae
-    #   metric: mae
-    #   threshold_kwargs: null
-    #   target_nan_mask: 0
-
-    # - name: averageprecision
-    #   metric: averageprecision
-    #   target_nan_mask: 0
-    #   threshold_kwargs:
-    #     threshold: null
-    #     th_on_preds: False
-    #     th_on_target: False
-    #     target_to_int: True
-
-    # - name: auroc
-    #   metric: auroc
-    #   target_nan_mask: 0
-    #   threshold_kwargs:
-    #     threshold: null
-    #     th_on_preds: False
-    #     th_on_target: False
-    #     target_to_int: True
-
-    # - name: f1 > 0.5
-    #   metric: f1
-    #   num_classes: 2
-    #   average: macro
-    #   target_nan_mask: ignore-mean-label
-    #   threshold_kwargs: &threshold_1
-    #     operator: greater
-    #     threshold: 0.5
-    #     th_on_preds: True
-    #     th_on_target: False
-    #     target_to_int: True
-
-    # - name: f1 > 0.25
-    #   metric: f1
-    #   num_classes: 2
-    #   average: macro
-    #   target_nan_mask: ignore-mean-label
-    #   threshold_kwargs: &threshold_2
-    #     operator: greater
-    #     threshold: 0.25
-    #     th_on_preds: True
-    #     th_on_target: False
-    #     target_to_int: True
-
-    # - name: f1 > 0.75
-    #   metric: f1
-    #   num_classes: 2
-    #   average: macro
-    #   target_nan_mask: ignore-mean-label
-    #   threshold_kwargs: &threshold_3
-    #     operator: greater
-    #     threshold: 0.75
-    #     th_on_preds: True
-    #     th_on_target: False
-    #     target_to_int: True
-
-    # - name: recall > 0.5
-    #   metric: recall
-    #   target_nan_mask: ignore-mean-label
-    #   threshold_kwargs: *threshold_1
-
-    # - name: recall > 0.25
-    #   metric: recall
-    #   target_nan_mask: ignore-mean-label
-    #   threshold_kwargs: *threshold_2
-
-    # - name: recall > 0.75
-    #   metric: recall
-    #   target_nan_mask: ignore-mean-label
-    #   threshold_kwargs: *threshold_3
-
-    # - name: precision > 0.5
-    #   metric: precision
-    #   target_nan_mask: ignore-mean-label
-    #   threshold_kwargs: *threshold_1
-
-    # - name: precision > 0.25
-    #   metric: precision
-    #   target_nan_mask: ignore-mean-label
-    #   threshold_kwargs: *threshold_2
-
-    # - name: precision > 0.25
-    #   metric: precision
-    #   target_nan_mask: ignore-mean-label
-    #   threshold_kwargs: *threshold_3
-
-trainer:
-  logger:
-    save_dir: logs/ipu-reproduce
-    name: *name
-  #early_stopping:
-  #  monitor: *monitor
-  #  min_delta: 0
-  #  patience: 10
-  #  mode: &mode min
-  model_checkpoint:
-    dirpath: models_checkpoints/ipu-reproduce/
-    filename: *name
-    #monitor: *monitor
-    #mode: *mode
-    save_top_k: 1
-    every_n_epochs: 1
-  trainer:
-    precision: 16
-    max_epochs: 1000
-    min_epochs: 100
-    accumulate_grad_batches: 1000
diff --git a/expts/configs/config_molPCQM4Mv2.yaml b/expts/configs/config_molPCQM4Mv2.yaml
deleted file mode 100644
index 4d4c62fa3..000000000
--- a/expts/configs/config_molPCQM4Mv2.yaml
+++ /dev/null
@@ -1,210 +0,0 @@
-# Testing the multitask pipeline with the QM9 dataset on IPU, by splitting it up into three tasks: homo, alpha and cv.
-constants:
-  name: &name pcqm4mv2_goli_ipu
-  seed: &seed 42
-  raise_train_error: true   # Whether the code should raise an error if it crashes during training
-  accelerator:
-    type: ipu  #cpu or ipu or gpu
-
-
-datamodule:
-  module_type: "GraphOGBDataModule"
-  args: # Matches that in the test_multitask_datamodule.py case.
-    task_specific_args:   # To be replaced by a new class "DatasetParams"
-      pcqm4mv2:
-        df: null
-        dataset_name: "ogbg-lsc-pcqm4mv2"
-        sample_size: null
-
-    # Featurization
-    prepare_dict_or_graph: pyg:graph
-    featurization_n_jobs: 32
-    featurization_progress: True
-    # featurization_backend: "threading"
-    featurization:
-      atom_property_list_onehot: [atomic-number, valence]
-      atom_property_list_float: [mass, electronegativity, in-ring]
-      edge_property_list: [bond-type-onehot, stereo, in-ring]
-      add_self_loop: False
-      explicit_H: False
-      use_bonds_weights: False
-      pos_encoding_as_features:
-        pos_types:
-          la_pos: &pos_enc  #use same name as pe_encoder
-            pos_type: laplacian_eigvec_eigval #laplacian_eigvec
-            num_pos: 3
-            normalization: "none"
-            disconnected_comp: True
-          rw_pos: #use same name as pe_encoder
-            pos_type: rwse
-            ksteps: 16
-      # pos_encoding_as_directions: *pos_enc # Only for DGN or directional pooling
-
-    # Data handling-related
-    batch_size_training: 30
-    batch_size_inference: 30
-    # cache_data_path: null
-
-    num_workers: 4
-
-    ipu_dataloader_training_opts:
-      max_num_nodes_per_graph: &max_nodes 16
-      max_num_edges_per_graph: &max_edges 34
-
-    ipu_dataloader_inference_opts:
-      max_num_nodes_per_graph: *max_nodes
-      max_num_edges_per_graph: *max_edges
-
-architecture:
-  model_type: FullGraphMultiTaskNetwork
-  pre_nn:   # Set as null to avoid a pre-nn network
-    out_dim: &gnn_dim 400
-    hidden_dims: *gnn_dim
-    depth: 3
-    activation: relu
-    last_activation: none
-    dropout: &dropout 0.1
-    normalization: &normalization "none"
-    last_normalization: *normalization
-    residual_type: none
-
-  pre_nn_edges:   # Set as null to avoid a pre-nn network
-    out_dim: 32
-    hidden_dims: 32
-    depth: 3
-    activation: relu
-    last_activation: none
-    dropout: *dropout
-    normalization: *normalization
-    last_normalization: *normalization
-    residual_type: none
-
-  pe_encoders:
-    out_dim: 100
-    pool: "sum" #"mean" "max"
-    last_norm: None #"batch_norm", "layer_norm"
-    encoders: #la_pos |  rw_pos
-      la_pos:  # Set as null to avoid a pre-nn network
-        encoder_type: "laplacian_pe"
-        input_keys: ["eigvecs", "eigvals"]
-        hidden_dim: 64
-        model_type: 'DeepSet' #'Transformer' or 'DeepSet'
-        num_layers: 1
-        num_layers_post: 0 # Num. layers to apply after pooling
-        dropout: 0.1
-        first_normalization: "none" #"batch_norm" or "layer_norm"
-      rw_pos:
-        encoder_type: "mlp"
-        input_keys: ["rwse"]
-        output_keys: "node"
-        hidden_dim: 64
-        num_layers: 1
-        dropout: 0.1
-        normalization: "none" #"batch_norm" or "layer_norm"
-        first_normalization: "none" #"batch_norm" or "layer_norm"
-
-
-  gnn:  # Set as null to avoid a post-nn network
-    out_dim: *gnn_dim
-    hidden_dims: *gnn_dim
-    depth: 16
-    activation: relu
-    last_activation: none
-    dropout: *dropout
-    normalization: *normalization
-    last_normalization: *normalization
-    residual_type: simple
-    pooling: [sum, mean, max]
-    virtual_node: 'none'
-    layer_type: 'pyg:gps' #pyg:gine #'pyg:gps' # pyg:gated-gcn, pyg:gine,pyg:gps
-    layer_kwargs:  # Parameters for the model itself. You could define dropout_attn: 0.1
-      mpnn_type: 'pyg:gine'
-      mpnn_kwargs: null
-        #out_dim_edges: 10
-      attn_type: "full-attention" #"none"
-      attn_kwargs: null
-
-
-      # out_dim_edges: 16
-      # aggregators: [mean, max]
-      # scalers: [identity, amplification, attenuation]
-      # num_heads: 3
-
-  post_nn: null
-    #out_dim: 1
-    #hidden_dims: 32
-    #depth: 2
-    #activation: relu
-    #last_activation: none
-    #dropout: *dropout
-    #normalization: *normalization
-    #last_normalization: "none"
-    #residual_type: none
-
-  task_heads:
-    - task_name: "pcqm4mv2"
-      out_dim: 1
-      hidden_dims: *gnn_dim
-      depth: 3                          # Not needed if we have hidden_dims
-      activation: relu
-      last_activation: none
-      dropout: *dropout
-      normalization: *normalization
-      last_normalization: "none"
-      residual_type: none
-
-#Task-specific
-predictor:
-  metrics_on_progress_bar:
-    pcqm4mv2: ["mae", "pearsonr"]
-  loss_fun:
-    pcqm4mv2: mse_ipu
-  random_seed: *seed
-  optim_kwargs:
-    lr: 1.e-5
-    # weight_decay: 1.e-7
-  torch_scheduler_kwargs:
-    #module_type: ReduceLROnPlateau
-    #factor: 0.5
-    #patience: 7
-  scheduler_kwargs: null
-  #  monitor: &monitor loss/val
-  #  mode: min
-  #  frequency: 1
-  target_nan_mask: null # null: no mask, 0: 0 mask, ignore: ignore nan values from loss
-  flag_kwargs:
-    n_steps: 0 #1
-    alpha: 0.0 #0.01
-
-# Task-specific
-metrics:
-  pcqm4mv2:
-    - name: mae
-      metric: mae
-      threshold_kwargs: null
-    - name: pearsonr
-      metric: pearsonr
-      threshold_kwargs: null
-      target_nan_mask: ignore-mean-label
-
-trainer:
-  logger:
-    save_dir: logs/QM9
-    name: *name
-  #early_stopping:
-  #  monitor: *monitor
-  #  min_delta: 0
-  #  patience: 10
-  #  mode: &mode min
-  model_checkpoint:
-    dirpath: models_checkpoints/QM9/
-    filename: *name
-    #monitor: *monitor
-    #mode: *mode
-    save_top_k: 1
-    every_n_epochs: 1
-  trainer:
-    precision: 32
-    max_epochs: 20 #50
-    min_epochs: 1 #5
-    accumulate_grad_batches: 4
diff --git a/expts/configs/config_mpnn_10M_b3lyp.yaml b/expts/configs/config_mpnn_10M_b3lyp.yaml
new file mode 100644
index 000000000..f4107a331
--- /dev/null
+++ b/expts/configs/config_mpnn_10M_b3lyp.yaml
@@ -0,0 +1,252 @@
+# Testing the mpnn only model with the PCQMv2 dataset on IPU.
+constants:
+  name: &name pcqm4mv2_mpnn_4layer
+  seed: &seed 42
+  raise_train_error: true   # Whether the code should raise an error if it crashes during training
+  accelerator:
+    type: cpu  # cpu or ipu or gpu
+
+datamodule:
+  module_type: "MultitaskFromSmilesDataModule"
+  # module_type: "FakeDataModule"  # Option to use generated data
+  args: # Matches that in the test_multitask_datamodule.py case.
+    task_specific_args:   # To be replaced by a new class "DatasetParams"
+      betagap:
+        df: null
+        df_path: goli/data/b3lyp/b3lyp_mini.parquet #goli/data/b3lyp/b3lyp_mini.parquet
+        # wget https://storage.googleapis.com/goli-public/datasets/b3lyp/b3lyp_mini.parquet
+        # or set path as https://storage.googleapis.com/goli-public/datasets/b3lyp/b3lyp_mini.parquet directly
+        smiles_col: "smiles"
+        label_cols: ["beta_gap"]
+        # sample_size: 30000 # use sample_size for test
+        # splits_path: goli/data/PCQM4Mv2/split_dict.pt  # Download with `wget https://storage.googleapis.com/datasets-public-research/PCQM4M/cxsmiles/split_dict.pt`
+        split_val: 0.1
+        split_test: 0.1
+      alphagap:
+        df: null
+        df_path: goli/data/b3lyp/b3lyp_mini.parquet #goli/data/b3lyp/b3lyp_mini.parquet
+        # wget https://storage.googleapis.com/goli-public/datasets/b3lyp/b3lyp_mini.parquet
+        # or set path as https://storage.googleapis.com/goli-public/datasets/b3lyp/b3lyp_mini.parquet directly
+        smiles_col: "smiles"
+        label_cols: ["alpha_gap"]
+        # sample_size: 30000 # use sample_size for test
+        # splits_path: goli/data/PCQM4Mv2/split_dict.pt  # Download with `wget https://storage.googleapis.com/datasets-public-research/PCQM4M/cxsmiles/split_dict.pt`
+        split_val: 0.1
+        split_test: 0.1
+
+    # Featurization
+    prepare_dict_or_graph: pyg:graph
+    featurization_n_jobs: 30
+    featurization_progress: True
+    featurization_backend: "loky"
+    featurization:
+    # OGB: ['atomic_num', 'degree', 'possible_formal_charge', 'possible_numH' (total-valence),
+    # 'possible_number_radical_e', 'possible_is_aromatic', 'possible_is_in_ring',
+    # 'num_chiral_centers (not included yet)']
+      atom_property_list_onehot: [atomic-number, group, period, total-valence]
+      atom_property_list_float: [degree, formal-charge, radical-electron, aromatic, in-ring]
+      # OGB: ['possible_bond_type', 'possible_bond_stereo', 'possible_is_in_ring']
+      edge_property_list: [bond-type-onehot, stereo, in-ring]
+      add_self_loop: False
+      explicit_H: False # if H is included
+      use_bonds_weights: False
+      pos_encoding_as_features: # encoder dropout 0.18
+        pos_types:
+          la_pos: &pos_enc
+            pos_type: laplacian_eigvec_eigval #laplacian_eigvec
+            num_pos: 8
+            normalization: "none" # nomrlization already applied on the eigen vectors
+            disconnected_comp: True # if eigen values/vector for disconnected graph are included
+          rw_pos: # use same name as pe_encoder
+            pos_type: rwse
+            ksteps: 16
+      # pos_encoding_as_directions: *pos_enc # Only for DGN or directional pooling
+
+    # Data handling-related
+    batch_size_training: 64
+    batch_size_inference: 16
+    # cache_data_path: .
+    num_workers: 0 # -1 to use all
+    persistent_workers: False # if use persistent worker at the start of each epoch.
+    # Using persistent_workers false might make the start of each epoch very long.
+    featurization_backend: "loky"
+
+    # ipu_dataloader_training_opts:
+    #   mode: async
+    #   max_num_nodes_per_graph: 20 # train max nodes: 20, max_edges: 54
+    #   max_num_edges_per_graph: 60
+
+    # ipu_dataloader_inference_opts:
+    #   mode: async
+    #   max_num_nodes_per_graph: 20 # valid max nodes: 51, max_edges: 118
+    #   max_num_edges_per_graph: 120
+    #   # test-dev max nodes: 50, max_edges: 116
+    #   # test-challenge max nodes: 51, max_edges: 106
+
+architecture:
+  model_type: FullGraphMultiTaskNetwork
+  mup_base_path: null
+  pre_nn:   # Set as null to avoid a pre-nn network
+    out_dim: 32
+    hidden_dims: 64
+    depth: 2
+    activation: relu
+    last_activation: none
+    dropout: &dropout 0.1
+    normalization: &normalization layer_norm
+    last_normalization: *normalization
+    residual_type: none
+
+  pre_nn_edges:   # Set as null to avoid a pre-nn network
+    out_dim: 16
+    hidden_dims: 32
+    depth: 2
+    activation: relu
+    last_activation: none
+    dropout: *dropout
+    normalization: *normalization
+    last_normalization: *normalization
+    residual_type: none
+
+  pe_encoders:
+    out_dim: 32
+    pool: "sum" #"mean" "max"
+    last_norm: None #"batch_norm", "layer_norm"
+    encoders: #la_pos |  rw_pos
+      la_pos:  # Set as null to avoid a pre-nn network
+        encoder_type: "laplacian_pe"
+        input_keys: ["eigvecs", "eigvals"]
+        output_keys: ["feat"]
+        hidden_dim: 64
+        out_dim: 32
+        model_type: 'DeepSet' #'Transformer' or 'DeepSet'
+        num_layers: 2
+        num_layers_post: 1 # Num. layers to apply after pooling
+        dropout: 0.1
+        first_normalization: "none" #"batch_norm" or "layer_norm"
+      rw_pos:
+        encoder_type: "mlp"
+        input_keys: ["rwse"]
+        output_keys: ["feat"]
+        hidden_dim: 64
+        out_dim: 32
+        num_layers: 2
+        dropout: 0.1
+        normalization: "layer_norm" #"batch_norm" or "layer_norm"
+        first_normalization: "layer_norm" #"batch_norm" or "layer_norm"
+
+
+  gnn:  # Set as null to avoid a post-nn network
+    out_dim: 32
+    hidden_dims: 32
+    depth: 4
+    activation: gelu
+    last_activation: none
+    dropout: 0.0
+    normalization: "layer_norm"
+    last_normalization: *normalization
+    residual_type: simple
+    pooling: [sum]
+    virtual_node: 'none'
+    layer_type: 'pyg:gps' #pyg:gine #'pyg:gps' # pyg:gated-gcn, pyg:gine,pyg:gps
+    layer_kwargs:  # Parameters for the model itself. You could define dropout_attn: 0.1
+      node_residual: false
+      mpnn_type: 'pyg:mpnnplus'
+      mpnn_kwargs:
+        in_dim: 32
+        out_dim: 32
+        in_dim_edges: 16
+        out_dim_edges: 16
+      attn_type: "none" # "full-attention", "none"
+      # biased_attention: false
+      attn_kwargs: null
+
+
+  post_nn: null
+
+  task_heads:
+    alphagap:
+      out_dim: 1
+      hidden_dims: 256
+      depth: 2                          # Not needed if we have hidden_dims
+      activation: relu
+      last_activation: none
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none
+    betagap:
+      out_dim: 1
+      hidden_dims: 256
+      depth: 2                          # Not needed if we have hidden_dims
+      activation: relu
+      last_activation: none
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none
+
+#Task-specific
+predictor:
+  metrics_on_progress_bar:
+    alphagap: ["mae", "pearsonr"]
+    betagap: ["mae", "pearsonr"]
+  loss_fun:
+    alphagap: mse_ipu
+    betagap: mse_ipu
+  random_seed: *seed
+  optim_kwargs:
+    lr: 4.e-4 # warmup can be scheduled using torch_scheduler_kwargs
+    # weight_decay: 1.e-7
+  torch_scheduler_kwargs:
+    module_type: WarmUpLinearLR
+    max_num_epochs: &max_epochs 100
+    warmup_epochs: 10
+    verbose: False
+  scheduler_kwargs:
+  #  monitor: &monitor homolumo/mae/train
+  #  mode: min
+  #  frequency: 1
+  target_nan_mask: null # null: no mask, 0: 0 mask, ignore: ignore nan values from loss
+  flag_kwargs:
+    n_steps: 0 # 1
+    alpha: 0.0 # 0.01
+
+# Task-specific
+metrics:
+  alphagap: &alpha_metrics
+    - name: mae
+      metric: mae_ipu
+      target_nan_mask: null
+      multitask_handling: flatten
+      threshold_kwargs: null
+    - name: pearsonr
+      metric: pearsonr_ipu
+      threshold_kwargs: null
+      target_nan_mask: null
+      multitask_handling: mean-per-label
+  betagap: *alpha_metrics
+
+trainer:
+  logger:
+    save_dir: logs/PCQMv2
+    name: *name
+    project: PCQMv2_mpnn
+  #early_stopping:
+  #  monitor: *monitor
+  #  min_delta: 0
+  #  patience: 10
+  #  mode: &mode min
+  model_checkpoint:
+    dirpath: models_checkpoints/PCMQv2/
+    filename: *name
+    #monitor: *monitor
+    #mode: *mode
+    save_top_k: 1
+    every_n_epochs: 100
+  trainer:
+    precision: 32
+    max_epochs: *max_epochs
+    min_epochs: 1
+    accumulate_grad_batches: 2
+    check_val_every_n_epoch: 20
diff --git a/expts/configs/config_pcqmv2_mpnn.yaml b/expts/configs/config_mpnn_10M_pcqm4m.yaml
similarity index 84%
rename from expts/configs/config_pcqmv2_mpnn.yaml
rename to expts/configs/config_mpnn_10M_pcqm4m.yaml
index 44580b59b..e99d76131 100644
--- a/expts/configs/config_pcqmv2_mpnn.yaml
+++ b/expts/configs/config_mpnn_10M_pcqm4m.yaml
@@ -4,7 +4,7 @@ constants:
   seed: &seed 42
   raise_train_error: true   # Whether the code should raise an error if it crashes during training
   accelerator:
-    type: ipu  # cpu or ipu or gpu
+    type: cpu  # cpu or ipu or gpu
 
 datamodule:
   module_type: "MultitaskFromSmilesDataModule"
@@ -13,16 +13,15 @@ datamodule:
     task_specific_args:   # To be replaced by a new class "DatasetParams"
       homolumo:
         df: null
-        df_path: goli/data/PCQM4Mv2/pcqm4mv2.csv #goli/data/PCQM4Mv2/pcqm4mv2.csv
-        # wget https://storage.googleapis.com/datasets-public-research/PCQM4M/cxsmiles/pcqm4mv2.csv
-        # or set path as https://storage.googleapis.com/datasets-public-research/PCQM4M/cxsmiles/pcqm4mv2.csv directly
+        df_path: goli/data/PCQM4M/pcqm4mv2-20k.csv
+        # wget https://storage.googleapis.com/datasets-public-research/PCQM4M/cxsmiles/pcqm4mv2-20k.csv
+        # or set path as https://storage.googleapis.com/datasets-public-research/PCQM4M/cxsmiles/pcqm4mv2-20k.csv directly
         smiles_col: "cxsmiles"
         label_cols: ["homo_lumo_gap"]
         # sample_size: 30000 # use sample_size for test
-        splits_path: goli/data/PCQM4Mv2/split_dict.pt  # Download with `wget https://storage.googleapis.com/datasets-public-research/PCQM4M/cxsmiles/split_dict.pt`
-        #goli/data/PCQM4Mv2/split_dict.pt
-        # goli/data/PCQM4Mv2/pcqm4m_split.csv
-        split_names: ["train", "val", "test"]
+        # splits_path: goli/data/PCQM4Mv2/split_dict.pt  # Download with `wget https://storage.googleapis.com/datasets-public-research/PCQM4M/cxsmiles/split_dict.pt`
+        split_val: 0.1
+        split_test: 0.1
 
     # Featurization
     prepare_dict_or_graph: pyg:graph
@@ -56,29 +55,29 @@ datamodule:
     batch_size_training: 64
     batch_size_inference: 16
     # cache_data_path: .
-    num_workers: 10 # -1 to use all
+    num_workers: 0 # -1 to use all
     persistent_workers: False # if use persistent worker at the start of each epoch.
     # Using persistent_workers false might make the start of each epoch very long.
     featurization_backend: "loky"
 
-    ipu_dataloader_training_opts:
-      mode: async
-      max_num_nodes_per_graph: 20 # train max nodes: 20, max_edges: 54
-      max_num_edges_per_graph: 60
+    # ipu_dataloader_training_opts:
+    #   mode: async
+    #   max_num_nodes_per_graph: 20 # train max nodes: 20, max_edges: 54
+    #   max_num_edges_per_graph: 60
 
-    ipu_dataloader_inference_opts:
-      mode: async
-      max_num_nodes_per_graph: 20 # valid max nodes: 51, max_edges: 118
-      max_num_edges_per_graph: 120
-      # test-dev max nodes: 50, max_edges: 116
-      # test-challenge max nodes: 51, max_edges: 106
+    # ipu_dataloader_inference_opts:
+    #   mode: async
+    #   max_num_nodes_per_graph: 20 # valid max nodes: 51, max_edges: 118
+    #   max_num_edges_per_graph: 120
+    #   # test-dev max nodes: 50, max_edges: 116
+    #   # test-challenge max nodes: 51, max_edges: 106
 
 architecture:
   model_type: FullGraphMultiTaskNetwork
   mup_base_path: null
   pre_nn:   # Set as null to avoid a pre-nn network
-    out_dim: 256
-    hidden_dims: 1024
+    out_dim: 32
+    hidden_dims: 64
     depth: 2
     activation: relu
     last_activation: none
@@ -88,8 +87,8 @@ architecture:
     residual_type: none
 
   pre_nn_edges:   # Set as null to avoid a pre-nn network
-    out_dim: 128
-    hidden_dims: 512
+    out_dim: 16
+    hidden_dims: 32
     depth: 2
     activation: relu
     last_activation: none
@@ -127,8 +126,8 @@ architecture:
 
 
   gnn:  # Set as null to avoid a post-nn network
-    out_dim: 256
-    hidden_dims: 256
+    out_dim: 32
+    hidden_dims: 32
     depth: 4
     activation: gelu
     last_activation: none
@@ -143,10 +142,10 @@ architecture:
       node_residual: false
       mpnn_type: 'pyg:mpnnplus'
       mpnn_kwargs:
-        in_dim: 256
-        out_dim: 256
-        in_dim_edges: 128
-        out_dim_edges: 128
+        in_dim: 32
+        out_dim: 32
+        in_dim_edges: 16
+        out_dim_edges: 16
       attn_type: "none" # "full-attention", "none"
       # biased_attention: false
       attn_kwargs: null
@@ -222,7 +221,7 @@ trainer:
     save_top_k: 1
     every_n_epochs: 100
   trainer:
-    precision: 16
+    precision: 32
     max_epochs: *max_epochs
     min_epochs: 1
     accumulate_grad_batches: 2
diff --git a/expts/configs/config_pcqmv2_gpspp.yaml b/expts/configs/config_pcqmv2_gpspp.yaml
deleted file mode 100644
index 3f9f7ecea..000000000
--- a/expts/configs/config_pcqmv2_gpspp.yaml
+++ /dev/null
@@ -1,261 +0,0 @@
-# Testing the multitask pipeline with the PCQMv2 dataset on IPU.
-constants:
-  name: &name pcqm
-  seed: &seed 42
-  raise_train_error: true   # Whether the code should raise an error if it crashes during training
-  accelerator:
-    type: ipu  #cpu or ipu or gpu
-
-datamodule:
-  module_type: "MultitaskFromSmilesDataModule"
-  args: # Matches that in the test_multitask_datamodule.py case.
-    task_specific_args:   # To be replaced by a new class "DatasetParams"
-      homo:
-        df: null
-        df_path: &df_path goli/data/PCQM4Mv2/pcqm4m-v2-train.csv  # https://storage.googleapis.com/datasets-public-research/PCQM4M/cxsmiles/pcqm4mv2-mini.csv
-        smiles_col: "cxsmiles" #"smiles"
-        label_cols: ["homo_lumo_gap"] #["homo"]
-        split_val: 0.2
-        split_test: 0.2
-        split_seed: *seed
-        splits_path: null                 # This may not always be provided
-        sample_size: null                 # This may not always be provided
-        idx_col: null                     # This may not always be provided
-        weights_col: null                 # This may not always be provided
-        weights_type: null                # This may not always be provided
-
-    # Featurization
-    prepare_dict_or_graph: pyg:graph
-    featurization_n_jobs: 32
-    featurization_progress: True
-    featurization_backend: loky
-    featurization:
-      atom_property_list_onehot: [atomic-number, valence]
-      atom_property_list_float: [mass, electronegativity, in-ring]
-      conformer_property_list: [positions_3d]
-      edge_property_list: [bond-type-onehot, stereo, in-ring]
-      add_self_loop: False
-      explicit_H: False
-      use_bonds_weights: False
-      pos_encoding_as_features:
-        pos_types:
-          la_pos: &pos_enc  #use same name as pe_encoder
-            pos_type: laplacian_eigvec_eigval #laplacian_eigvec
-            num_pos: 3
-            normalization: "none"
-            disconnected_comp: True
-          rw_pos: #use same name as pe_encoder
-            pos_type: rwse
-            ksteps: 16
-
-      # pos_encoding_as_directions: *pos_enc # Only for DGN or directional pooling
-
-    # Data handling-related
-    batch_size_training: 6
-    batch_size_inference: 6
-    cache_data_path: temp_cache
-
-    num_workers: -1
-
-    ipu_dataloader_training_opts:
-      max_num_nodes_per_graph: &max_num_nodes 30
-      max_num_edges_per_graph: &max_num_edges 60
-
-    ipu_dataloader_inference_opts:
-      max_num_nodes_per_graph: *max_num_nodes
-      max_num_edges_per_graph: *max_num_edges
-
-architecture:
-  model_type: FullGraphMultiTaskNetwork
-  mup_base_path: null
-  pre_nn:   # Set as null to avoid a pre-nn network
-    out_dim: 64
-    hidden_dims: 64
-    depth: 1
-    activation: relu
-    last_activation: none
-    dropout: &dropout 0.1
-    normalization: &normalization layer_norm
-    last_normalization: *normalization
-    residual_type: none
-
-  pre_nn_edges:   # Set as null to avoid a pre-nn network
-    out_dim: 32
-    hidden_dims: 32
-    depth: 1
-    activation: relu
-    last_activation: none
-    dropout: *dropout
-    normalization: *normalization
-    last_normalization: *normalization
-    residual_type: none
-
-  pe_encoders:
-    out_dim: 64
-    pool: "sum" #"mean" "max"
-    last_norm: None #"batch_norm", "layer_norm"
-    encoders: #la_pos |  rw_pos | gaussian_pos
-      la_pos:  # Set as null to avoid a pre-nn network
-        encoder_type: "laplacian_pe"
-        input_keys: ["eigvecs", "eigvals"]
-        output_keys: ["feat"]
-        hidden_dim: 64
-        model_type: 'DeepSet' #'Transformer' or 'DeepSet'
-        num_layers: 1
-        num_layers_post: 0 # Num. layers to apply after pooling
-        dropout: 0.1
-        first_normalization: "none" #"batch_norm" or "layer_norm"
-      rw_pos:
-        encoder_type: "mlp"
-        input_keys: ["rwse"]
-        output_keys: ["feat"]
-        hidden_dim: 64
-        num_layers: 1
-        dropout: 0.1
-        normalization: "none" #"batch_norm" or "layer_norm"
-        first_normalization: "layer_norm" #"batch_norm" or "layer_norm"
-      gaussian_pos:
-        encoder_type: "gaussian_kernel"
-        input_keys: ["positions_3d"]
-        output_keys: ["feat", "graph_gaussian_bias_3d"]
-        num_heads: 1
-        num_layers: 1 #2
-        embed_dim: 64
-        use_input_keys_prefix: False
-
-
-
-  gnn:  # Set as null to avoid a post-nn network
-    out_dim: 64
-    hidden_dims: 64
-    depth: 4
-    activation: gelu
-    last_activation: none
-    dropout: 0.0
-    normalization: "layer_norm"
-    last_normalization: *normalization
-    residual_type: simple
-    pooling: [sum]
-    virtual_node: "sum"
-    use_virtual_edges: True
-    layer_type: 'pyg:gps' #pyg:gine #'pyg:gps' # pyg:gated-gcn, pyg:gine,pyg:gps
-    layer_kwargs:  # Parameters for the model itself. You could define dropout_attn: 0.1
-      out_dim_edges: 32
-      node_residual: false
-      mpnn_type: 'pyg:mpnnplus'
-      mpnn_kwargs:
-        in_dim: 64
-        out_dim: 64
-        in_dim_edges: 32
-        out_dim_edges: 32
-        droppath_rate: 0.
-      attn_type: "full-attention" # "full-attention", "none"
-      biased_attention_key: "graph_gaussian_bias_3d" #same key from gaussian_pos/output_keys/"graph_xxx"
-      attn_kwargs: null
-      droppath_rate_attn: 0.3
-      droppath_rate_ffn: 0.3
-
-  post_nn:
-    out_dim: 32
-    hidden_dims: 400
-    depth: 2
-    activation: relu
-    last_activation: none
-    dropout: *dropout
-    normalization: *normalization
-    last_normalization: "none"
-    residual_type: none
-
-  task_heads:
-    homo:
-      out_dim: 1
-      hidden_dims: 32
-      depth: 2                          # Not needed if we have hidden_dims
-      activation: relu
-      last_activation: none
-      dropout: *dropout
-      normalization: *normalization
-      last_normalization: "none"
-      residual_type: none
-
-#Task-specific
-predictor:
-  metrics_on_progress_bar:
-    homo: ["mae", "pearsonr", "spearmanr", "f1 > 0.5", "precision > 0.5"]
-  loss_fun:  #<---- loss for the task head
-    homo: mse_ipu
-  random_seed: *seed
-  optim_kwargs:
-    lr: 1.e-3
-    # weight_decay: 1.e-7
-  torch_scheduler_kwargs:
-    #module_type: ReduceLROnPlateau
-    #factor: 0.5
-    #patience: 7
-  scheduler_kwargs: null
-  #  monitor: &monitor loss/val
-  #  mode: min
-  #  frequency: 1
-  target_nan_mask: null # null: no mask, 0: 0 mask, ignore: ignore nan values from loss
-  flag_kwargs:
-    n_steps: 0 #1
-    alpha: 0.0 #0.01
-
-# Task-specific
-metrics:
-  homo: # <--- task head
-    - name: mae
-      metric: mae_ipu
-      target_nan_mask: null
-      multitask_handling: flatten
-      threshold_kwargs: null
-    - name: pearsonr
-      metric: pearsonr_ipu
-      threshold_kwargs: null
-      target_nan_mask: null
-      multitask_handling: mean-per-label
-    - name: spearmanr
-      metric: spearmanr_ipu
-      threshold_kwargs: null
-      target_nan_mask: null
-      multitask_handling: mean-per-label
-    - name: f1 > 0.5
-      metric: f1_ipu
-      num_classes: 2
-      average: micro
-      target_nan_mask: null
-      multitask_handling: mean-per-label
-      threshold_kwargs: &threshold_05
-        operator: greater
-        threshold: 0.5
-        th_on_preds: True
-        th_on_target: True
-        target_to_int: False
-    - name: precision > 0.5
-      metric: precision_ipu
-      average: micro
-      target_nan_mask: null
-      multitask_handling: mean-per-label
-      threshold_kwargs: *threshold_05
-
-trainer:
-  logger:
-    save_dir: logs/PCQMv2
-    name: *name
-  #early_stopping:
-  #  monitor: *monitor
-  #  min_delta: 0
-  #  patience: 10
-  #  mode: &mode min
-  model_checkpoint:
-    dirpath: models_checkpoints/PCQMv2/
-    filename: *name
-    #monitor: *monitor
-    #mode: *mode
-    save_top_k: 1
-    every_n_epochs: 1
-  trainer:
-    precision: 32
-    max_epochs: 5 #50
-    min_epochs: 1 #5
-    accumulate_grad_batches: 1
diff --git a/expts/main_run_multitask.py b/expts/main_run_multitask.py
index 9be864b2c..2b6ee278a 100644
--- a/expts/main_run_multitask.py
+++ b/expts/main_run_multitask.py
@@ -28,10 +28,8 @@
 # Set up the working directory
 MAIN_DIR = dirname(dirname(abspath(goli.__file__)))
 
-# CONFIG_FILE = "expts/configs/config_ipu_qm9.yaml"
-# CONFIG_FILE = "expts/configs/config_gpu_qm9.yaml"
-# CONFIG_FILE = "expts/configs/config_pcqmv2_gpspp.yaml"
-CONFIG_FILE = "expts/configs/config_pcqmv2_mpnn.yaml"
+# CONFIG_FILE = "expts/configs/config_mpnn_10M_b3lyp.yaml"
+CONFIG_FILE = "expts/configs/config_mpnn_10M_pcqm4m.yaml"
 os.chdir(MAIN_DIR)
 
 
diff --git a/goli/config/_loader.py b/goli/config/_loader.py
index 9081f30fc..5cbfca482 100644
--- a/goli/config/_loader.py
+++ b/goli/config/_loader.py
@@ -194,10 +194,10 @@ def load_metrics(config: Union[omegaconf.DictConfig, Dict[str, Any]]) -> Dict[st
     """
 
     task_metrics = {}
-    cfg_metrics = deepcopy(config["metrics"])
+    cfg_metrics = config.get("metrics", None)
     if cfg_metrics is None:
         return task_metrics
-
+    cfg_metrics = {key: deepcopy(value) for key, value in config["metrics"].items()}
     # Wrap every metric in the class `MetricWrapper` to standardize them
     for task in cfg_metrics:
         task_metrics[task] = {}
diff --git a/goli/data/collate.py b/goli/data/collate.py
index 0ae1c5e20..549f1b550 100644
--- a/goli/data/collate.py
+++ b/goli/data/collate.py
@@ -169,7 +169,10 @@ def collate_labels(
         for this_label in labels:
             empty_task_labels = set(labels_size_dict.keys()) - set(this_label.keys())
             for task in empty_task_labels:
-                this_label[task] = torch.full((len(labels), *labels_size_dict[task]), torch.nan)
+                this_label[task] = torch.full([*labels_size_dict[task]], torch.nan)
+            for task in this_label.keys():
+                if not isinstance(task, torch.Tensor):
+                    this_label[task] = torch.as_tensor(this_label[task])
     labels_dict = default_collate(labels)
 
     return labels_dict
diff --git a/goli/data/datamodule.py b/goli/data/datamodule.py
index aa4e2dca1..e394eab4d 100644
--- a/goli/data/datamodule.py
+++ b/goli/data/datamodule.py
@@ -6,6 +6,8 @@
 import zipfile
 from copy import deepcopy
 import time
+from tqdm import tqdm
+import gc
 
 from loguru import logger
 import fsspec
@@ -14,6 +16,7 @@
 import pandas as pd
 import numpy as np
 import datamol as dm
+from fastparquet import ParquetFile
 
 from sklearn.model_selection import train_test_split
 
@@ -216,7 +219,7 @@ def predict_ds(self, value):
         """Set the dataset for the prediction"""
         self._predict_ds = value
 
-    def get_first_graph(self):
+    def get_fake_graph(self):
         raise NotImplementedError()
 
     # Private methods
@@ -252,22 +255,107 @@ def _read_csv(
         return df
 
     @staticmethod
-    def _read_parquet(path: str, **kwargs) -> pd.DataFrame:
+    def _get_data_file_type(path):
+        # Extract the extension
+        name, ext = os.path.splitext(path)
+
+        # support compressed files
+        _, ext2 = os.path.splitext(name)
+        if ext2 != "":
+            ext = f"{ext2}{ext}"
+
+        if ext.endswith((".parquet")):  # Support parquet files. Compression is implicit
+            return "parquet"
+        elif ".sdf" in ext:  # support compressed sdf files
+            return "sdf"
+        elif ".csv" in ext:  # support compressed csv files
+            return "csv"
+        elif ".tsv" in ext:  # support compressed tsv files
+            return "tsv"
+        elif ".pkl" in ext:  # support compressed pickle files
+            return "pkl"
+        elif ext.endswith(".pt"):  # Pytorch tensor files, used for storing index splits
+            return "pt"
+        else:
+            raise ValueError(f"unsupported file `{path}`")
+
+    @staticmethod
+    def _get_table_columns(path: str) -> List[str]:
         """
-        read the parquet file int a pandas dataframe
+        Get the columns of a table without reading all the data.
+        Might be slow to decompress the file if the file is compressed.
+
         Parameters:
-            path: path to the parquet file
-            kwargs: keyword arguments for pd.read_parquet
+            path: path to the table file
+
         Returns:
-            pd.DataFrame: the panda dataframe storing molecules
+            List[str]: the column names
         """
 
-        path = str(path)
+        datafile_type = BaseDataModule._get_data_file_type(path)
 
-        if path.startswith("goli://"):
-            path = goli_package_path(path)
+        if datafile_type == "parquet":
+            # Read the schema of a parquet file
+            file = ParquetFile(path)
+            schema = file.pandas_metadata["columns"]
+            column_names = [s["name"] for s in schema if s["name"] is not None]
+        elif datafile_type == "sdf":
+            df = BaseDataModule._read_sdf(path, max_num_mols=5, discard_invalud=True, n_jobs=1)
+            column_names = df.columns
+        elif datafile_type in ["csv", "tsv"]:
+            # Read the schema of a csv / tsv file
+            df = BaseDataModule._read_csv(path, nrows=5)
+            column_names = df.columns
+        return column_names
+
+    @staticmethod
+    def _read_parquet(path, **kwargs):
+        kwargs.pop("dtype", None)  # Only useful for csv
+        column_names = BaseDataModule._get_table_columns(path)
+
+        # Change the 'usecols' parameter to 'columns'
+        columns = kwargs.pop("columns", None)
+        if "usecols" in kwargs.keys():
+            assert columns is None, "Ambiguous value of `columns`"
+            columns = kwargs.pop("usecols")
+        if columns is None:
+            columns = column_names
+        for column in columns:
+            assert (
+                column in column_names
+            ), f"Column `{column}` is not in the parquet file with columns {column_names}"
+
+        # Read the parquet file per column, and convert the data to float16 to reduce memory consumption
+        all_series = {}
+        progress = tqdm(columns)
+        for col in progress:
+            # Read single column
+            progress.set_description(f"Reading parquet column `{col}`")
+            this_series = pd.read_parquet(path, columns=[col], engine="fastparquet", **kwargs)[col]
+
+            # Check if the data is float
+            first_elem = this_series.iloc[0]
+            is_float = False
+            if isinstance(first_elem, (list, tuple)):
+                is_float = isinstance(first_elem[0], np.floating)
+            elif isinstance(first_elem, np.ndarray):
+                is_float = isinstance(first_elem, np.floating)
+
+            # Convert floats to float16
+            if is_float:
+                if isinstance(first_elem, np.ndarray):
+                    this_series.update([elem.astype(np.float16) for elem in this_series])
+                elif isinstance(first_elem, list):
+                    this_series.update([np.asarray(elem).astype(np.float16) for elem in this_series])
+                else:
+                    this_series = this_series.astype(np.float16)
+
+            all_series[col] = this_series
+            gc.collect()  # Reset memory after each column
+
+        # Merge columns into a dataframe
+        df = pd.concat(all_series, axis=1)
 
-        df = pd.read_parquet(path)
         return df
 
     @staticmethod
@@ -317,7 +405,6 @@ def _read_sdf(path: str, mol_col_name: str = "_rdkit_molecule_obj", **kwargs):
 
         return df
 
-    @staticmethod
     def _read_table(self, path: str, **kwargs) -> pd.DataFrame:
         """
         a general read file function which determines if which function to use, either _read_csv or _read_parquet
@@ -327,10 +414,15 @@ def _read_table(self, path: str, **kwargs) -> pd.DataFrame:
         Returns:
             pd.DataFrame: the panda dataframe storing molecules
         """
-        if str(path).endswith((".parquet")):
-            return self._read_parquet(path)
+        file_type = self._get_data_file_type(path)
+        if file_type == "parquet":
+            return self._read_parquet(path, **kwargs)
+        elif file_type == "sdf":  # support compressed sdf files
+            return self._read_sdf(path, **kwargs)
+        elif file_type in ["csv", "tsv"]:  # support compressed csv and tsv files
+            return self._read_csv(path, **kwargs)
         else:
-            return self._read_csv(path)
+            raise ValueError(f"unsupported file `{path}`")
 
     def get_dataloader_kwargs(self, stage: RunningStage, shuffle: bool, **kwargs) -> Dict[str, Any]:
         """
@@ -772,7 +864,7 @@ def prepare_data(self):
                     + check_arg_iterator(args.weights_col, enforce_type=list)
                 )
                 label_dtype = {col: np.float32 for col in label_cols}
-                task_df[task] = self._read_csv(args.df_path, usecols=usecols, dtype=label_dtype)
+                task_df[task] = self._read_table(args.df_path, usecols=usecols, dtype=label_dtype)
 
             else:
                 label_cols = self._parse_label_cols(
@@ -1149,12 +1241,9 @@ def _parse_label_cols(
             the parsed label columns
         """
         if df is None:
-            # Only load the useful columns, as some dataset can be very large
-            # when loading all columns
-            data_frame = self._read_csv(df_path, nrows=0)
+            cols = BaseDataModule._get_table_columns(df_path)
         else:
-            data_frame = df
-        cols = list(data_frame.columns)
+            cols = list(df.columns)
 
         # A star `*` at the beginning or end of the string specifies to look for all
         # columns that starts/end with a specific string
@@ -1186,7 +1275,7 @@ def is_setup(self):
     @property
     def num_node_feats(self):
         """Return the number of node features in the first graph"""
-        graph = self.get_first_graph()
+        graph = self.get_fake_graph()
         num_feats = graph.feat.shape[1]
         return num_feats
 
@@ -1198,7 +1287,7 @@ def in_dims(self):
         raw positional encoding dimensions such eigval, eigvec, rwse and more
         """
 
-        graph = self.get_first_graph()
+        graph = self.get_fake_graph()
         if isinstance(graph, (GraphDict)):
             graph = graph.ndata
 
@@ -1217,56 +1306,23 @@ def in_dims(self):
     def num_edge_feats(self):
         """Return the number of edge features in the first graph"""
 
-        graph = self.get_first_graph()
+        graph = self.get_fake_graph()
         empty = torch.Tensor([])
         num_feats = graph.get("edge_feat", empty).shape[-1]
 
         return num_feats
 
-    def get_first_graph(self):
-        """
-        Low memory footprint method to get the first datapoint graph.
-        The first 10 rows of the data are read in case the first one has a featurization
-        error. If all 20 first element, then `None` is returned, otherwise the first
-        graph to not fail is returned.
+    def get_fake_graph(self):
         """
-        keys = list(self.task_dataset_processing_params.keys())
-        task = keys[0]
-        args = self.task_dataset_processing_params[task]
-        if args.df is None:
-            df = self._read_csv(args.df_path, nrows=20)
-        else:
-            df = args.df.iloc[0:20, :]
+        Low memory footprint method to get the featurization of a fake graph
+        without reading the dataset. Useful for getting the number of node/edge features.
 
-        label_cols = self._parse_label_cols(
-            df, df_path=None, label_cols=args.label_cols, smiles_col=args.smiles_col
-        )
-
-        smiles, labels, sample_idx, extras = self._extract_smiles_labels(
-            df,
-            smiles_col=args.smiles_col,
-            label_cols=label_cols,
-            idx_col=args.idx_col,
-            weights_col=args.weights_col,
-            weights_type=args.weights_type,
-        )
+        Returns:
+            graph: A fake graph with the right featurization
+        """
 
-        graph = None
-        for s in smiles:
-            graph = self.smiles_transformer(s, mask_nan=0.0)
-            if (graph is None) or isinstance(graph, str):
-                continue
-            else:
-                num_nodes = graph.num_nodes
-                num_edges = graph.num_edges
-            if (graph is not None) and not isinstance(graph, str) and (num_edges > 0) and (num_nodes > 0):
-                break
-        if graph is None:
-            raise ValueError("No valid graph found in the first 20 rows of the dataset")
-        elif isinstance(graph, str):
-            raise ValueError(
-                f"No valid graph found in the first 20 rows of the dataset, with error:\n {graph}"
-            )
+        smiles = "C1=CC=CC=C1"
+        graph = self.smiles_transformer(smiles, mask_nan=0.0)
         return graph
 
     ########################## Private methods ######################################
@@ -1411,18 +1467,18 @@ def _get_split_indices(
 
         else:
             # Split from an indices file
-            name, ext = os.path.splitext(splits_path)
-            _, ext2 = os.path.splitext(name)
-            if ext2 != "":
-                ext = f"{ext2}{ext}"
-            if ext == ".pt":
+            file_type = self._get_data_file_type(splits_path)
+
+            if file_type == "pt":
                 splits = torch.load(splits_path)
-            elif (".csv" in ext) or (".tsv" in ext):
+            elif file_type in ["csv", "tsv"]:
                 with fsspec.open(str(splits_path)) as f:
                     splits = self._read_csv(splits_path)
                 splits = splits.dropna()
             else:
-                raise ValueError(f"file extension {ext} not recognised, please use .pt or .csv.")
+                raise ValueError(
+                    f"file type `{file_type}` for `{splits_path}` not recognised, please use .pt, .csv or .tsv"
+                )
             train, val, test = split_names
             train_indices = splits[train].astype("int").tolist()
             val_indices = splits[val].astype("int").tolist()
@@ -2079,7 +2135,7 @@ def setup(self, stage=None):
         if default_labels_size_dict is None:
             self.collate_fn.keywords["labels_size_dict"] = labels_size
 
-    def get_first_graph(self):
+    def get_fake_graph(self):
         """
         Low memory footprint method to get the first datapoint DGL graph.
         The first 10 rows of the data are read in case the first one has a featurization
diff --git a/mkdocs.yml b/mkdocs.yml
index 3a7311408..7ff75b215 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -12,12 +12,14 @@ docs_dir: "docs"
 nav:
   - Overview: index.md
   - Tutorials:
-    - basics:
-      - Using GNN layers: tutorials/basics/using_gnn_layers.ipynb
-      - Implementing GNN layers: tutorials/basics/implementing_gnn_layers.ipynb
-      - Making GNN networks: tutorials/basics/making_gnn_networks.ipynb
-      - Convert CSV to Parquet files: tutorials/basics/csv_to_parquet.ipynb
-      - Reading Parquet files: tutorials/basics/parquet_reader.ipynb
+    - feature_processing:
+      - Add new positional encoding: tutorials/feature_processing/add_new_positional_encoding.ipynb
+      - Convert CSV to Parquet files: tutorials/feature_processing/csv_to_parquet.ipynb
+      - Timing parallel processing: tutorials/feature_processing/timing_parallel.ipynb
+    - gnn_layers:
+      - Add new gnn layers: tutorials/gnn/add_new_gnn_layers.ipynb
+      - Making GNN networks: tutorials/gnn/making_gnn_networks.ipynb
+      - Using GNN layers: tutorials/gnn/using_gnn_layers.ipynb
     - model_training:
       - Simple Molecular Model: tutorials/model_training/simple-molecular-model.ipynb
       - Training on IPU: tutorials/model_training/ipu_training.ipynb
diff --git a/news/collate_parquet_cleanup.rst b/news/collate_parquet_cleanup.rst
new file mode 100644
index 000000000..3651d0066
--- /dev/null
+++ b/news/collate_parquet_cleanup.rst
@@ -0,0 +1,12 @@
+**Added:**
+
+* Added 2 more tutorials
+
+**Changed:**
+
+* Deleted most outdated config files
+* Added a config file that everyone can use for multitasking
+* Fixed the parquet file loading (went to a very old branch that was never merged, and took most of the code)
+* Added deepcopy to the metrics due to using the &* notation in the yaml
+* Fixed the collating of labels from the latest commits
+* Checked that the model works with both CSV and parquet file, and created different configs
diff --git a/requirements.txt b/requirements.txt
index fb0e9d309..a6df98dd9 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -38,7 +38,7 @@ loguru
 appdirs
 matplotlib
 tk
-pytorch-lightning >1.8
+pytorch-lightning >1.8,<2.0
 datamol
 PyYAML
 wandb
@@ -46,3 +46,4 @@ optuna
 google-cloud-storage
 mup
 pyg-nightly
+fastparquet


Test patch:
diff --git a/tests/test_collate.py b/tests/test_collate.py
new file mode 100644
index 000000000..cc96dfd1c
--- /dev/null
+++ b/tests/test_collate.py
@@ -0,0 +1,75 @@
+"""
+Unit tests for the collate
+"""
+
+import unittest as ut
+from copy import deepcopy
+import torch
+import numpy as np
+
+from goli.data.collate import collate_labels, goli_collate_fn
+
+
+class test_Collate(ut.TestCase):
+    def test_collate_labels(self):
+        # Create fake labels
+        labels_size_dict = {"label1": [1], "label2": [5], "label3": [5, 2]}
+        fake_label = {
+            "label1": torch.FloatTensor([1]),
+            "label2": torch.FloatTensor([1, 2, 3, 4, 5]),
+            "label3": torch.FloatTensor([[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]]),
+        }
+        fake_labels = []
+        num_labels = 10
+        for i in range(num_labels):
+            fake_labels.append(deepcopy({key: val + 17 * i for key, val in fake_label.items()}))
+
+        # Collate labels and check for the right shapes
+        collated_labels = collate_labels(deepcopy(fake_labels), labels_size_dict)
+        self.assertEqual(collated_labels["label1"].shape, torch.Size([num_labels, 1]))
+        self.assertEqual(collated_labels["label2"].shape, torch.Size([num_labels, 5]))
+        self.assertEqual(collated_labels["label3"].shape, torch.Size([num_labels, 5, 2]))
+
+        # Check that the values are correct
+        label1_true = deepcopy(torch.stack([this_label["label1"] for this_label in fake_labels]))
+        label2_true = deepcopy(torch.stack([this_label["label2"] for this_label in fake_labels]))
+        label3_true = deepcopy(torch.stack([this_label["label3"] for this_label in fake_labels]))
+        np.testing.assert_array_equal(collated_labels["label1"].numpy(), label1_true.numpy())
+        np.testing.assert_array_equal(collated_labels["label2"].numpy(), label2_true.numpy())
+        np.testing.assert_array_equal(collated_labels["label3"].numpy(), label3_true.numpy())
+
+        # Remove some labels and check that the collation still works and puts `nan` in the right places
+        missing_labels = {"label1": [1, 3, 5, 7, 9], "label2": [0, 2, 4, 6, 8], "label3": [0, 1, 2, 3, 4]}
+        for key, missing_idx in missing_labels.items():
+            for idx in missing_idx:
+                fake_labels[idx].pop(key)
+        label1_true[missing_labels["label1"]] = float("nan")
+        label2_true[missing_labels["label2"]] = float("nan")
+        label3_true[missing_labels["label3"]] = float("nan")
+
+        # Collate labels and check for the right shapes
+        collated_labels = collate_labels(deepcopy(fake_labels), labels_size_dict)
+        self.assertEqual(collated_labels["label1"].shape, torch.Size([num_labels, 1]))
+        self.assertEqual(collated_labels["label2"].shape, torch.Size([num_labels, 5]))
+        self.assertEqual(collated_labels["label3"].shape, torch.Size([num_labels, 5, 2]))
+
+        # Check that the values are correct when some labels are missing
+        np.testing.assert_array_equal(collated_labels["label1"].numpy(), label1_true.numpy())
+        np.testing.assert_array_equal(collated_labels["label2"].numpy(), label2_true.numpy())
+        np.testing.assert_array_equal(collated_labels["label3"].numpy(), label3_true.numpy())
+
+        # Now test the `goli_collate_fn` function when only labels are given
+        fake_labels2 = [{"labels": this_label} for this_label in fake_labels]
+        collated_labels = goli_collate_fn(deepcopy(fake_labels2), labels_size_dict=labels_size_dict)["labels"]
+        self.assertEqual(collated_labels["label1"].shape, torch.Size([num_labels, 1]))
+        self.assertEqual(collated_labels["label2"].shape, torch.Size([num_labels, 5]))
+        self.assertEqual(collated_labels["label3"].shape, torch.Size([num_labels, 5, 2]))
+
+        # Check that the values are correct when some labels are missing
+        np.testing.assert_array_equal(collated_labels["label1"].numpy(), label1_true.numpy())
+        np.testing.assert_array_equal(collated_labels["label2"].numpy(), label2_true.numpy())
+        np.testing.assert_array_equal(collated_labels["label3"].numpy(), label3_true.numpy())
+
+
+if __name__ == "__main__":
+    ut.main()