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

Issue:
Transformer packing
Checklist:

- [x] Implementing packing within a local-batch
- [x] Implement different `to_dense_batch` and `to_sparse_batch` for packed and regular version
- [ ] Change the dataloading pipeline to pack the graphs
- [ ] Change the attention-bias masking when using packs
- [ ] Pre-define the tensor dimension when using packing
- [ ] Added a `news` entry: _copy `news/TEMPLATE.rst` to `news/my-feature-or-branch.rst`) and edit it._

---



Discussion / hints:

Fix patch:
diff --git a/expts/configs/config_ipu_qm9.yaml b/expts/configs/config_ipu_qm9.yaml
index 12c6494bc..940e3d1b2 100644
--- a/expts/configs/config_ipu_qm9.yaml
+++ b/expts/configs/config_ipu_qm9.yaml
@@ -14,7 +14,7 @@ datamodule:
       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
+        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
@@ -76,8 +76,9 @@ datamodule:
       # pos_encoding_as_directions: *pos_enc # Only for DGN or directional pooling
 
     # Data handling-related
-    batch_size_training: 6
-    batch_size_inference: 6
+    batch_size_training: 9
+    batch_size_inference: 9
+    batch_size_per_pack: 3
     cache_data_path: "datacache"
 
     num_workers: -1
@@ -92,6 +93,7 @@ datamodule:
 
 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
@@ -122,6 +124,7 @@ architecture:
       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
@@ -131,7 +134,7 @@ architecture:
       rw_pos:
         encoder_type: "mlp"
         input_keys: ["rwse"]
-        output_keys: "node"
+        output_keys: ["feat"]
         hidden_dim: 64
         num_layers: 1
         dropout: 0.1
@@ -316,4 +319,4 @@ trainer:
     precision: 32
     max_epochs: 50 #50
     min_epochs: 1 #5
-    accumulate_grad_batches: 4
+    accumulate_grad_batches: 2
diff --git a/goli/data/collate.py b/goli/data/collate.py
index fdfc20685..a21c90378 100644
--- a/goli/data/collate.py
+++ b/goli/data/collate.py
@@ -9,6 +9,7 @@
 from torch_geometric.data import Data, Batch
 
 from goli.features import GraphDict, to_dense_array
