diff --git a/n3fit/runcards/reproduce_nnpdf40/NNPDF40_nnlo_as_01180_1000.yml b/n3fit/runcards/reproduce_nnpdf40/NNPDF40_nnlo_as_01180_1000.yml
index da30fc1b2e..4e2295470f 100644
--- a/n3fit/runcards/reproduce_nnpdf40/NNPDF40_nnlo_as_01180_1000.yml
+++ b/n3fit/runcards/reproduce_nnpdf40/NNPDF40_nnlo_as_01180_1000.yml
@@ -163,4 +163,4 @@ integrability:
 
 ############################################################
 debug: false
-maxcores: 4
+maxcores: 4
\ No newline at end of file
diff --git a/n3fit/src/n3fit/backends/keras_backend/operations.py b/n3fit/src/n3fit/backends/keras_backend/operations.py
index efff803b92..c44faebc71 100644
--- a/n3fit/src/n3fit/backends/keras_backend/operations.py
+++ b/n3fit/src/n3fit/backends/keras_backend/operations.py
@@ -215,6 +215,10 @@ def flatten(x):
     """Flatten tensor x"""
     return tf.reshape(x, (-1,))
 
+@tf.function
+def reshape(x, shape):
+    """ reshape tensor x """
+    return tf.reshape(x, shape)
 
 def boolean_mask(*args, **kwargs):
     """
diff --git a/n3fit/src/n3fit/checks.py b/n3fit/src/n3fit/checks.py
index dc9c6a9d22..3363839b48 100644
--- a/n3fit/src/n3fit/checks.py
+++ b/n3fit/src/n3fit/checks.py
@@ -362,11 +362,11 @@ def check_consistent_parallel(parameters, parallel_models, same_trvl_per_replica
     """
     if not parallel_models:
         return
-    if not same_trvl_per_replica:
-        raise CheckError(
-            "Replicas cannot be run in parallel with different training/validation "
-            " masks, please set `same_trvl_per_replica` to True in the runcard"
-        )
+#    if not same_trvl_per_replica:
+#        raise CheckError(
+#            "Replicas cannot be run in parallel with different training/validation "
+#            " masks, please set `same_trvl_per_replica` to True in the runcard"
+#        )
     if parameters.get("layer_type") != "dense":
         raise CheckError("Parallelization has only been tested with layer_type=='dense'")
 
diff --git a/n3fit/src/n3fit/layers/losses.py b/n3fit/src/n3fit/layers/losses.py
index 4315919907..cf8befe4f3 100644
--- a/n3fit/src/n3fit/layers/losses.py
+++ b/n3fit/src/n3fit/layers/losses.py
@@ -57,7 +57,7 @@ def build(self, input_shape):
         weights of the layers"""
         init = MetaLayer.init_constant(self._invcovmat)
         self.kernel = self.builder_helper(
-            "invcovmat", (self._ndata, self._ndata), init, trainable=False
+            "invcovmat", self._invcovmat.shape, init, trainable=False
         )
         mask_shape = (1, 1, self._ndata)
         if self._mask is None:
@@ -85,10 +85,17 @@ def call(self, y_pred, **kwargs):
         tmp = op.op_multiply([tmp_raw, self.mask])
         if tmp.shape[1] == 1:
             # einsum is not well suited for CPU, so use tensordot if not multimodel
-            right_dot = op.tensor_product(self.kernel, tmp[0, 0, :], axes=1)
-            res = op.tensor_product(tmp[0, :, :], right_dot, axes=1)
+            if len(self.kernel.shape) == 3:
+                right_dot = op.tensor_product(self.kernel[0, ...], tmp[0, 0, :], axes=1)
+                res = op.tensor_product(tmp[0, :, :], right_dot, axes=1)
+            else:
+                right_dot = op.tensor_product(self.kernel, tmp[0, 0, :], axes=1)
+                res = op.tensor_product(tmp[0, :, :], right_dot, axes=1)
         else:
-            res = op.einsum("bri, ij, brj -> r", tmp, self.kernel, tmp)
+            if len(self.kernel.shape) == 3:
+                res = op.einsum("bri, rij, brj -> r", tmp, self.kernel, tmp)
+            else:
+                res = op.einsum("bri, ij, brj -> r", tmp, self.kernel, tmp)
         return res
 
 
diff --git a/n3fit/src/n3fit/layers/mask.py b/n3fit/src/n3fit/layers/mask.py
index 877df3b502..cc2ad4e552 100644
--- a/n3fit/src/n3fit/layers/mask.py
+++ b/n3fit/src/n3fit/layers/mask.py
@@ -1,6 +1,6 @@
 from n3fit.backends import MetaLayer
 from n3fit.backends import operations as op
-
+from numpy import count_nonzero
 
 class Mask(MetaLayer):
     """
@@ -19,14 +19,20 @@ class Mask(MetaLayer):
         c: float
             constant multiplier for every output
         axis: int
-            axis in which to apply the mask
+            axis in which to apply the mask. Currently,
+            only the last axis gives the correct output shape
     """
 
     def __init__(self, bool_mask=None, c=None, axis=None, **kwargs):
         if bool_mask is None:
             self.mask = None
+            self.last_dim = -1
         else:
             self.mask = op.numpy_to_tensor(bool_mask, dtype=bool)
+            if len(bool_mask.shape) == 1:
+                self.last_dim = count_nonzero(bool_mask)
+            else:
+                self.last_dim = count_nonzero(bool_mask[0, ...])
         self.c = c
         self.axis = axis
         super().__init__(**kwargs)
@@ -41,7 +47,10 @@ def build(self, input_shape):
 
     def call(self, ret):
         if self.mask is not None:
-            ret = op.boolean_mask(ret, self.mask, axis=self.axis)
+            flat_res = op.boolean_mask(ret, self.mask, axis=self.axis)
+            output_shape = ret.get_shape().as_list()
+            output_shape[-1] = self.last_dim
+            ret = op.reshape(flat_res, shape=output_shape)
         if self.c is not None:
             ret = ret * self.kernel
         return ret
diff --git a/n3fit/src/n3fit/model_gen.py b/n3fit/src/n3fit/model_gen.py
index 2bd898021b..c665006311 100644
--- a/n3fit/src/n3fit/model_gen.py
+++ b/n3fit/src/n3fit/model_gen.py
@@ -25,13 +25,13 @@
     FlavourToEvolution,
     ObsRotation,
     Preprocessing,
+    Mask,
     losses,
 )
 from n3fit.layers.observable import is_unique
 from n3fit.msr import generate_msr_model_and_grid
 from validphys.photon.compute import Photon  # only used for type hint here
 
-
 @dataclass
 class ObservableWrapper:
     """Wraps many observables into an experimental layer once the PDF model is prepared
@@ -47,6 +47,7 @@ class ObservableWrapper:
 
     name: str
     observables: list
+    trvl_mask_layers: list
     dataset_xsizes: list
     invcovmat: np.array = None
     covmat: np.array = None
@@ -61,8 +62,7 @@ def _generate_loss(self, mask=None):
         was initialized with"""
         if self.invcovmat is not None:
             loss = losses.LossInvcovmat(
-                self.invcovmat, self.data, mask, covmat=self.covmat, name=self.name
-            )
+                self.invcovmat, self.data, mask, covmat=self.covmat, name=self.name)
         elif self.positivity:
             loss = losses.LossPositivity(name=self.name, c=self.multiplier)
         elif self.integrability:
@@ -86,8 +86,15 @@ def _generate_experimental_layer(self, pdf):
         else:
             output_layers = [obs(pdf) for obs in self.observables]
 
-        # Finally concatenate all observables (so that experiments are one single entitiy)
-        ret = op.concatenate(output_layers, axis=2)
+        masked_output_layers = []
+        if self.trvl_mask_layers is not None:
+            for output_layer, mask_layer in zip(output_layers, self.trvl_mask_layers):
+                masked_output_layers.append(mask_layer(output_layer))
+        else:
+            masked_output_layers = output_layers
+
+        # Finally concatenate all observables (so that experiments are one single entity)
+        ret = op.concatenate(masked_output_layers)
         if self.rotation is not None:
             ret = self.rotation(ret)
         return ret
@@ -98,8 +105,14 @@ def __call__(self, pdf_layer, mask=None):
         return loss_f(experiment_prediction)
 
 
-def observable_generator(
-    spec_dict, positivity_initial=1.0, integrability=False
+def observable_generator(spec_dict,
+                         mask_array=None,
+                         training_data=None,
+                         validation_data=None,
+                         invcovmat_tr=None,
+                         invcovmat_vl=None,
+                         positivity_initial=1.0,
+                         integrability=False
 ):  # pylint: disable=too-many-locals
     """
     This function generates the observable models for each experiment.