+from goli.utils.packing import fast_packing, get_pack_sizes, node_to_pack_indices_mask
 
 
 def goli_collate_fn(
@@ -52,8 +53,15 @@ def goli_collate_fn(
 
         do_not_batch_keys:
             Keys to ignore for the collate
+
+        batch_size_per_pack: The number of graphs to pack together.
+            This is useful for using packing with the Transformer.
+            If None, no packing is done.
+            Otherwise, indices are generated to map the nodes to the pack they belong to under the key `"pack_from_node_idx"`,
+            with an additional mask to indicate which nodes are from the same graph under the key `"pack_attn_mask"`.
+
     Returns:
-        A dictionary of the batch
+        The batched elements. See `torch.utils.data.dataloader.default_collate`.
     """
 
     elem = elements[0]
@@ -70,7 +78,7 @@ def goli_collate_fn(
             # If a PyG Graph is provided, use the PyG batching
             elif isinstance(elem[key], Data):
                 pyg_graphs = [d[key] for d in elements]
-                batch[key] = collage_pyg_graph(pyg_graphs)
+                batch[key] = collage_pyg_graph(pyg_graphs, batch_size_per_pack=batch_size_per_pack)
 
             # Ignore the collate for specific keys
             elif key in do_not_collate_keys:
@@ -124,6 +132,18 @@ def collage_pyg_graph(pyg_graphs: Iterable[Union[Data, Dict]], batch_size_per_pa
         pyg_graph.edge_index = pyg_graph.edge_index.to(torch.int64)
         pyg_batch.append(pyg_graph)
 
+    # Apply the packing at the mini-batch level. This is useful for using packing with the Transformer,
+    # especially in the case of the large graphs being much larger than the small graphs.
+    if batch_size_per_pack is not None:
+        num_nodes = [g.num_nodes for g in pyg_batch]
+        packed_graph_idx = fast_packing(num_nodes, batch_size_per_pack)
+
+        # Get the node to pack indices and the mask
+        pack_from_node_idx, pack_attn_mask = node_to_pack_indices_mask(packed_graph_idx, num_nodes)
+        for pyg_graph in pyg_batch:
+            pyg_graph.pack_from_node_idx = pack_from_node_idx
+            pyg_graph.pack_attn_mask = pack_attn_mask
+
     return Batch.from_data_list(pyg_batch)
 
 
diff --git a/goli/data/datamodule.py b/goli/data/datamodule.py
index 582fc7e08..e4c918b6e 100644
--- a/goli/data/datamodule.py
+++ b/goli/data/datamodule.py
@@ -592,6 +592,7 @@ def __init__(
         self,
         batch_size_training: int = 16,
         batch_size_inference: int = 16,
+        batch_size_per_pack: Optional[int] = None,
         num_workers: int = 0,
         pin_memory: bool = True,
         persistent_workers: bool = False,
@@ -612,6 +613,15 @@ def __init__(
 
         self.batch_size_training = batch_size_training
         self.batch_size_inference = batch_size_inference
+        self.batch_size_per_pack = batch_size_per_pack
+        if self.batch_size_per_pack is not None:
+            # Check that batch_size_per_pack is a divisor of batch_size_training and batch_size_inference
+            assert (
+                self.batch_size_training % self.batch_size_per_pack == 0
+            ), f"batch_size_training must be a multiple of batch_size_per_pack, provided batch_size_training={self.batch_size_training}, batch_size_per_pack={self.batch_size_per_pack}"
+            assert (
+                self.batch_size_inference % self.batch_size_per_pack == 0
+            ), f"batch_size_inference must be a multiple of batch_size_per_pack, provided batch_size_inference={self.batch_size_inference}, batch_size_per_pack={self.batch_size_per_pack}"
 
         self.num_workers = num_workers
         self.pin_memory = pin_memory
@@ -676,11 +686,10 @@ def predict_dataloader(self, **kwargs):
             **kwargs,
         )
 
-    @staticmethod
-    def get_collate_fn(collate_fn):
+    def get_collate_fn(self, collate_fn):
         if collate_fn is None:
             # Some values become `inf` when changing data type. `mask_nan` deals with that
-            collate_fn = partial(goli_collate_fn, mask_nan=0)
+            collate_fn = partial(goli_collate_fn, mask_nan=0, batch_size_per_pack=self.batch_size_per_pack)
             collate_fn.__name__ = goli_collate_fn.__name__
 
         return collate_fn
@@ -1102,6 +1111,7 @@ def __init__(
         featurization: Optional[Union[Dict[str, Any], omegaconf.DictConfig]] = None,
         batch_size_training: int = 16,
         batch_size_inference: int = 16,
+        batch_size_per_pack: Optional[int] = None,
         num_workers: int = 0,
         pin_memory: bool = True,
         persistent_workers: bool = False,
@@ -1194,6 +1204,7 @@ def __init__(
             self,
             batch_size_training=batch_size_training,
             batch_size_inference=batch_size_inference,
+            batch_size_per_pack=batch_size_per_pack,
             num_workers=num_workers,
             pin_memory=pin_memory,
             persistent_workers=persistent_workers,
@@ -1525,11 +1536,15 @@ def get_dataloader(
 
         return loader
 
-    @staticmethod
-    def get_collate_fn(collate_fn):
+    def get_collate_fn(self, collate_fn):
         if collate_fn is None:
             # Some values become `inf` when changing data type. `mask_nan` deals with that
-            collate_fn = partial(goli_collate_fn, mask_nan=0, do_not_collate_keys=["smiles", "mol_ids"])
+            collate_fn = partial(
+                goli_collate_fn,
+                mask_nan=0,
+                do_not_collate_keys=["smiles", "mol_ids"],
+                batch_size_per_pack=self.batch_size_per_pack,
+            )
             collate_fn.__name__ = goli_collate_fn.__name__
         return collate_fn
 
@@ -2140,6 +2155,7 @@ def to_dict(self) -> Dict[str, Any]:
         obj_repr["test_size"] = len(self.test_indices) if self.test_indices is not None else None
         obj_repr["batch_size_training"] = self.batch_size_training
         obj_repr["batch_size_inference"] = self.batch_size_inference
+        obj_repr["batch_size_per_pack"] = self.batch_size_per_pack
         obj_repr["num_node_feats"] = self.num_node_feats
         obj_repr["num_node_feats_with_positional_encoding"] = self.num_node_feats_with_positional_encoding
         obj_repr["num_edge_feats"] = self.num_edge_feats
@@ -2167,6 +2183,7 @@ def __init__(
         featurization: Optional[Union[Dict[str, Any], omegaconf.DictConfig]] = None,
         batch_size_training: int = 16,
         batch_size_inference: int = 16,
+        batch_size_per_pack: Optional[int] = None,
         num_workers: int = 0,
         pin_memory: bool = True,
         persistent_workers: bool = False,
@@ -2239,6 +2256,7 @@ def __init__(
         dm_args["featurization"] = featurization
         dm_args["batch_size_training"] = batch_size_training
         dm_args["batch_size_inference"] = batch_size_inference
+        dm_args["batch_size_per_pack"] = batch_size_per_pack
         dm_args["num_workers"] = num_workers
         dm_args["pin_memory"] = pin_memory
         dm_args["featurization_n_jobs"] = featurization_n_jobs
diff --git a/goli/ipu/ipu_dataloader.py b/goli/ipu/ipu_dataloader.py
index ab58b1717..a02c45173 100644
--- a/goli/ipu/ipu_dataloader.py
+++ b/goli/ipu/ipu_dataloader.py
@@ -1,6 +1,6 @@
 # Copyright (c) 2022 Graphcore Ltd. All rights reserved.
 
-from typing import Callable, Iterable, Optional, List
+from typing import Callable, Iterable, Optional, List, Tuple, Dict, Any, Union
 from copy import deepcopy
 from dataclasses import dataclass
 import numpy as np
@@ -12,6 +12,13 @@
 from torch_geometric.transforms import BaseTransform
 
 from goli.ipu.ipu_utils import import_poptorch
+from goli.utils.packing import (
+    fast_packing,
+    hybrid_packing,
+    get_pack_sizes,
+    node_to_pack_indices_mask,
+    estimate_max_pack_node_size,
+)
 
 
 @dataclass
@@ -110,15 +117,17 @@ def __init__(
         self.dataset_max_nodes_per_graph = dataset_max_nodes_per_graph
         self.dataset_max_edges_per_graph = dataset_max_edges_per_graph
 
-    def __call__(self, batch: Batch) -> Batch:
+    def __call__(
+        self, batch: List[Dict[str, Union[Data, Dict[str, Tensor]]]]
+    ) -> Dict[str, Union[Batch, Dict[str, Tensor], Any]]:
         """
-        padding option to pad each batch to be same size.
+        Stack tensors, batch the pyg graphs, and pad each tensor to be same size.
 
         Parameters:
-            batch: The batch of pyg-graphs to be padded
+            batch: The batch of data, including pyg-graphs `Data` and labels `Dict[str, Tensor]` to be padded
 
         Returns:
-            batch: The padded batch
+            out_batch: A dictionary where the graphs are batched and the labels or other Tensors are stacked
         """
 
         # Sort the batch such that large graphs are paired with small graphs
@@ -155,13 +164,12 @@ def __call__(self, batch: Batch) -> Batch:
             if isinstance(val, torch.Tensor):
                 stacked_features[key] = torch.stack([this_graph[key] for this_graph in out_graphs], dim=0)
 
-        # TODO: Make this more robust, instead of hard-coding the keys
         out_batch["features"] = stacked_features
-        out_batch["_batch_idx"] = torch.as_tensor(range(len(all_batches)), dtype=torch.int32).unsqueeze(-1)
         for key in all_batches[0].keys():
-            if key not in ("features", "labels", "_types_conversion"):
+            if key not in ("features", "labels"):
                 out_batch[key] = [this_batch[key] for this_batch in all_batches]
 
+        #
         for data_key, data_val in out_batch.items():
             if isinstance(data_val, Batch):
                 for sub_key, sub_val in data_val.items():
@@ -400,256 +408,3 @@ def __repr__(self) -> str:
         s += f"node_value={self.node_value}, "
         s += f"edge_value={self.edge_value})"
         return s
-
-
-class MolPack:
-    """
-    Class that keeps track of the number of atoms and indices that are added
-    to each pack. Useful when doing packing, or other forms of smart batching.
-    A pack is a batch, but with optimized memory consumption.
-    """
-
-    def __init__(self):
-        self.num_nodes = 0
-        self.num_graphs = 0
-        self.average_atom = 0
-        self.indices = []
-
-    def add_mol(self, num_nodes: int, idx: int) -> "MolPack":
-        """
-        Add a molecule and it's index to the batch
-
-        Parameters:
-            num_nodes: Number of atoms of the new molecule
-
-            idx: Index associated to the molecule
-        """
-        self.num_nodes += num_nodes
-        self.num_graphs += 1
-        self.average_atom = self.num_nodes / self.num_graphs
-        self.indices.append(idx)
-        return self
-
-    def expected_atoms(self, remaining_mean_num_nodes: float, batch_size: int) -> float:
-        """
-        Given a desired batch size, and given the remaining mean number of
-        atoms, find the expected number of atoms of the current batch when it is full
-
-        Parameters:
-            remaining_mean_num_nodes: Average number of atoms per molecule
-                left to be sampled and distributed across tasks.
-
-            batch_size: Desired batch size
-
-        Returns:
-            expected_atoms: The expected number of atoms in this batch if we
-                sample randomly the remaining molecules.
-        """
-        return self.num_nodes + ((batch_size - self.num_graphs) * remaining_mean_num_nodes)
-
-    def __repr__(self) -> str:
-        """
-        Print the main attributes of the current class
-        """
-        return f"{self.__class__.__name__}(m: {self.num_graphs},\ta: {self.num_nodes},\tav: {self.average_atom:.1f})"
-
-
-def smart_packing(num_nodes: List[int], batch_size: int) -> List[List[int]]:
-    """
-    Simple and fast algorithm for packing graphs such that each batch has roughly the
-    same number of atoms.
-    Has for-loop scalability issues `O(num_graphs * ipu_batch_size)` = `O(num_graphs^2 / batch_size)`
-
-    Parameters:
-        num_nodes: List of the number of atoms per molecule for the entire global batch.
-            Must be of length `batch_size * ipu_batch_size`.
-
-        batch_size: The batch size per iteration, considering a single device and single
-            forward pass.
-            The global batch size is `batch_size * device_iterations * replication_factor * gradient_accumulation`
-
-    Returns:
-        packed_indices: A list of packs, each containing a list of indices, such that
-            if we collect `num_nodes` from the indices, then each pack has roughly the
-            same total number of atoms.
-    """
-
-    # Sort the list
-    num_nodes = np.asarray(num_nodes)
-    argsort_num_nodes = np.argsort(num_nodes)
-    sorted_num_nodes = num_nodes[argsort_num_nodes]
-    ipu_batch_size = int(len(num_nodes) / batch_size)
-    sorted_num_nodes, initial_num_nodes = (
-        sorted_num_nodes[:-ipu_batch_size],
-        sorted_num_nodes[-ipu_batch_size:],
-    )
-    reverse_cumsum = np.sum(sorted_num_nodes) - np.cumsum(sorted_num_nodes) + sorted_num_nodes[-1]
-
-    # Start with the largest element in separate packs
-    mol_batches = [
-        MolPack().add_mol(initial_num_nodes[-ii - 1], argsort_num_nodes[-ii - 1])
-        for ii in range(ipu_batch_size)
-    ]
-
-    # Loop from smallest to largest molecule, and add each molecule to the pack with smallest expected sum
-    for ii, num_atom in enumerate(sorted_num_nodes):
-        remaining_mean = reverse_cumsum[ii] / (len(sorted_num_nodes) - ii)
-        max_expected, idx_max_expected = 0, 0
-        for jj, m in enumerate(mol_batches):
-            if m.num_graphs >= batch_size:
-                continue
-            expected = m.num_nodes + (
-                (batch_size - m.num_graphs) * remaining_mean
-            )  # Faster than calling m.expected_atoms
-            if expected > max_expected:
-                max_expected = expected
-                idx_max_expected = jj
-        mol_batches[idx_max_expected].add_mol(num_atom, argsort_num_nodes[ii])
-
-    packed_indices = [batch.indices for batch in mol_batches]
-
-    return packed_indices
-
-
-def fast_packing(num_nodes: List[int], batch_size: int) -> List[List[int]]:
-    """
-    Super fast algorithm for packing graphs such that each batch has roughly the
-    same number of atoms. Not as good as `smart_packing` but
-    faster and more scalable for-loop complexity of `O(batch_size)`.
-
-    Parameters:
-        num_nodes: List of the number of atoms per molecule for the entire global batch.
-            Must be of length `batch_size * ipu_batch_size`.
-
-        batch_size: The batch size per iteration, considering a single device and single
-            forward pass.
-            The global batch size is `batch_size * device_iterations * replication_factor * gradient_accumulation`
-
-    Returns:
-        packed_indices: A list of packs, each containing a list of indices, such that
-            if we collect `num_nodes` from the indices, then each pack has roughly the
-            same total number of atoms.
-    """
-    num_nodes = np.asarray(num_nodes)
-    argsort_num_nodes = np.argsort(num_nodes)
-    ipu_batch_size = int(len(num_nodes) / batch_size)
-
-    groups = []
-    for ii in range(batch_size):
-        group = argsort_num_nodes[ii * ipu_batch_size : (ii + 1) * ipu_batch_size]
-        np.random.shuffle(group)
-        groups.append(group)
-
-    packed_indices = np.stack(groups, axis=1).tolist()
-    return packed_indices
-
-
-def hybrid_packing(num_nodes: List[int], batch_size: int) -> List[List[int]]:
-    """
-    Uses a combination of the `smart_packing` `O(n^2)` on the most important data points,
-    and the `fast_packing` `O(n)` on the average-sized data points.
-
-    Depending on the expected complexity
-
-    Parameters:
-        num_nodes: List of the number of atoms per molecule for the entire global batch.
-            Must be of length `batch_size * ipu_batch_size`.
-
-        batch_size: The batch size per iteration, considering a single device and single
-            forward pass.
-            The global batch size is `batch_size * device_iterations * replication_factor * gradient_accumulation`
-
-    Returns:
-        packed_indices: A list of packs, each containing a list of indices, such that
-            if we collect `num_nodes` from the indices, then each pack has roughly the
-            same total number of atoms.
-    """
-
-    # Determine the parameters based on the complexity of the smart-packing.
-    # The bigger the complexity, the more the `fast_packing` algorithm becomes
-    # statistically powerful, and the more speed benefits it provides.
-    smart_packing_complexity = len(num_nodes) ** 2 / batch_size
-    if smart_packing_complexity < 1e4:
-        return smart_packing(num_nodes=num_nodes, batch_size=batch_size)
-    elif smart_packing_complexity < 1e6:
-        big, small = 3, 6
-    elif smart_packing_complexity < 1e8:
-        big, small = 2, 4
-    elif smart_packing_complexity < 1e10:
-        big, small = 1, 2
-    else:
-        return fast_packing(num_nodes=num_nodes, batch_size=batch_size)
-
-    # Small datasets benefit from smart-packing, without compute burden
-    ipu_batch_size = int(len(num_nodes) / batch_size)
-    if len(num_nodes) < (big + small) * ipu_batch_size:
-        return smart_packing(num_nodes=num_nodes, batch_size=batch_size)
-
-    # Sort the list
-    num_nodes = np.asarray(num_nodes)
-    argsort_num_nodes = np.argsort(num_nodes)
-
-    # Smallest and biggest graphs are often outliers and will benefit from the `smart_packing`
-    biggest_graphs = argsort_num_nodes[-big * ipu_batch_size :]
-    smallest_graphs = argsort_num_nodes[: small * ipu_batch_size]
-    big_n_small_graphs = np.concatenate([biggest_graphs, smallest_graphs])
-    big_n_small_packs = smart_packing(num_nodes[big_n_small_graphs], batch_size=big + small)
-    big_n_small_indices = [big_n_small_graphs[pack] for pack in big_n_small_packs]
-    big_n_small_nodes = [num_nodes[pack] for pack in big_n_small_indices]
-
-    # Medium graphs will be packed faster
-    medium_graphs = argsort_num_nodes[small * ipu_batch_size : -big * ipu_batch_size]
-    medium_packs = fast_packing(num_nodes[medium_graphs], batch_size=batch_size - big - small)
-    medium_indices = [medium_graphs[pack] for pack in medium_packs]
-    medium_nodes = [num_nodes[pack] for pack in medium_indices]
-
-    # Pack the big/small with the medium in a smart way
-    big_n_small_sort = np.argsort(np.sum(np.stack(big_n_small_nodes, axis=1), axis=0))
-    medium_sort = np.argsort(np.sum(np.stack(medium_nodes, axis=1), axis=0))
-    packed_indices = [
-        np.concatenate([medium_indices[medium_sort[ii]], big_n_small_indices[big_n_small_sort[-ii]]])
-        for ii in range(len(medium_indices))
-    ]
-
-    return packed_indices
-
-
-def get_pack_sizes(packed_indices, num_nodes):
-    """
-    Get the number of atoms of each pack
-    """
-    pack_sums = []
-    for pack in packed_indices:
-        pack_sum = 0
-        for idx in pack:
-            pack_sum += num_nodes[idx]
-        pack_sums.append(pack_sum)
-    return pack_sums
-
-
-def estimate_max_pack_node_size(num_nodes: Iterable[int], batch_size: int, combined_batch_size: int):
-    """
-    Estimate the value of `max_num_nodes`, which represents the maximum number of nodes
-    needed in a batch to fit the data.
-
-    Parameters:
-        num_nodes: Number of nodes for all the graphs in the dataset
-        batch_size: The regular batch size per IPU
-        combined_batch_size: batch_size * device_iterations
-                             * replication_factor * gradient_accumulation
-
-    """
-
-    # Estimate the packing size needed
-    rand_indices = np.arange(len(num_nodes))
-    np.random.shuffle(rand_indices)
-    max_pack_size = 0
-    for ii in range(0, len(num_nodes), combined_batch_size):
-        this_indices = rand_indices[ii : ii + combined_batch_size]
-        choice = num_nodes[this_indices]
-        if len(choice) == combined_batch_size:
-            packed_indices = hybrid_packing(choice, batch_size)
-            max_pack_size = max(max_pack_size, max(get_pack_sizes(packed_indices, num_nodes[this_indices])))
-    max_pack_size_per_graph = max_pack_size / batch_size
-
-    return max_pack_size, max_pack_size_per_graph
diff --git a/goli/ipu/to_dense_batch.py b/goli/ipu/to_dense_batch.py
index bad023418..c136fee45 100644
--- a/goli/ipu/to_dense_batch.py
+++ b/goli/ipu/to_dense_batch.py
@@ -12,6 +12,13 @@ def to_sparse_batch(x: Tensor, mask_idx: Tensor):
     return torch.index_select(x.reshape(-1, x.shape[-1]), 0, mask_idx)
 
 
+def to_sparse_batch_from_packed(x: Tensor, pack_from_node_idx: Tensor):
+    """
+    Reverse function of `to_packed_dense_batch`
+    """
+    return x[pack_from_node_idx[:, 0], pack_from_node_idx[:, 1]]
+
+
 def to_dense_batch(
     x: Tensor,
     batch: Optional[Tensor] = None,
@@ -63,7 +70,7 @@ def to_dense_batch(
     num_nodes = scatter_add(batch.new_ones(x.size(0)), batch, dim=0, dim_size=batch_size)
     cum_nodes = torch.cat([batch.new_zeros(1), num_nodes.cumsum(dim=0)])
 
-    if max_num_nodes_per_graph is None:
+    if max_num_nodes_per_graph is None:  # Must be provided on IPU
         max_num_nodes_per_graph = int(num_nodes.max())
 
     idx = torch.arange(batch.size(0), dtype=torch.long, device=x.device)
@@ -105,3 +112,48 @@ def to_dense_batch(
     mask = mask.view(batch_size, max_num_nodes_per_graph)
 
     return out, mask, idx  # Added `idx` as a return
+
+
+def to_packed_dense_batch(
+    x: Tensor,
+    pack_from_node_idx: Tensor,
+    pack_attn_mask: Tensor,
+    fill_value: float = 0.0,
+    max_num_nodes_per_pack: Optional[int] = None,
+) -> Tuple[Tensor, Tensor]:
+    r"""Given a sparse batch of node features
+    :math:`\mathbf{X} \in \mathbb{R}^{(N_1 + \ldots + N_B) \times F}` (with
+    :math:`N_i` indicating the number of nodes in graph :math:`i`), creates a
+    dense node feature tensor
+    :math:`\mathbf{X} \in \mathbb{R}^{B \times N_{\max} \times F}` (with
+    :math:`N_{\max} = \max_i^B N_i`).
+    In addition, a mask of shape :math:`\mathbf{M} \in \{ 0, 1 \}^{B \times
+    N_{\max}}` is returned, holding information about the existence of
+    fake-nodes in the dense representation.
+
+    Parameters: # TODO: Update docstring
+        x: Node feature matrix
+            :math:`\mathbf{X} \in \mathbb{R}^{(N_1 + \ldots + N_B) \times F}`.
+        batch: Batch vector
+            :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns each
+            node to a specific example. Must be ordered. (default: :obj:`None`)
+        fill_value: The value for invalid entries in the
+            resulting dense output tensor. (default: :obj:`0`)
+        max_num_nodes_per_graph: The size of the output node dimension.
+            (default: :obj:`None`)
+        batch_size: The batch size. (default: :obj:`None`)
+        drop_nodes_last_graph: Whether to drop the nodes of the last graphs that exceed
+            the `max_num_nodes_per_graph`. Useful when the last graph is a padding.
+
+    :rtype: (:class:`Tensor`, :class:`BoolTensor`)
+    """
+
+    if max_num_nodes_per_pack is None:  # Must be provided on IPU
+        max_num_nodes_per_pack = pack_attn_mask.shape[-1]
+
+    size = [pack_attn_mask[0], max_num_nodes_per_pack] + list(x.size())[1:]
+
+    out = x.new_full(size, fill_value)
+    out[pack_from_node_idx[:, 0], pack_from_node_idx[:, 1]] = x
+
+    return out
diff --git a/goli/nn/pyg_layers/gps_pyg.py b/goli/nn/pyg_layers/gps_pyg.py
index b4648bdfd..3020976dd 100644
--- a/goli/nn/pyg_layers/gps_pyg.py
+++ b/goli/nn/pyg_layers/gps_pyg.py
@@ -5,7 +5,7 @@
 from torch import Tensor
 from torch_geometric.data import Batch
 from goli.nn.base_graph_layer import BaseGraphModule
-from goli.nn.base_layers import FCLayer, MultiheadAttentionMup, MLP, get_activation_str
+from goli.nn.base_layers import FCLayer, MultiheadAttentionMup, MLP
 from goli.nn.pyg_layers import (
     GatedGCNPyg,
     GINConvPyg,
@@ -14,7 +14,12 @@
     MPNNPlusPyg,
 )
 from goli.utils.decorators import classproperty
-from goli.ipu.to_dense_batch import to_dense_batch, to_sparse_batch
+from goli.ipu.to_dense_batch import (
+    to_dense_batch,
+    to_sparse_batch,
+    to_packed_dense_batch,
+    to_sparse_batch_from_packed,
+)
 from goli.ipu.ipu_utils import is_running_on_ipu
 
 PYG_LAYERS_DICT = {
@@ -277,6 +282,62 @@ def _parse_attn_layer(
             attn_layer = attn_class(biased_attention_key, **attn_kwargs)
         return attn_layer
 
+    def _use_packing(self, batch: Batch) -> bool:
+        """
+        Check if we should use packing for the batch of graphs.
+        """
+        return "pack_from_node_idx" in batch.keys and "pack_attn_mask" in batch.keys
+
+    def _to_dense_batch(
+        self,
+        h: Tensor,
+        batch: Batch,
+        batch_size: Optional[int] = None,
+        max_num_nodes_per_graph: Optional[int] = None,
+        on_ipu: bool = False,
+    ) -> Tensor:
+        """
+        Convert the batch of graphs to a dense batch.
+        """
+
+        if self._use_packing(batch):
+            attn_mask = batch.pack_attn_mask
+            key_padding_mask = None
+            idx = batch.pack_from_node_idx
+            h_dense = to_packed_dense_batch(
+                h,
+                pack_from_node_idx=idx,
+                pack_attn_mask=attn_mask,
+                max_num_nodes_per_pack=100,  # TODO: This should be a parameter
+            )
+        else:
+            attn_mask = None
+            h_dense, key_padding_mask, idx = to_dense_batch(
+                h,
+                batch=batch.batch,  # The batch index as a vector that indicates for nodes of which graph it belongs to
+                batch_size=batch_size,
+                max_num_nodes_per_graph=max_num_nodes_per_graph,
+                drop_nodes_last_graph=on_ipu,
+            )
+            key_padding_mask = ~key_padding_mask
+        return h_dense, attn_mask, key_padding_mask, idx
+
+    def _to_sparse_batch(self, batch: Batch, h_dense: Tensor, idx: Tensor) -> Tensor:
+        """
+        Convert the dense batch back to a sparse batch.
+        """
+        if self._use_packing(batch):
+            h = to_sparse_batch_from_packed(
+                h_dense,
+                pack_from_node_idx=idx,
+            )
+        else:
+            h = to_sparse_batch(
+                h_dense,
+                mask_idx=idx,
+            )
+        return h
+
     def _self_attention_block(self, feat: Tensor, feat_in: Tensor, batch: Batch) -> Tensor:
         """
         Applying the multi-head self-attention to the batch of graphs.
@@ -294,24 +355,26 @@ def _self_attention_block(self, feat: Tensor, feat_in: Tensor, batch: Batch) ->
         # Convert the tensor to a dense batch, then back to a sparse batch
         batch_size = None if feat.device.type != "ipu" else batch.graph_is_true.shape[0]
 
-        # feat[num_nodes, hidden_dim] -> feat_dense[num_graphs, max_num_nodes, hidden_dim]
-        h_dense, mask, idx = to_dense_batch(
+        # h[num_nodes, hidden_dim] -> h_dense[num_graphs, max_num_nodes, hidden_dim]
+        feat_dense, attn_mask, key_padding_mask, idx = self._to_dense_batch(
             feat,
-            batch=batch.batch,  # The batch index as a vector that indicates for nodes of which graph it belongs to
+            batch=batch,  # The batch index as a vector that indicates for nodes of which graph it belongs to
             batch_size=batch_size,
             max_num_nodes_per_graph=max_num_nodes_per_graph,
-            drop_nodes_last_graph=on_ipu,
+            on_ipu=on_ipu,
         )
 
         attn_bias = None
         if self.biased_attention_key is not None:
             attn_bias = batch[self.biased_attention_key]
 
-        # feat_dense[num_graphs, max_num_nodes, hidden_dim] -> feat_attn[num_graphs, max_num_nodes, hidden_dim]
-        feat_attn = self._sa_block(h_dense, attn_bias=attn_bias, attn_mask=None, key_padding_mask=~mask)
+        # h_dense[num_graphs, max_num_nodes, hidden_dim] -> feat_attn[num_graphs, max_num_nodes, hidden_dim]
+        feat_attn = self._sa_block(
+            feat_dense, attn_bias=attn_bias, attn_mask=attn_mask, key_padding_mask=key_padding_mask
+        )
 
         # feat_attn[num_graphs, max_num_nodes, hidden_dim] -> feat_attn[num_nodes, hidden_dim]
-        feat_attn = to_sparse_batch(feat_attn, idx)
+        feat_attn = self._to_sparse_batch(batch, feat_attn, idx)
 
         # Dropout, residual, norm
         if self.dropout_attn is not None:
diff --git a/goli/nn/pyg_layers/pooling_pyg.py b/goli/nn/pyg_layers/pooling_pyg.py
index 28d52da80..ddaecfb7e 100644
--- a/goli/nn/pyg_layers/pooling_pyg.py
+++ b/goli/nn/pyg_layers/pooling_pyg.py
@@ -87,9 +87,9 @@ def forward(self, g: Batch, feature: Tensor, *args, **kwargs):
             the pooled features
         """
         dim_size = g.num_graphs
-        if self.feat_type is "node":
+        if self.feat_type == "node":
             index = g.batch
-        elif self.feat_type is "edge":
+        elif self.feat_type == "edge":
             index = g.batch[g.edge_index][0]
         else:
             index = g.batch
diff --git a/goli/utils/packing.py b/goli/utils/packing.py
new file mode 100644
index 000000000..16ff1fc0f
--- /dev/null
+++ b/goli/utils/packing.py
@@ -0,0 +1,322 @@
+# Copyright (c) 2022 Graphcore Ltd. All rights reserved.
+
+from typing import List, Tuple, Iterable, Optional
+import numpy as np
+import torch
+
+
+class MolPack:
+    """
+    Class that keeps track of the number of atoms and indices that are added
+    to each pack. Useful when doing packing, or other forms of smart batching.
+    A pack is a batch, but with optimized memory consumption.
+    """
+
+    def __init__(self):
+        self.num_nodes = 0
+        self.num_graphs = 0
+        self.average_atom = 0
+        self.indices = []
+
+    def add_mol(self, num_nodes: int, idx: int) -> "MolPack":
+        """
+        Add a molecule and it's index to the batch
+
+        Parameters:
+            num_nodes: Number of atoms of the new molecule
+
+            idx: Index associated to the molecule
+        """
+        self.num_nodes += num_nodes
+        self.num_graphs += 1
+        self.average_atom = self.num_nodes / self.num_graphs
+        self.indices.append(idx)
+        return self
+
+    def expected_atoms(self, remaining_mean_num_nodes: float, batch_size: int) -> float:
+        """
+        Given a desired batch size, and given the remaining mean number of
+        atoms, find the expected number of atoms of the current batch when it is full
+
+        Parameters:
+            remaining_mean_num_nodes: Average number of atoms per molecule
+                left to be sampled and distributed across tasks.
+
+            batch_size: Desired batch size
+
+        Returns:
+            expected_atoms: The expected number of atoms in this batch if we
+                sample randomly the remaining molecules.
+        """
+        return self.num_nodes + ((batch_size - self.num_graphs) * remaining_mean_num_nodes)
+
+    def __repr__(self) -> str:
+        """
+        Print the main attributes of the current class
+        """
+        return f"{self.__class__.__name__}(m: {self.num_graphs},\ta: {self.num_nodes},\tav: {self.average_atom:.1f})"
+
+
+def smart_packing(num_nodes: List[int], batch_size: int) -> List[List[int]]:
+    """
+    Simple and fast algorithm for packing graphs such that each batch has roughly the
+    same number of atoms.
+    Has for-loop scalability issues `O(num_graphs * ipu_batch_size)` = `O(num_graphs^2 / batch_size)`
+
+    Parameters:
+        num_nodes: List of the number of atoms per molecule for the entire global batch.
+            Must be of length `batch_size * ipu_batch_size`.
+
+        batch_size: The batch size per iteration, considering a single device and single
+            forward pass.
+            The global batch size is `batch_size * device_iterations * replication_factor * gradient_accumulation`
+
+    Returns:
+        packed_indices: A list of packs, each containing a list of indices, such that
+            if we collect `num_nodes` from the indices, then each pack has roughly the
+            same total number of atoms.
+    """
+
+    # Sort the list
+    num_nodes = np.asarray(num_nodes)
+    argsort_num_nodes = np.argsort(num_nodes)
+    sorted_num_nodes = num_nodes[argsort_num_nodes]
+    ipu_batch_size = int(len(num_nodes) / batch_size)
+    sorted_num_nodes, initial_num_nodes = (
+        sorted_num_nodes[:-ipu_batch_size],
+        sorted_num_nodes[-ipu_batch_size:],
+    )
+    reverse_cumsum = np.sum(sorted_num_nodes) - np.cumsum(sorted_num_nodes) + sorted_num_nodes[-1]
+
+    # Start with the largest element in separate packs
+    mol_batches = [
+        MolPack().add_mol(initial_num_nodes[-ii - 1], argsort_num_nodes[-ii - 1])
+        for ii in range(ipu_batch_size)
+    ]
+
+    # Loop from smallest to largest molecule, and add each molecule to the pack with smallest expected sum
+    for ii, num_atom in enumerate(sorted_num_nodes):
+        remaining_mean = reverse_cumsum[ii] / (len(sorted_num_nodes) - ii)
+        max_expected, idx_max_expected = 0, 0
+        for jj, m in enumerate(mol_batches):
+            if m.num_graphs >= batch_size:
+                continue
+            expected = m.num_nodes + (
+                (batch_size - m.num_graphs) * remaining_mean
+            )  # Faster than calling m.expected_atoms
+            if expected > max_expected:
+                max_expected = expected
+                idx_max_expected = jj
+        mol_batches[idx_max_expected].add_mol(num_atom, argsort_num_nodes[ii])
+
+    packed_indices = [batch.indices for batch in mol_batches]
+
+    return packed_indices
+
+
+def fast_packing(num_nodes: List[int], batch_size: int) -> List[List[int]]:
+    """
+    Super fast algorithm for packing graphs such that each batch has roughly the
+    same number of atoms. Not as good as `smart_packing` but
+    faster and more scalable for-loop complexity of `O(batch_size)`.
+
+    Parameters:
+        num_nodes: List of the number of atoms per molecule for the entire global batch.
+            Must be of length `batch_size * ipu_batch_size`.
+
+        batch_size: The batch size per iteration, considering a single device and single
+            forward pass.
+            The global batch size is `batch_size * device_iterations * replication_factor * gradient_accumulation`
+
+    Returns:
+        packed_indices: A list of packs, each containing a list of indices, such that
+            if we collect `num_nodes` from the indices, then each pack has roughly the
+            same total number of atoms.
+    """
+    num_nodes = np.asarray(num_nodes)
+    argsort_num_nodes = np.argsort(num_nodes)
+    ipu_batch_size = int(len(num_nodes) / batch_size)
+
+    groups = []
+    for ii in range(batch_size):
+        group = argsort_num_nodes[ii * ipu_batch_size : (ii + 1) * ipu_batch_size]
+        np.random.shuffle(group)
+        groups.append(group)
+
+    packed_indices = np.stack(groups, axis=1).tolist()
+    return packed_indices
+
+
+def hybrid_packing(num_nodes: List[int], batch_size: int) -> List[List[int]]:
+    """
+    Uses a combination of the `smart_packing` `O(n^2)` on the most important data points,
+    and the `fast_packing` `O(n)` on the average-sized data points.
+
+    Depending on the expected complexity
+
+    Parameters:
+        num_nodes: List of the number of atoms per molecule for the entire global batch.
+            Must be of length `batch_size * ipu_batch_size`.
+
+        batch_size: The batch size per iteration, considering a single device and single
+            forward pass.
+            The global batch size is `batch_size * device_iterations * replication_factor * gradient_accumulation`
+
+    Returns:
+        packed_indices: A list of packs, each containing a list of indices, such that
+            if we collect `num_nodes` from the indices, then each pack has roughly the
+            same total number of atoms.
+    """
+
+    # Determine the parameters based on the complexity of the smart-packing.
+    # The bigger the complexity, the more the `fast_packing` algorithm becomes
+    # statistically powerful, and the more speed benefits it provides.
+    smart_packing_complexity = len(num_nodes) ** 2 / batch_size
+    if smart_packing_complexity < 1e4:
+        return smart_packing(num_nodes=num_nodes, batch_size=batch_size)
+    elif smart_packing_complexity < 1e6:
+        big, small = 3, 6
+    elif smart_packing_complexity < 1e8:
+        big, small = 2, 4
+    elif smart_packing_complexity < 1e10:
+        big, small = 1, 2
+    else:
+        return fast_packing(num_nodes=num_nodes, batch_size=batch_size)
+
+    # Small datasets benefit from smart-packing, without compute burden
+    ipu_batch_size = int(len(num_nodes) / batch_size)
+    if len(num_nodes) < (big + small) * ipu_batch_size:
+        return smart_packing(num_nodes=num_nodes, batch_size=batch_size)
+
+    # Sort the list
+    num_nodes = np.asarray(num_nodes)
+    argsort_num_nodes = np.argsort(num_nodes)
+
+    # Smallest and biggest graphs are often outliers and will benefit from the `smart_packing`
+    biggest_graphs = argsort_num_nodes[-big * ipu_batch_size :]
+    smallest_graphs = argsort_num_nodes[: small * ipu_batch_size]
+    big_n_small_graphs = np.concatenate([biggest_graphs, smallest_graphs])
+    big_n_small_packs = smart_packing(num_nodes[big_n_small_graphs], batch_size=big + small)
+    big_n_small_indices = [big_n_small_graphs[pack] for pack in big_n_small_packs]
+    big_n_small_nodes = [num_nodes[pack] for pack in big_n_small_indices]
+
+    # Medium graphs will be packed faster
+    medium_graphs = argsort_num_nodes[small * ipu_batch_size : -big * ipu_batch_size]
+    medium_packs = fast_packing(num_nodes[medium_graphs], batch_size=batch_size - big - small)
+    medium_indices = [medium_graphs[pack] for pack in medium_packs]
+    medium_nodes = [num_nodes[pack] for pack in medium_indices]
+
+    # Pack the big/small with the medium in a smart way
+    big_n_small_sort = np.argsort(np.sum(np.stack(big_n_small_nodes, axis=1), axis=0))
+    medium_sort = np.argsort(np.sum(np.stack(medium_nodes, axis=1), axis=0))
+    packed_indices = [
+        np.concatenate([medium_indices[medium_sort[ii]], big_n_small_indices[big_n_small_sort[-ii]]])
+        for ii in range(len(medium_indices))
+    ]
+
+    return packed_indices
+
+
+def get_pack_sizes(packed_indices, num_nodes):
+    """
+    Get the number of atoms of each pack
+    """
+    pack_sums = []
+    for pack in packed_indices:
+        pack_sum = 0
+        for idx in pack:
+            pack_sum += num_nodes[idx]
+        pack_sums.append(pack_sum)
+    return pack_sums
+
+
+def estimate_max_pack_node_size(num_nodes: Iterable[int], batch_size: int, combined_batch_size: int):
+    """
+    Estimate the value of `max_num_nodes`, which represents the maximum number of nodes
+    needed in a batch to fit the data.
+
+    Parameters:
+        num_nodes: Number of nodes for all the graphs in the dataset
+        batch_size: The regular batch size per IPU
+        combined_batch_size: batch_size * device_iterations
+                             * replication_factor * gradient_accumulation
+
+    """
+
+    # Estimate the packing size needed
+    rand_indices = np.arange(len(num_nodes))
+    np.random.shuffle(rand_indices)
+    max_pack_size = 0
+    for ii in range(0, len(num_nodes), combined_batch_size):
+        this_indices = rand_indices[ii : ii + combined_batch_size]
+        choice = num_nodes[this_indices]
+        if len(choice) == combined_batch_size:
+            packed_indices = hybrid_packing(choice, batch_size)
+            max_pack_size = max(max_pack_size, max(get_pack_sizes(packed_indices, num_nodes[this_indices])))
+    max_pack_size_per_graph = max_pack_size / batch_size
+
+    return max_pack_size, max_pack_size_per_graph
+
+
+def node_to_pack_indices_mask(
+    packed_indices: Iterable[Iterable[int]], all_num_nodes: Iterable[int], max_pack_size: Optional[int] = None
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Given a list of packed indices, and the number of nodes in each graph,
+    return a tensor of shape (sum(all_num_nodes), 2) where the first column
+    is the pack index, and the second column is the node index within the pack.
+
+    Can be used to generate a dense packing of the nodes as follows:
+    ```
+    # node_features: A tensor of shape (num_nodes, num_node_features)
+    # num_packs: The number of packs desired
+    # max_nodes_per_pack: The maximum number of nodes per pack
+    # dense_pack: A tensor of shape (num_packs, max_nodes_per_pack, num_node_features)
+
+    dense_pack = torch.zeros([num_packs, max_nodes_per_pack, num_node_features])
+    dense_pack[pack_from_node_idx[:, 0], pack_from_node_idx[:, 1]] = node_features
+    ```
+
+    This is useful when using a Transformer, to avoid wasteful padding when the
+    the longest sequence is much longer than the average sequence length.
+
+    Parameters:
+        packed_indices: A list of lists of graph indices, where each sub-list
+                        represents a pack of graphs
+        all_num_nodes: The number of nodes in each graph
+        max_pack_size: The maximum number of nodes per pack. If None, will be
+                          infered from the provided packs.
+                          Useful to determine the shape of the `pack_attn_mask`.
+
+    Returns:
+        pack_from_node_idx: A tensor of shape (num_nodes, 2) where the first column
+                        is the pack index, and the second column is the node index within the pack.
+
+        pack_attn_mask: A tensor of shape (num_packs, max_pack_size, max_pack_size),
+                            that represents the attention masking for each pack,
+                            such that the graphs in the pack are masked out from each other.
+    """
+
+    all_num_nodes = torch.as_tensor(all_num_nodes, dtype=torch.long)
+    cumsum_num_nodes = torch.cumsum(all_num_nodes, dim=0)
+    if max_pack_size is None:
+        pack_sizes = get_pack_sizes(packed_indices, all_num_nodes)
+        max_pack_size = max(pack_sizes)
+
+    # Get the node indices associated to the packs, with 0 padding
+    pack_from_node_idx = torch.zeros(sum(all_num_nodes), 2, dtype=torch.long)
+    pack_attn_mask = []  # masks for the attention
+    for ii, pack in enumerate(packed_indices):
+        jj = 0  # Counter for the number of nodes in the pack
+        this_pack_attn_mask = torch.ones((max_pack_size, max_pack_size), dtype=torch.bool)
+        for graph_idx in pack:
+            num_nodes = all_num_nodes[graph_idx]
+            node_idx = torch.arange(cumsum_num_nodes[graph_idx] - num_nodes, cumsum_num_nodes[graph_idx])
+            this_pack_attn_mask[jj : jj + num_nodes, jj : jj + num_nodes] = False
+            pack_from_node_idx[node_idx, 0] = ii
+            pack_from_node_idx[node_idx, 1] = jj + torch.arange(num_nodes)
+            jj += num_nodes
+        pack_attn_mask.append(this_pack_attn_mask)
+    pack_attn_mask = torch.stack(pack_attn_mask, dim=0)
+
+    return pack_from_node_idx, pack_attn_mask


Test patch:
diff --git a/tests/test_ipu_dataloader.py b/tests/test_ipu_dataloader.py
index 280696909..9037301f4 100644
--- a/tests/test_ipu_dataloader.py
+++ b/tests/test_ipu_dataloader.py
@@ -12,9 +12,7 @@
 from torch.utils.data.dataloader import default_collate
 
 # Current library imports
-from goli.ipu.ipu_dataloader import smart_packing, get_pack_sizes, fast_packing, hybrid_packing
 from goli.config._loader import load_datamodule, load_metrics, load_architecture
-from goli.ipu.ipu_utils import import_poptorch
 
 
 def random_packing(num_nodes, batch_size):
@@ -34,126 +32,6 @@ def global_batch_collator(batch_size, batches):
     return global_batch
 
 
-class test_Packing(ut.TestCase):
-    def test_smart_packing(self):
-        np.random.seed(42)
-
-        batch_sizes = [2, 4, 8, 16, 32, 64]
-        ipu_batch_sizes = [2, 3, 4, 8, 16, 32, 64]
-
-        for batch_size in batch_sizes:
-            for ipu_batch_size in ipu_batch_sizes:
-                err_msg = f"bz={batch_size}, ipu_bz={ipu_batch_size}"
-
-                # Generate random batch size
-                global_batch = batch_size * ipu_batch_size
-                num_nodes = np.abs(np.random.gamma(2, 20, size=global_batch)).astype(int)
-
-                # Use the smart packing
-                packed_indices = smart_packing(num_nodes=num_nodes, batch_size=batch_size)
-                pack_num_nodes = get_pack_sizes(packed_indices, num_nodes)
-
-                # Use the random packing
-                rand_packed_indices = random_packing(num_nodes=num_nodes, batch_size=batch_size)
-                rand_pack_num_nodes = get_pack_sizes(rand_packed_indices, num_nodes)
-
-                # Assert that the smart packing is better than the random packing
-                self.assertLessEqual(max(pack_num_nodes), max(rand_pack_num_nodes), msg=err_msg)
-                self.assertGreaterEqual(min(pack_num_nodes), min(rand_pack_num_nodes), msg=err_msg)
-
-                # Assert that the total number of atoms is right
-                self.assertEqual(sum(pack_num_nodes), sum(num_nodes), msg=err_msg)
-                self.assertEqual(sum(rand_pack_num_nodes), sum(num_nodes), msg=err_msg)
-
-                # Assert that all index are there
-                self.assertListEqual(
-                    np.sort(np.asarray(packed_indices).flatten()).tolist(), np.arange(len(num_nodes)).tolist()
-                )
-                self.assertListEqual(
-                    np.sort(np.asarray(rand_packed_indices).flatten()).tolist(),
-                    np.arange(len(num_nodes)).tolist(),
-                )
-
-    def test_fast_packing(self):
-        np.random.seed(42)
-
-        # Start at 4 for fast_packing for better statistical significance
-        batch_sizes = [4, 8, 16, 32, 64]
-        ipu_batch_sizes = [4, 8, 16, 32, 64]
-
-        for batch_size in batch_sizes:
-            for ipu_batch_size in ipu_batch_sizes:
-                err_msg = f"bz={batch_size}, ipu_bz={ipu_batch_size}"
-
-                # Generate random batch size
-                global_batch = batch_size * ipu_batch_size
-                num_nodes = np.abs(np.random.gamma(2, 20, size=global_batch)).astype(int)
-
-                # Use the smart packing
-                packed_indices = fast_packing(num_nodes=num_nodes, batch_size=batch_size)
-                pack_num_nodes = get_pack_sizes(packed_indices, num_nodes)
-
-                # Use the random packing
-                rand_packed_indices = random_packing(num_nodes=num_nodes, batch_size=batch_size)
-                rand_pack_num_nodes = get_pack_sizes(rand_packed_indices, num_nodes)
-
-                # Assert that the smart packing is better than the random packing
-                self.assertLessEqual(max(pack_num_nodes), max(rand_pack_num_nodes), msg=err_msg)
-                self.assertGreaterEqual(min(pack_num_nodes), min(rand_pack_num_nodes), msg=err_msg)
-
-                # Assert that the total number of atoms is right
-                self.assertEqual(sum(pack_num_nodes), sum(num_nodes), msg=err_msg)
-                self.assertEqual(sum(rand_pack_num_nodes), sum(num_nodes), msg=err_msg)
-
-                # Assert that all index are there
-                self.assertListEqual(
-                    np.sort(np.asarray(packed_indices).flatten()).tolist(), np.arange(len(num_nodes)).tolist()
-                )
-                self.assertListEqual(
-                    np.sort(np.asarray(rand_packed_indices).flatten()).tolist(),
-                    np.arange(len(num_nodes)).tolist(),
-                )
-
-    def test_hybrid_packing(self):
-        np.random.seed(42)
-
-        batch_sizes = [2, 4, 8, 16, 32, 64]
-        ipu_batch_sizes = [2, 3, 4, 8, 16, 32, 64]
-
-        for batch_size in batch_sizes:
-            for ipu_batch_size in ipu_batch_sizes:
-                err_msg = f"bz={batch_size}, ipu_bz={ipu_batch_size}"
-
-                # Generate random batch size
-                global_batch = batch_size * ipu_batch_size
-                num_nodes = np.abs(np.random.gamma(2, 20, size=global_batch)).astype(int)
-
-                # Use the smart packing
-                packed_indices = hybrid_packing(num_nodes=num_nodes, batch_size=batch_size)
-                pack_num_nodes = get_pack_sizes(packed_indices, num_nodes)
-
-                # Use the random packing
-                rand_packed_indices = random_packing(num_nodes=num_nodes, batch_size=batch_size)
-                rand_pack_num_nodes = get_pack_sizes(rand_packed_indices, num_nodes)
-
-                # Assert that the smart packing is better than the random packing
-                self.assertLessEqual(max(pack_num_nodes), max(rand_pack_num_nodes), msg=err_msg)
-                self.assertGreaterEqual(min(pack_num_nodes), min(rand_pack_num_nodes), msg=err_msg)
-
-                # Assert that the total number of atoms is right
-                self.assertEqual(sum(pack_num_nodes), sum(num_nodes), msg=err_msg)
-                self.assertEqual(sum(rand_pack_num_nodes), sum(num_nodes), msg=err_msg)
-
-                # Assert that all index are there
-                self.assertListEqual(
-                    np.sort(np.asarray(packed_indices).flatten()).tolist(), np.arange(len(num_nodes)).tolist()
-                )
-                self.assertListEqual(
-                    np.sort(np.asarray(rand_packed_indices).flatten()).tolist(),
-                    np.arange(len(num_nodes)).tolist(),
-                )
-
-
 class test_DataLoading(ut.TestCase):
     class TestSimpleLightning(LightningModule):
         # Create a basic Ligthning for testing the batch sizes
diff --git a/tests/test_packing.py b/tests/test_packing.py
new file mode 100644
index 000000000..9057b19b7
--- /dev/null
+++ b/tests/test_packing.py
@@ -0,0 +1,220 @@
+# General imports
+import unittest as ut
+import numpy as np
+
+import torch
+from torch_geometric.data import Data, Batch
+
+# Current library imports
+from goli.utils.packing import (
+    smart_packing,
+    get_pack_sizes,
+    fast_packing,
+    hybrid_packing,
+    node_to_pack_indices_mask,
+)
+
+
+def random_packing(num_nodes, batch_size):
+    ipu_batch_size = int(len(num_nodes) / batch_size)
+    indices = np.arange(len(num_nodes))
+    np.random.shuffle(indices)
+    indices = np.reshape(indices, (ipu_batch_size, batch_size)).tolist()
+    return indices
+
+
+class test_Packing(ut.TestCase):
+    def test_smart_packing(self):
+        np.random.seed(42)
+
+        batch_sizes = [2, 4, 8, 16, 32, 64]
+        ipu_batch_sizes = [2, 3, 4, 8, 16, 32, 64]
+
+        for batch_size in batch_sizes:
+            for ipu_batch_size in ipu_batch_sizes:
+                err_msg = f"bz={batch_size}, ipu_bz={ipu_batch_size}"
+
+                # Generate random batch size
+                global_batch = batch_size * ipu_batch_size
+                num_nodes = np.abs(np.random.gamma(2, 20, size=global_batch)).astype(int)
+
+                # Use the smart packing
+                packed_indices = smart_packing(num_nodes=num_nodes, batch_size=batch_size)
+                pack_num_nodes = get_pack_sizes(packed_indices, num_nodes)
+
+                # Use the random packing
+                rand_packed_indices = random_packing(num_nodes=num_nodes, batch_size=batch_size)
+                rand_pack_num_nodes = get_pack_sizes(rand_packed_indices, num_nodes)
+
+                # Assert that the smart packing is better than the random packing
+                self.assertLessEqual(max(pack_num_nodes), max(rand_pack_num_nodes), msg=err_msg)
+                self.assertGreaterEqual(min(pack_num_nodes), min(rand_pack_num_nodes), msg=err_msg)
+
+                # Assert that the total number of atoms is right
+                self.assertEqual(sum(pack_num_nodes), sum(num_nodes), msg=err_msg)
+                self.assertEqual(sum(rand_pack_num_nodes), sum(num_nodes), msg=err_msg)
+
+                # Assert that all index are there
+                self.assertListEqual(
+                    np.sort(np.asarray(packed_indices).flatten()).tolist(), np.arange(len(num_nodes)).tolist()
+                )
+                self.assertListEqual(
+                    np.sort(np.asarray(rand_packed_indices).flatten()).tolist(),
+                    np.arange(len(num_nodes)).tolist(),
+                )
+
+    def test_fast_packing(self):
+        np.random.seed(42)
+
+        # Start at 4 for fast_packing for better statistical significance
+        batch_sizes = [4, 8, 16, 32, 64]
+        ipu_batch_sizes = [4, 8, 16, 32, 64]
+
+        for batch_size in batch_sizes:
+            for ipu_batch_size in ipu_batch_sizes:
+                err_msg = f"bz={batch_size}, ipu_bz={ipu_batch_size}"
+
+                # Generate random batch size
+                global_batch = batch_size * ipu_batch_size
+                num_nodes = np.abs(np.random.gamma(2, 20, size=global_batch)).astype(int)
+
+                # Use the smart packing
+                packed_indices = fast_packing(num_nodes=num_nodes, batch_size=batch_size)
+                pack_num_nodes = get_pack_sizes(packed_indices, num_nodes)
+
+                # Use the random packing
+                rand_packed_indices = random_packing(num_nodes=num_nodes, batch_size=batch_size)
+                rand_pack_num_nodes = get_pack_sizes(rand_packed_indices, num_nodes)
+
+                # Assert that the smart packing is better than the random packing
+                self.assertLessEqual(max(pack_num_nodes), max(rand_pack_num_nodes), msg=err_msg)
+                self.assertGreaterEqual(min(pack_num_nodes), min(rand_pack_num_nodes), msg=err_msg)
+
+                # Assert that the total number of atoms is right
+                self.assertEqual(sum(pack_num_nodes), sum(num_nodes), msg=err_msg)
+                self.assertEqual(sum(rand_pack_num_nodes), sum(num_nodes), msg=err_msg)
+
+                # Assert that all index are there
+                self.assertListEqual(
+                    np.sort(np.asarray(packed_indices).flatten()).tolist(), np.arange(len(num_nodes)).tolist()
+                )
+                self.assertListEqual(
+                    np.sort(np.asarray(rand_packed_indices).flatten()).tolist(),
+                    np.arange(len(num_nodes)).tolist(),
+                )
+
+    def test_hybrid_packing(self):
+        np.random.seed(42)
+
+        batch_sizes = [2, 4, 8, 16, 32, 64]
+        ipu_batch_sizes = [2, 3, 4, 8, 16, 32, 64]
+
+        for batch_size in batch_sizes:
+            for ipu_batch_size in ipu_batch_sizes:
+                err_msg = f"bz={batch_size}, ipu_bz={ipu_batch_size}"
+
+                # Generate random batch size
+                global_batch = batch_size * ipu_batch_size
+                num_nodes = np.abs(np.random.gamma(2, 20, size=global_batch)).astype(int)
+
+                # Use the smart packing
+                packed_indices = hybrid_packing(num_nodes=num_nodes, batch_size=batch_size)
+                pack_num_nodes = get_pack_sizes(packed_indices, num_nodes)
+
+                # Use the random packing
+                rand_packed_indices = random_packing(num_nodes=num_nodes, batch_size=batch_size)
+                rand_pack_num_nodes = get_pack_sizes(rand_packed_indices, num_nodes)
+
+                # Assert that the smart packing is better than the random packing
+                self.assertLessEqual(max(pack_num_nodes), max(rand_pack_num_nodes), msg=err_msg)
+                self.assertGreaterEqual(min(pack_num_nodes), min(rand_pack_num_nodes), msg=err_msg)
+
+                # Assert that the total number of atoms is right
+                self.assertEqual(sum(pack_num_nodes), sum(num_nodes), msg=err_msg)
+                self.assertEqual(sum(rand_pack_num_nodes), sum(num_nodes), msg=err_msg)
+
+                # Assert that all index are there
+                self.assertListEqual(
+                    np.sort(np.asarray(packed_indices).flatten()).tolist(), np.arange(len(num_nodes)).tolist()
+                )
+                self.assertListEqual(
+                    np.sort(np.asarray(rand_packed_indices).flatten()).tolist(),
+                    np.arange(len(num_nodes)).tolist(),
+                )
+
+    def test_node_to_pack_indices_mask(self):
+        # Create a dummy batch
+        in_dim = 7
+        in_dim_edges = 11
+        max_num_nodes_per_graph = 20
+        batch_size_per_pack = 5
+
+        torch.manual_seed(42)
+
+        # Create a dummy batch of graphs
+        batch, all_num_nodes = [], []
+        for ii in range(100):
+            num_nodes = torch.randint(1, max_num_nodes_per_graph, (1,)).item()
+            all_num_nodes.append(num_nodes)
+            num_edges = abs(round(2.2 * num_nodes) + torch.randint(-2, 2, (1,)).item()) + 1
+            x = torch.randn(num_nodes, in_dim, dtype=torch.float32)
+            edge_idx = torch.randint(0, num_nodes, (2, num_edges))
+            e = torch.randn(edge_idx.shape[-1], in_dim_edges, dtype=torch.float32)
+            g = Data(h=x, edge_index=edge_idx, edge_attr=e)
+            batch.append(g)
+        batch = Batch.from_data_list(batch)
+
+        # Get the packing
+        packed_graph_idx = fast_packing(all_num_nodes, batch_size_per_pack)
+        pack_sizes = get_pack_sizes(packed_graph_idx, all_num_nodes)
+        max_pack_size = max(pack_sizes)
+        num_packs = len(pack_sizes)
+
+        # Get the node to pack indices and the mask
+        pack_from_node_idx, pack_attn_mask = node_to_pack_indices_mask(packed_graph_idx, all_num_nodes)
+
+        # Assert that the nodes to pack indices are correct
+        h = torch.arange(batch.num_nodes, dtype=torch.float32)
+        packed_shape = [num_packs, max_pack_size]
+        h_packed = torch.zeros(packed_shape)
+        h_packed[pack_from_node_idx[:, 0], pack_from_node_idx[:, 1]] = h
+        h_packed_unique = torch.sort(torch.unique(h_packed))[0]
+        np.testing.assert_array_equal(h_packed_unique, torch.arange(batch.num_nodes))
+        self.assertEqual(h_packed.sum(), h.sum())
+
+        # Test again with additional h dimension
+        h = batch.h
+        packed_shape = [num_packs, max_pack_size] + list(h.shape[1:])
+        h_packed = torch.zeros(packed_shape)
+        h_packed[pack_from_node_idx[:, 0], pack_from_node_idx[:, 1]] = h
+        h_packed_unique = torch.sort(torch.unique(h_packed))[0]
+        h_packed_unique = h_packed_unique[h_packed_unique != 0]
+        np.testing.assert_array_almost_equal(h_packed_unique, torch.unique(h))
+        self.assertAlmostEqual(h_packed.sum().item(), h.sum().item(), places=3)
+
+        # Assert that the mask is correct by counting the number of False values (the sum of squared number of nodes per pack)
+        num_false = (~pack_attn_mask).sum([1, 2])
+        num_expected = torch.as_tensor(
+            [sum([all_num_nodes[graph_idx] ** 2 for graph_idx in pack]) for pack in packed_graph_idx]
+        )
+        np.testing.assert_array_equal(num_false, num_expected)
+
+        # Assert that the mask is correct by counting the number of elements in each row and column
+        num_expected = []
+        for pack in packed_graph_idx:
+            pack_num_expected = []
+            for graph_idx in pack:
+                num_nodes = all_num_nodes[graph_idx]
+                for ii in range(num_nodes):
+                    pack_num_expected.append(num_nodes)
+            pack_num_expected.extend([0] * (max_pack_size - len(pack_num_expected)))
+            num_expected.append(pack_num_expected)
+        num_expected = torch.as_tensor(num_expected)
+        num_false_row = (~pack_attn_mask).sum([2])
+        num_false_col = (~pack_attn_mask).sum([1])
+        np.testing.assert_array_equal(num_false_row, num_expected)
+        np.testing.assert_array_equal(num_false_col, num_expected)
+
+
+if __name__ == "__main__":
+    ut.main()