@@ -148,10 +161,13 @@ def observable_generator(
     spec_name = spec_dict["name"]
     dataset_xsizes = []
     model_inputs = []
-    model_obs_tr = []
-    model_obs_vl = []
-    model_obs_ex = []
+    model_observables = []
+    tr_mask_layers = []
+    vl_mask_layers = []
+    offset = 0
+    apply_masks = spec_dict.get("data_transformation_tr") is None and mask_array is not None
     # The first step is to compute the observable for each of the datasets
+    masks = []
     for dataset in spec_dict["datasets"]:
         # Get the generic information of the dataset
         dataset_name = dataset.name
@@ -165,56 +181,29 @@ def observable_generator(
         # Set the operation (if any) to be applied to the fktables of this dataset
         operation_name = dataset.operation
 
+        # Extract the masks that will end up in the observable wrappers...
+        if apply_masks:
+            trmask = mask_array[:, offset:offset + dataset.ndata]
+            masks.append(trmask)
+            tr_mask_layers.append(Mask(trmask, axis=1, name=f"trmask_{dataset_name}"))
+            vl_mask_layers.append(Mask(~trmask, axis=1, name=f"vlmask_{dataset_name}"))
+
         # Now generate the observable layer, which takes the following information:
         # operation name
         # dataset name
         # list of validphys.coredata.FKTableData objects
         #   these will then be used to check how many different pdf inputs are needed
         #   (and convolutions if given the case)
-
-        if spec_dict["positivity"]:
-            # Positivity (and integrability, which is a special kind of positivity...)
-            # enters only at the "training" part of the models
-            obs_layer_tr = Obs_Layer(
-                dataset.fktables_data,
-                dataset.training_fktables(),
-                operation_name,
-                name=f"dat_{dataset_name}",
-            )
-            obs_layer_ex = obs_layer_vl = None
-        elif spec_dict.get("data_transformation_tr") is not None:
-            # Data transformation needs access to the full array of output data
-            obs_layer_ex = Obs_Layer(
-                dataset.fktables_data,
-                dataset.fktables(),
-                operation_name,
-                name=f"exp_{dataset_name}",
-            )
-            obs_layer_tr = obs_layer_vl = obs_layer_ex
-        else:
-            obs_layer_tr = Obs_Layer(
-                dataset.fktables_data,
-                dataset.training_fktables(),
-                operation_name,
-                name=f"dat_{dataset_name}",
-            )
-            obs_layer_ex = Obs_Layer(
-                dataset.fktables_data,
-                dataset.fktables(),
-                operation_name,
-                name=f"exp_{dataset_name}",
-            )
-            obs_layer_vl = Obs_Layer(
-                dataset.fktables_data,
-                dataset.validation_fktables(),
-                operation_name,
-                name=f"val_{dataset_name}",
-            )
+        obs_layer = Obs_Layer(
+            dataset.fktables_data,
+            dataset.fktables(),
+            operation_name,
+            name=f"dat_{dataset_name}")
 
         # If the observable layer found that all input grids are equal, the splitting will be None
         # otherwise the different xgrids need to be stored separately
         # Note: for pineappl grids, obs_layer_tr.splitting should always be None
-        if obs_layer_tr.splitting is None:
+        if obs_layer.splitting is None:
             xgrid = dataset.fktables_data[0].xgrid
             model_inputs.append(xgrid)
             dataset_xsizes.append(len(xgrid))
@@ -223,9 +212,10 @@ def observable_generator(
             model_inputs += xgrids
             dataset_xsizes.append(sum([len(i) for i in xgrids]))
 
-        model_obs_tr.append(obs_layer_tr)
-        model_obs_vl.append(obs_layer_vl)
-        model_obs_ex.append(obs_layer_ex)
+        model_observables.append(obs_layer)
+
+        # shift offset for new mask array
+        offset = offset + dataset.ndata
 
     # Check whether all xgrids of all observables in this experiment are equal
     # if so, simplify the model input
@@ -240,7 +230,8 @@ def observable_generator(
     if spec_dict["positivity"]:
         out_positivity = ObservableWrapper(
             spec_name,
-            model_obs_tr,
+            model_observables,
+            tr_mask_layers if apply_masks else None,
             dataset_xsizes,
             multiplier=positivity_initial,
             positivity=not integrability,
@@ -265,23 +256,26 @@ def observable_generator(
 
     out_tr = ObservableWrapper(
         spec_name,
-        model_obs_tr,
+        model_observables,
+        tr_mask_layers if apply_masks else None,
         dataset_xsizes,
-        invcovmat=spec_dict["invcovmat"],
-        data=spec_dict["expdata"],
+        invcovmat=invcovmat_tr,
+        data=training_data,
         rotation=obsrot_tr,
     )
     out_vl = ObservableWrapper(
         f"{spec_name}_val",
-        model_obs_vl,
+        model_observables,
+        vl_mask_layers if apply_masks else None,
         dataset_xsizes,
-        invcovmat=spec_dict["invcovmat_vl"],
-        data=spec_dict["expdata_vl"],
+        invcovmat=invcovmat_vl,
+        data=validation_data,
         rotation=obsrot_vl,
     )
     out_exp = ObservableWrapper(
         f"{spec_name}_exp",
-        model_obs_ex,
+        model_observables,
+        None,
         dataset_xsizes,
         invcovmat=spec_dict["invcovmat_true"],
         covmat=spec_dict["covmat"],
diff --git a/n3fit/src/n3fit/model_trainer.py b/n3fit/src/n3fit/model_trainer.py
index 2f77590a6e..c1402756c3 100644
--- a/n3fit/src/n3fit/model_trainer.py
+++ b/n3fit/src/n3fit/model_trainer.py
@@ -146,15 +146,13 @@ def __init__(
                 list with the replicas ids to be fitted
         """
         # Save all input information
-        self.exp_info = exp_info
-        if pos_info is None:
-            pos_info = []
+        self.exp_info = list(exp_info)
         self.pos_info = pos_info
         self.integ_info = integ_info
         if self.integ_info is not None:
-            self.all_info = exp_info + pos_info + integ_info
+            self.all_info = self.exp_info[0] + pos_info + integ_info
         else:
-            self.all_info = exp_info + pos_info
+            self.all_info = self.exp_info[0] + pos_info
         self.flavinfo = flavinfo
         self.fitbasis = fitbasis
         self._nn_seeds = nnseeds
@@ -229,6 +227,7 @@ def __init__(
             "model": None,
             "folds": [],
         }
+        self.tr_masks = []
 
         self._fill_the_dictionaries()
 
@@ -278,7 +277,7 @@ def _fill_the_dictionaries(self):
             - ``name``: names of the experiment
             - ``ndata``: number of experimental points
         """
-        for exp_dict in self.exp_info:
+        for index, exp_dict in enumerate(self.exp_info[0]):
             self.training["expdata"].append(exp_dict["expdata"])
             self.validation["expdata"].append(exp_dict["expdata_vl"])
             self.experimental["expdata"].append(exp_dict["expdata_true"])
@@ -502,7 +501,7 @@ def _reset_observables(self):
             self.experimental[key] = []
 
     ############################################################################
-    # # Parametizable functions                                                #
+    # # Parameterizable functions                                                #
     #                                                                          #
     # The functions defined in this block accept a 'params' dictionary which   #
     # defines the fit and the behaviours of the Neural Networks                #
@@ -543,11 +542,23 @@ def _generate_observables(
         log.info("Generating layers")
 
         # Now we need to loop over all dictionaries (First exp_info, then pos_info and integ_info)
-        for exp_dict in self.exp_info:
+        for index, exp_dict in enumerate(self.exp_info[0]):
             if not self.mode_hyperopt:
                 log.info("Generating layers for experiment %s", exp_dict["name"])
 
-            exp_layer = model_gen.observable_generator(exp_dict)
+            # Stacked tr-vl mask array for all replicas for this dataset
+            replica_masks = np.stack([e[index]["trmask"] for e in self.exp_info])
+            training_data = np.stack([e[index]["expdata"].flatten() for e in self.exp_info])
+            validation_data = np.stack([e[index]["expdata_vl"].flatten() for e in self.exp_info])
+            invcovmat = np.stack([e[index]["invcovmat"] for e in self.exp_info])
+            invcovmat_vl = np.stack([e[index]["invcovmat_vl"] for e in self.exp_info])
+
+            exp_layer = model_gen.observable_generator(exp_dict,
+                                                       mask_array=replica_masks,
+                                                       training_data=training_data,
+                                                       validation_data=validation_data,
+                                                       invcovmat_tr=invcovmat,
+                                                       invcovmat_vl=invcovmat_vl)
 
             # Save the input(s) corresponding to this experiment
             self.input_list.append(exp_layer["inputs"])
@@ -568,8 +579,14 @@ def _generate_observables(
             pos_initial, pos_multiplier = _LM_initial_and_multiplier(
                 all_pos_initial, all_pos_multiplier, max_lambda, positivity_steps
             )
-
-            pos_layer = model_gen.observable_generator(pos_dict, positivity_initial=pos_initial)
+            replica_masks = np.stack([pos_dict["trmask"] for i in range(len(self.exp_info))])
+            training_data = np.stack([pos_dict["expdata"].flatten() for i in range(len(self.exp_info))])
+
+            pos_layer = model_gen.observable_generator(pos_dict,
+                                                       positivity_initial=pos_initial,
+                                                       mask_array=replica_masks,
+                                                       training_data=training_data,
+                                                       validation_data=training_data)
             # The input list is still common
             self.input_list.append(pos_layer["inputs"])
 
diff --git a/n3fit/src/n3fit/performfit.py b/n3fit/src/n3fit/performfit.py
index 431bb1f5a6..b38e235ada 100644
--- a/n3fit/src/n3fit/performfit.py
+++ b/n3fit/src/n3fit/performfit.py
@@ -159,31 +159,33 @@ def performfit(
         replicas, replica_experiments, nnseeds = zip(*replicas_nnseed_fitting_data_dict)
         # Parse the experiments so that the output data contain information for all replicas
         # as the only different from replica to replica is the experimental training/validation data
-        all_experiments = copy.deepcopy(replica_experiments[0])
-        for i_exp in range(len(all_experiments)):
-            training_data = []
-            validation_data = []
-            for i_rep in range(n_models):
-                training_data.append(replica_experiments[i_rep][i_exp]['expdata'])
-                validation_data.append(replica_experiments[i_rep][i_exp]['expdata_vl'])
-            all_experiments[i_exp]['expdata'] = np.concatenate(training_data, axis=0)
-            all_experiments[i_exp]['expdata_vl'] = np.concatenate(validation_data, axis=0)
+        # all_experiments = copy.deepcopy(replica_experiments[0])
+        # n_experiments dicts
+        # for i_exp in range(len(all_experiments)):
+        #    training_data = []
+        #    validation_data = []
+        #    for i_rep in range(n_models):
+        #        training_data.append(replica_experiments[i_rep][i_exp]['expdata'])
+        #        validation_data.append(replica_experiments[i_rep][i_exp]['expdata_vl'])
+        #    all_experiments[i_exp]['expdata'] = np.concatenate(training_data, axis=0)
+        #    all_experiments[i_exp]['expdata_vl'] = np.concatenate(validation_data, axis=0)
         log.info(
             "Starting parallel fits from replica %d to %d",
             replicas[0],
             replicas[0] + n_models - 1,
         )
-        replicas_info = [(replicas, all_experiments, nnseeds)]
+        replicas_info = [(replicas, replica_experiments, nnseeds)]
     else:
+        # Cases 1 and 2 above are a special case of 3 where the replica idx and the seed should
+        # be a list of just one element
         replicas_info = replicas_nnseed_fitting_data_dict
+        for i, info_tuple in enumerate(replicas_info):
+            replica_idxs = info_tuple[0]
+            nnseeds = info_tuple[2]
+            replicas_info[i] = (tuple([replica_idxs]), [info_tuple[1]], tuple([nnseeds]))
 
     for replica_idxs, exp_info, nnseeds in replicas_info:
-        if not parallel_models or n_models == 1:
-            # Cases 1 and 2 above are a special case of 3 where the replica idx and the seed should
-            # be a list of just one element
-            replica_idxs = [replica_idxs]
-            nnseeds = [nnseeds]
-            log.info("Starting replica fit %d", replica_idxs[0])
+        log.info("Starting replica fit " + str(replica_idxs))
 
         # Generate a ModelTrainer object
         # this object holds all necessary information to train a PDF (up to the NN definition)
diff --git a/validphys2/src/validphys/commondata.py b/validphys2/src/validphys/commondata.py
index b67bc3ab4f..1bbd7c0c2f 100644
--- a/validphys2/src/validphys/commondata.py
+++ b/validphys2/src/validphys/commondata.py
@@ -8,8 +8,9 @@
 """
 from reportengine import collect
 from validphys.commondataparser import load_commondata
+import functools
 
-
+@functools.lru_cache
 def loaded_commondata_with_cuts(commondata, cuts):
     """Load the commondata and apply cuts.
 
diff --git a/validphys2/src/validphys/config.py b/validphys2/src/validphys/config.py
index 6530471051..724ea0dfdc 100644
--- a/validphys2/src/validphys/config.py
+++ b/validphys2/src/validphys/config.py
@@ -45,10 +45,11 @@
 )
 from validphys.paramfits.config import ParamfitsConfig
 from validphys.plotoptions import get_info
+from validphys.utils import freezeargs
 import validphys.scalevariations
 
-log = logging.getLogger(__name__)
 
+log = logging.getLogger(__name__)
 
 class Environment(Environment):
     """Container for information to be filled at run time"""
@@ -561,6 +562,7 @@ def _produce_similarity_cuts(self, commondata):
             inps.append((ds, pdf))
         return SimilarCuts(tuple(inps), cut_similarity_threshold)
 
+    @functools.lru_cache
     def produce_cuts(self, *, commondata, use_cuts):
         """Obtain cuts for a given dataset input, based on the
         appropriate policy.
@@ -1304,7 +1306,9 @@ def parse_default_filter_rules_recorded_spec_(self, spec):
         it reportengine detects a conflict in the `dataset` key.
         """
         return spec
-
+    
+    @freezeargs
+    @functools.lru_cache
     def produce_rules(
         self,
         theoryid,
diff --git a/validphys2/src/validphys/core.py b/validphys2/src/validphys/core.py
index 8d59e44c97..6c65a845e2 100644
--- a/validphys2/src/validphys/core.py
+++ b/validphys2/src/validphys/core.py
@@ -493,8 +493,9 @@ def load_cfactors(self):
 
         return [[parse_cfactor(c.open("rb")) for c in cfacs] for cfacs in self.cfactors]
 
+    @functools.lru_cache()
     def load_with_cuts(self, cuts):
-        """Load the fktable and apply cuts inmediately. Returns a FKTableData"""
+        """Load the fktable and apply cuts immediately. Returns a FKTableData"""
         return load_fktable(self).with_cuts(cuts)
 
 
diff --git a/validphys2/src/validphys/covmats.py b/validphys2/src/validphys/covmats.py
index 910156be7f..7d3a300698 100644
--- a/validphys2/src/validphys/covmats.py
+++ b/validphys2/src/validphys/covmats.py
@@ -2,6 +2,7 @@
 matrices on different levels of abstraction
 """
 import logging
+import functools
 
 import numpy as np
 import pandas as pd
@@ -227,6 +228,7 @@ def dataset_inputs_covmat_from_systematics(
 
 
 @check_cuts_considered
+@functools.lru_cache
 def dataset_t0_predictions(dataset, t0set):
     """Returns the t0 predictions for a ``dataset`` which are the predictions
     calculated using the central member of ``pdf``. Note that if ``pdf`` has
diff --git a/validphys2/src/validphys/filters.py b/validphys2/src/validphys/filters.py
index 91c80ce2fb..e486339b6a 100644
--- a/validphys2/src/validphys/filters.py
+++ b/validphys2/src/validphys/filters.py
@@ -6,6 +6,7 @@
 from importlib.resources import read_text
 import logging
 import re
+import functools
 
 import numpy as np
 
@@ -13,6 +14,7 @@
 from reportengine.compat import yaml
 from validphys.commondatawriter import write_commondata_to_file, write_systype_to_file
 import validphys.cuts
+from validphys.utils import freezeargs
 
 log = logging.getLogger(__name__)
 
@@ -555,6 +557,8 @@ def _make_point_namespace(self, dataset, idat) -> dict:
         return ns
 
 
+@freezeargs
+@functools.lru_cache
 def get_cuts_for_dataset(commondata, rules) -> list:
     """Function to generate a list containing the index
     of all experimental points that passed kinematic
diff --git a/validphys2/src/validphys/n3fit_data.py b/validphys2/src/validphys/n3fit_data.py
index 38f946248d..123c625904 100644
--- a/validphys2/src/validphys/n3fit_data.py
+++ b/validphys2/src/validphys/n3fit_data.py
@@ -74,7 +74,7 @@ def __iter__(self):
             yield m
 
 
-def tr_masks(data, replica_trvlseed):
+def tr_masks(data, replica_trvlseed, parallel_models=False):
     """Generate the boolean masks used to split data into training and
     validation points. Returns a list of 1-D boolean arrays, one for each
     dataset. Each array has length equal to N_data, the datapoints which
@@ -82,6 +82,9 @@ def tr_masks(data, replica_trvlseed):
 
         tr_data = data[tr_mask]
 
+    The single_datapoints_toss flag signals whether single-point datasets
+    should be always included in the training set only (True), or randomly
+    selected. The former is required for parallel replica fits.
     """
     nameseed = int(hashlib.sha256(str(data).encode()).hexdigest(), 16) % 10**8
     nameseed += replica_trvlseed
@@ -98,7 +101,7 @@ def tr_masks(data, replica_trvlseed):
         trmax = int(ndata * frac)
         if trmax == 0:
             # If that number is 0, then get 1 point with probability frac
-            trmax = int(rng.random() < frac)
+            trmax = int(rng.random() < frac) if not parallel_models else 1
         mask = np.concatenate([np.ones(trmax, dtype=bool), np.zeros(ndata - trmax, dtype=bool)])
         rng.shuffle(mask)
         trmask_partial.append(mask)
@@ -238,7 +241,7 @@ def fitting_data_dict(
 
     expdata = make_replica
     tr_masks = tr_masks.masks
-    covmat = dataset_inputs_fitting_covmat  # t0 covmat, or theory covmat or whatever was decided by the runcard
+    covmat = dataset_inputs_fitting_covmat # t0 covmat, or theory covmat or whatever was decided by the runcard
     inv_true = np.linalg.inv(covmat)
     fittable_datasets = fittable_datasets_masked
 
@@ -347,7 +350,7 @@ def pseudodata_table(groups_replicas_indexed_make_replica, replicas):
     Notes
     -----
     Whilst running ``n3fit``, this action will only be called if
-    `fitting::savepseudodata` is `true` (as per the default setting) and
+    `fitting::savepseudodata` is `true` (as per the default setting) and 
     replicas are fitted one at a time. The table can be found in the replica
     folder i.e. <fit dir>/nnfit/replica_*/
 
diff --git a/validphys2/src/validphys/n3fit_data_utils.py b/validphys2/src/validphys/n3fit_data_utils.py
index fe908d73cb..db52760099 100644
--- a/validphys2/src/validphys/n3fit_data_utils.py
+++ b/validphys2/src/validphys/n3fit_data_utils.py
@@ -7,6 +7,7 @@
 The ``validphys_group_extractor`` will loop over every dataset of a given group
 loading their fktables (and applying any necessary cuts).
 """
+import functools
 import dataclasses
 from itertools import zip_longest
 
diff --git a/validphys2/src/validphys/utils.py b/validphys2/src/validphys/utils.py
index e78fc866bc..296b63880a 100644
--- a/validphys2/src/validphys/utils.py
+++ b/validphys2/src/validphys/utils.py
@@ -8,10 +8,35 @@
 import pathlib
 import shutil
 import tempfile
+import functools
 
 import numpy as np
 from validobj import ValidationError, parse_input
 
+from frozendict import frozendict
+from frozenlist import FrozenList as frozenlist
+
+def immute(element):
+    if isinstance(element, dict):
+        return frozendict(element)
+    if isinstance(element, list):
+        ret = frozenlist(element)
+        ret.freeze()
+        return ret
+    return element 
+
+def freezeargs(func):
+    """Transform mutable dictionnary
+    Into immutable
+    Useful to be compatible with cache
+    """
+    @functools.wraps(func)
+    def wrapped(*args, **kwargs):
+        args = tuple([immute(arg) for arg in args])
+        kwargs = {k: immute(v) for k, v in kwargs.items()}
+        return func(*args, **kwargs)
+    return wrapped
+
 
 def parse_yaml_inp(inp, spec, path):
     """Helper function to parse yaml using the `validobj` library and print