From d8e8e1c9460da08c456e9779731ee83d2e8b2e9d Mon Sep 17 00:00:00 2001
From: gvdoord <g.vandenOord@gmail.com>
Date: Wed, 22 Jun 2022 22:01:08 +0200
Subject: [PATCH 01/18] Initial work to accommodate different tr-vl split per
 parallel replica.

# Conflicts:
#	n3fit/src/n3fit/model_trainer.py
---
 n3fit/src/n3fit/checks.py        |  10 +--
 n3fit/src/n3fit/model_gen.py     |   2 +-
 n3fit/src/n3fit/model_trainer.py | 112 ++++++++++++++++++-------------
 n3fit/src/n3fit/performfit.py    |  16 +++--
 4 files changed, 79 insertions(+), 61 deletions(-)

diff --git a/n3fit/src/n3fit/checks.py b/n3fit/src/n3fit/checks.py
index e07d27d74c..10caa5183e 100644
--- a/n3fit/src/n3fit/checks.py
+++ b/n3fit/src/n3fit/checks.py
@@ -359,11 +359,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/model_gen.py b/n3fit/src/n3fit/model_gen.py
index 2814c2f656..ffa0aba00a 100644
--- a/n3fit/src/n3fit/model_gen.py
+++ b/n3fit/src/n3fit/model_gen.py
@@ -161,7 +161,7 @@ def observable_generator(
         # 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)
-
+# TODO: Make these more memory-efficient, i.e. factoring FKTables and separate Mask layer...
         if spec_dict["positivity"]:
             # Positivity (and integrability, which is a special kind of positivity...)
             # enters only at the "training" part of the models
diff --git a/n3fit/src/n3fit/model_trainer.py b/n3fit/src/n3fit/model_trainer.py
index b756c0e26b..feee2a28a8 100644
--- a/n3fit/src/n3fit/model_trainer.py
+++ b/n3fit/src/n3fit/model_trainer.py
@@ -133,15 +133,13 @@ def __init__(
                 number of models to fit in parallel
         """
         # 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 + pos_info + integ_info
         else:
-            self.all_info = exp_info + pos_info
+            self.all_info = self.exp_info + pos_info
         self.flavinfo = flavinfo
         self.fitbasis = fitbasis
         self._nn_seeds = nnseeds
@@ -186,9 +184,9 @@ def __init__(
         # Initialize the dictionaries which contain all fitting information
         self.input_list = []
         self.training = {
-            "output": [],
-            "expdata": [],
-            "ndata": 0,
+            "output": [[] for _ in range(parallel_models)],
+            "expdata": [[] for _ in range(parallel_models)],
+            "ndata": [0 for _ in range(parallel_models)],
             "model": None,
             "posdatasets": [],
             "posmultipliers": [],
@@ -196,30 +194,31 @@ def __init__(
             "integdatasets": [],
             "integmultipliers": [],
             "integinitials": [],
-            "folds": [],
+            "folds": [[] for _ in range(parallel_models)],
         }
         self.validation = {
-            "output": [],
-            "expdata": [],
-            "ndata": 0,
+            "output": [[] for _ in range(parallel_models)],
+            "expdata": [[] for _ in range(parallel_models)],
+            "ndata": [0 for _ in range(parallel_models)],
             "model": None,
-            "folds": [],
+            "folds": [[] for _ in range(parallel_models)],
             "posdatasets": [],
         }
         self.experimental = {
-            "output": [],
-            "expdata": [],
-            "ndata": 0,
+            "output": [[] for _ in range(parallel_models)],
+            "expdata": [[] for _ in range(parallel_models)],
+            "ndata": [0 for _ in range(parallel_models)],
             "model": None,
-            "folds": [],
+            "folds": [[] for _ in range(parallel_models)],
         }
 
         self._fill_the_dictionaries()
 
-        if self.validation["ndata"] == 0:
+        if self.validation["ndata"][0] == 0:
             # If there is no validation, the validation chi2 = training chi2
             self.no_validation = True
-            self.validation["expdata"] = self.training["expdata"]
+            for replica in range(parallel_models):
+                self.validation["expdata"][replica] = self.training["expdata"][replica]
         else:
             # Consider the validation only if there is validation (of course)
             self.no_validation = False
@@ -262,34 +261,44 @@ def _fill_the_dictionaries(self):
             - ``name``: names of the experiment
             - ``ndata``: number of experimental points
         """
-        for exp_dict in self.exp_info:
-            self.training["expdata"].append(exp_dict["expdata"])
-            self.validation["expdata"].append(exp_dict["expdata_vl"])
-            self.experimental["expdata"].append(exp_dict["expdata_true"])
+        for replica in range(self._parallel_models):
+            replica_exp_info = self.exp_info[replica]
+            for exp_dict in replica_exp_info:
+                self.training["expdata"][replica].append(exp_dict["expdata"])
+                self.validation["expdata"][replica].append(exp_dict["expdata_vl"])
+                self.experimental["expdata"][replica].append(exp_dict["expdata_true"])
 
-            self.training["folds"].append(exp_dict["folds"]["training"])
-            self.validation["folds"].append(exp_dict["folds"]["validation"])
-            self.experimental["folds"].append(exp_dict["folds"]["experimental"])
+                nd_tr = exp_dict["ndata"]
+                nd_vl = exp_dict["ndata_vl"]
 
-            nd_tr = exp_dict["ndata"]
-            nd_vl = exp_dict["ndata_vl"]
-
-            self.training["ndata"] += nd_tr
-            self.validation["ndata"] += nd_vl
-            self.experimental["ndata"] += nd_tr + nd_vl
+                self.training["ndata"][replica] += nd_tr
+                self.validation["ndata"][replica] += nd_vl
+                self.experimental["ndata"][replica] += nd_tr + nd_vl
 
             for dataset in exp_dict["datasets"]:
-                self.all_datasets.append(dataset.name)
+                self.all_datasets.append(dataset["name"])
+
+            if replica == 0:
+                self.training["folds"].append(exp_dict["folds"]["training"])
+                self.validation["folds"].append(exp_dict["folds"]["validation"])
+                self.experimental["folds"].append(exp_dict["folds"]["experimental"])
+
+            for pos_dict in self.pos_info:
+                self.training["expdata"][replica].append(pos_dict["expdata"])
+                self.validation["expdata"][replica].append(pos_dict["expdata"])
+
+            if self.integ_info is not None:
+                for integ_dict in self.integ_info:
+                    self.training["expdata"][replica].append(integ_dict["expdata"])
+
         self.all_datasets = set(self.all_datasets)
 
         for pos_dict in self.pos_info:
-            self.training["expdata"].append(pos_dict["expdata"])
             self.training["posdatasets"].append(pos_dict["name"])
-            self.validation["expdata"].append(pos_dict["expdata"])
             self.validation["posdatasets"].append(pos_dict["name"])
+
         if self.integ_info is not None:
             for integ_dict in self.integ_info:
-                self.training["expdata"].append(integ_dict["expdata"])
                 self.training["integdatasets"].append(integ_dict["name"])
 
     def _xgrid_generation(self):
@@ -478,13 +487,17 @@ def _reset_observables(self):
         or be obliterated when/if the backend state is reset
         """
         self.input_list = []
-        for key in ["output", "posmultipliers", "integmultipliers"]:
+        self.input_sizes = []
+        self.training["output"] = [[] for _ in range(self._parallel_models)]
+        self.validation["output"] = [[] for _ in range(self._parallel_models)]
+        self.experimental["output"] = [[] for _ in range(self._parallel_models)]
+        for key in ["posmultipliers", "integmultipliers"]:
             self.training[key] = []
             self.validation[key] = []
             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                #
@@ -525,19 +538,22 @@ 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:
-            if not self.mode_hyperopt:
-                log.info("Generating layers for experiment %s", exp_dict["name"])
+        for replica in range(self._parallel_models):
+            for exp_dict in self.exp_info[replica]:
+                if not self.mode_hyperopt:
+                    log.info("Generating layers for experiment %s", exp_dict["name"])
+# TODO: Make this generator less memory-consuming by factorizing out the FK-tables mask into a separate layer...
+                exp_layer = model_gen.observable_generator(exp_dict)
 
-            exp_layer = model_gen.observable_generator(exp_dict)
+                # Save the input(s) corresponding to this experiment
+                if replica == 0:
+                    self.input_list.append(exp_layer["inputs"])
 
-            # Save the input(s) corresponding to this experiment
-            self.input_list.append(exp_layer["inputs"])
+                # Now save the observable layer, the losses and the experimental data
+                self.training["output"][replica].append(exp_layer["output_tr"])
+                self.validation["output"][replica].append(exp_layer["output_vl"])
+                self.experimental["output"][replica].append(exp_layer["output"])
 
-            # Now save the observable layer, the losses and the experimental data
-            self.training["output"].append(exp_layer["output_tr"])
-            self.validation["output"].append(exp_layer["output_vl"])
-            self.experimental["output"].append(exp_layer["output"])
 
         # Generate the positivity penalty
         for pos_dict in self.pos_info:
diff --git a/n3fit/src/n3fit/performfit.py b/n3fit/src/n3fit/performfit.py
index a98040dd96..b92460f3cf 100644
--- a/n3fit/src/n3fit/performfit.py
+++ b/n3fit/src/n3fit/performfit.py
@@ -154,6 +154,7 @@ def performfit(
         # 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])
+        # n_experiments dicts
         for i_exp in range(len(all_experiments)):
             training_data = []
             validation_data = []
@@ -167,17 +168,18 @@ def performfit(
             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)

From 685af8ab261c05b63a659c599a98df6988ffc7a2 Mon Sep 17 00:00:00 2001
From: gvdoord <g.vandenOord@gmail.com>
Date: Mon, 19 Dec 2022 14:34:53 +0100
Subject: [PATCH 02/18] Work on tr-vl masking layers, TF crashes with shape
 mismatch in gradient

---
 n3fit/src/n3fit/model_gen.py     |  95 ++++++++++++--------------
 n3fit/src/n3fit/model_trainer.py | 110 ++++++++++++++-----------------
 n3fit/src/n3fit/performfit.py    |  18 ++---
 3 files changed, 101 insertions(+), 122 deletions(-)

diff --git a/n3fit/src/n3fit/model_gen.py b/n3fit/src/n3fit/model_gen.py
index ffa0aba00a..c837e1b3b7 100644
--- a/n3fit/src/n3fit/model_gen.py
+++ b/n3fit/src/n3fit/model_gen.py
@@ -10,10 +10,12 @@
 
 """
 from dataclasses import dataclass
+
+import numpy
 import numpy as np
 from n3fit.msr import msr_impose
 from n3fit.layers import DIS, DY, ObsRotation, losses
-from n3fit.layers import Preprocessing, FkRotation, FlavourToEvolution
+from n3fit.layers import Preprocessing, FkRotation, FlavourToEvolution, Mask
 from n3fit.layers.observable import is_unique
 
 from n3fit.backends import MetaModel, Input
@@ -37,6 +39,7 @@ class ObservableWrapper:
 
     name: str
     observables: list
+    trvl_mask_layers: list
     dataset_xsizes: list
     invcovmat: np.array = None
     covmat: np.array = None
@@ -76,8 +79,16 @@ def _generate_experimental_layer(self, pdf):
         else:
             output_layers = [obs(pdf) for obs in self.observables]
 
+        masked_output_layers = []
+
+        for output_layer, mask_layer in zip(output_layers, self.trvl_mask_layers):
+            if mask_layer is None:
+                masked_output_layers.append(output_layer)
+            else:
+                masked_output_layers.append(mask_layer(output_layer))
+
         # Finally concatenate all observables (so that experiments are one single entitiy)
-        ret = op.concatenate(output_layers, axis=2)
+        ret = op.concatenate(masked_output_layers)
         if self.rotation is not None:
             ret = self.rotation(ret)
         return ret
@@ -89,7 +100,7 @@ def __call__(self, pdf_layer, mask=None):
 
 
 def observable_generator(
-    spec_dict, positivity_initial=1.0, integrability=False
+    spec_dict, mask_array=None, positivity_initial=1.0, integrability=False
 ):  # pylint: disable=too-many-locals
     """
     This function generates the observable models for each experiment.
@@ -138,9 +149,11 @@ 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
     for dataset in spec_dict["datasets"]:
         # Get the generic information of the dataset
@@ -155,56 +168,27 @@ 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...
+        trmask = mask_array[:, offset:offset + dataset.ndata] if apply_masks else None
+        tr_mask_layers.append(Mask(trmask, axis=1) if apply_masks else None)
+        vl_mask_layers.append(Mask(~trmask, axis=1) if apply_masks else None)
+
         # 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)
-# TODO: Make these more memory-efficient, i.e. factoring FKTables and separate Mask layer...
-        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))
@@ -213,9 +197,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
@@ -230,7 +215,8 @@ def observable_generator(
     if spec_dict["positivity"]:
         out_positivity = ObservableWrapper(
             spec_name,
-            model_obs_tr,
+            model_observables,
+            tr_mask_layers,
             dataset_xsizes,
             multiplier=positivity_initial,
             positivity=not integrability,
@@ -255,7 +241,8 @@ def observable_generator(
 
     out_tr = ObservableWrapper(
         spec_name,
-        model_obs_tr,
+        model_observables,
+        tr_mask_layers,
         dataset_xsizes,
         invcovmat=spec_dict["invcovmat"],
         data=spec_dict["expdata"],
@@ -263,7 +250,8 @@ def observable_generator(
     )
     out_vl = ObservableWrapper(
         f"{spec_name}_val",
-        model_obs_vl,
+        model_observables,
+        vl_mask_layers,
         dataset_xsizes,
         invcovmat=spec_dict["invcovmat_vl"],
         data=spec_dict["expdata_vl"],
@@ -271,7 +259,8 @@ def observable_generator(
     )
     out_exp = ObservableWrapper(
         f"{spec_name}_exp",
-        model_obs_ex,
+        model_observables,
+        [None] * len(model_observables),
         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 feee2a28a8..3df06e4cb0 100644
--- a/n3fit/src/n3fit/model_trainer.py
+++ b/n3fit/src/n3fit/model_trainer.py
@@ -11,6 +11,8 @@
 import logging
 from collections import namedtuple
 from itertools import zip_longest
+
+import numpy
 import numpy as np
 from scipy.interpolate import PchipInterpolator
 from n3fit import model_gen
@@ -137,9 +139,9 @@ def __init__(
         self.pos_info = pos_info
         self.integ_info = integ_info
         if self.integ_info is not None:
-            self.all_info = self.exp_info + pos_info + integ_info
+            self.all_info = self.exp_info[0] + pos_info + integ_info
         else:
-            self.all_info = self.exp_info + pos_info
+            self.all_info = self.exp_info[0] + pos_info
         self.flavinfo = flavinfo
         self.fitbasis = fitbasis
         self._nn_seeds = nnseeds
@@ -184,9 +186,9 @@ def __init__(
         # Initialize the dictionaries which contain all fitting information
         self.input_list = []
         self.training = {
-            "output": [[] for _ in range(parallel_models)],
-            "expdata": [[] for _ in range(parallel_models)],
-            "ndata": [0 for _ in range(parallel_models)],
+            "output": [],
+            "expdata": [],
+            "ndata": 0,
             "model": None,
             "posdatasets": [],
             "posmultipliers": [],
@@ -194,31 +196,31 @@ def __init__(
             "integdatasets": [],
             "integmultipliers": [],
             "integinitials": [],
-            "folds": [[] for _ in range(parallel_models)],
+            "folds": [],
         }
         self.validation = {
-            "output": [[] for _ in range(parallel_models)],
-            "expdata": [[] for _ in range(parallel_models)],
-            "ndata": [0 for _ in range(parallel_models)],
+            "output": [],
+            "expdata": [],
+            "ndata": 0,
             "model": None,
-            "folds": [[] for _ in range(parallel_models)],
+            "folds": [],
             "posdatasets": [],
         }
         self.experimental = {
-            "output": [[] for _ in range(parallel_models)],
-            "expdata": [[] for _ in range(parallel_models)],
-            "ndata": [0 for _ in range(parallel_models)],
+            "output": [],
+            "expdata": [],
+            "ndata": 0,
             "model": None,
-            "folds": [[] for _ in range(parallel_models)],
+            "folds": [],
         }
+        self.tr_masks = []
 
         self._fill_the_dictionaries()
 
-        if self.validation["ndata"][0] == 0:
+        if self.validation["ndata"] == 0:
             # If there is no validation, the validation chi2 = training chi2
             self.no_validation = True
-            for replica in range(parallel_models):
-                self.validation["expdata"][replica] = self.training["expdata"][replica]
+            self.validation["expdata"] = self.training["expdata"]
         else:
             # Consider the validation only if there is validation (of course)
             self.no_validation = False
@@ -261,44 +263,36 @@ def _fill_the_dictionaries(self):
             - ``name``: names of the experiment
             - ``ndata``: number of experimental points
         """
-        for replica in range(self._parallel_models):
-            replica_exp_info = self.exp_info[replica]
-            for exp_dict in replica_exp_info:
-                self.training["expdata"][replica].append(exp_dict["expdata"])
-                self.validation["expdata"][replica].append(exp_dict["expdata_vl"])
-                self.experimental["expdata"][replica].append(exp_dict["expdata_true"])
-
-                nd_tr = exp_dict["ndata"]
-                nd_vl = exp_dict["ndata_vl"]
 
-                self.training["ndata"][replica] += nd_tr
-                self.validation["ndata"][replica] += nd_vl
-                self.experimental["ndata"][replica] += nd_tr + nd_vl
-
-            for dataset in exp_dict["datasets"]:
-                self.all_datasets.append(dataset["name"])
+        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"])
 
-            if replica == 0:
-                self.training["folds"].append(exp_dict["folds"]["training"])
-                self.validation["folds"].append(exp_dict["folds"]["validation"])
-                self.experimental["folds"].append(exp_dict["folds"]["experimental"])
+            self.training["folds"].append(exp_dict["folds"]["training"])
+            self.validation["folds"].append(exp_dict["folds"]["validation"])
+            self.experimental["folds"].append(exp_dict["folds"]["experimental"])
 
-            for pos_dict in self.pos_info:
-                self.training["expdata"][replica].append(pos_dict["expdata"])
-                self.validation["expdata"][replica].append(pos_dict["expdata"])
+            nd_tr = exp_dict["ndata"]
+            nd_vl = exp_dict["ndata_vl"]
 
-            if self.integ_info is not None:
-                for integ_dict in self.integ_info:
-                    self.training["expdata"][replica].append(integ_dict["expdata"])
+            self.training["ndata"] += nd_tr
+            self.validation["ndata"] += nd_vl
+            self.experimental["ndata"] += nd_tr + nd_vl
 
+            for dataset in exp_dict["datasets"]:
+                self.all_datasets.append(dataset.name)
         self.all_datasets = set(self.all_datasets)
 
         for pos_dict in self.pos_info:
+            self.training["expdata"].append(pos_dict["expdata"])
             self.training["posdatasets"].append(pos_dict["name"])
+            self.validation["expdata"].append(pos_dict["expdata"])
             self.validation["posdatasets"].append(pos_dict["name"])
 
         if self.integ_info is not None:
             for integ_dict in self.integ_info:
+                self.training["expdata"].append(integ_dict["expdata"])
                 self.training["integdatasets"].append(integ_dict["name"])
 
     def _xgrid_generation(self):
@@ -487,11 +481,7 @@ def _reset_observables(self):
         or be obliterated when/if the backend state is reset
         """
         self.input_list = []
-        self.input_sizes = []
-        self.training["output"] = [[] for _ in range(self._parallel_models)]
-        self.validation["output"] = [[] for _ in range(self._parallel_models)]
-        self.experimental["output"] = [[] for _ in range(self._parallel_models)]
-        for key in ["posmultipliers", "integmultipliers"]:
+        for key in ["output", "posmultipliers", "integmultipliers"]:
             self.training[key] = []
             self.validation[key] = []
             self.experimental[key] = []
@@ -538,22 +528,22 @@ 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 replica in range(self._parallel_models):
-            for exp_dict in self.exp_info[replica]:
-                if not self.mode_hyperopt:
-                    log.info("Generating layers for experiment %s", exp_dict["name"])
-# TODO: Make this generator less memory-consuming by factorizing out the FK-tables mask into a separate layer...
-                exp_layer = model_gen.observable_generator(exp_dict)
+        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"])
+
+            # Stacked tr-vl mask array for all replicas for this dataset
+            replica_masks = numpy.stack([e[index]["trmask"] for e in self.exp_info])
 
-                # Save the input(s) corresponding to this experiment
-                if replica == 0:
-                    self.input_list.append(exp_layer["inputs"])
+            exp_layer = model_gen.observable_generator(exp_dict, mask_array=replica_masks)
 
-                # Now save the observable layer, the losses and the experimental data
-                self.training["output"][replica].append(exp_layer["output_tr"])
-                self.validation["output"][replica].append(exp_layer["output_vl"])
-                self.experimental["output"][replica].append(exp_layer["output"])
+            # Save the input(s) corresponding to this experiment
+            self.input_list.append(exp_layer["inputs"])
 
+            # Now save the observable layer, the losses and the experimental data
+            self.training["output"].append(exp_layer["output_tr"])
+            self.validation["output"].append(exp_layer["output_vl"])
+            self.experimental["output"].append(exp_layer["output"])
 
         # Generate the positivity penalty
         for pos_dict in self.pos_info:
diff --git a/n3fit/src/n3fit/performfit.py b/n3fit/src/n3fit/performfit.py
index b92460f3cf..d1e47f5ba9 100644
--- a/n3fit/src/n3fit/performfit.py
+++ b/n3fit/src/n3fit/performfit.py
@@ -153,16 +153,16 @@ 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])
+        # 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)
+        # 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],

From 94916ad50744217161ed233903be18a56780da02 Mon Sep 17 00:00:00 2001
From: goord <G.vandenOord@gmail.com>
Date: Sun, 25 Dec 2022 21:04:27 +0100
Subject: [PATCH 03/18] tr-vl masking layers working, validation needed

---
 n3fit/src/n3fit/backends/keras_backend/operations.py |  4 ++++
 n3fit/src/n3fit/layers/mask.py                       | 10 +++++++---
 n3fit/src/n3fit/model_gen.py                         |  7 ++++---
 3 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/n3fit/src/n3fit/backends/keras_backend/operations.py b/n3fit/src/n3fit/backends/keras_backend/operations.py
index 8cf1065849..f855b66a68 100644
--- a/n3fit/src/n3fit/backends/keras_backend/operations.py
+++ b/n3fit/src/n3fit/backends/keras_backend/operations.py
@@ -219,6 +219,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/layers/mask.py b/n3fit/src/n3fit/layers/mask.py
index 877df3b502..479afdeaff 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):
     """
@@ -25,8 +25,10 @@ class Mask(MetaLayer):
     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)
+            self.last_dim = count_nonzero(bool_mask[0,...])
         self.c = c
         self.axis = axis
         super().__init__(**kwargs)
@@ -34,14 +36,16 @@ def __init__(self, bool_mask=None, c=None, axis=None, **kwargs):
     def build(self, input_shape):
         if self.c is not None:
             initializer = MetaLayer.init_constant(value=self.c)
+            output_shape = list(input_shape)
+            output_shape[-1] = self.last_dim
             self.kernel = self.builder_helper(
-                "mask", (1,), initializer, trainable=False
+                "mask", output_shape, initializer, trainable=False
             )
         super(Mask, self).build(input_shape)
 
     def call(self, ret):
         if self.mask is not None:
-            ret = op.boolean_mask(ret, self.mask, axis=self.axis)
+            ret = op.reshape(op.boolean_mask(ret, self.mask, axis=self.axis), shape=(1, -1, self.last_dim))
         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 c837e1b3b7..206642fe83 100644
--- a/n3fit/src/n3fit/model_gen.py
+++ b/n3fit/src/n3fit/model_gen.py
@@ -23,6 +23,7 @@
 from n3fit.backends import MetaLayer, Lambda
 from n3fit.backends import base_layer_selector, regularizer_selector
 
+from keras.layers import Masking
 
 @dataclass
 class ObservableWrapper:
@@ -87,7 +88,7 @@ def _generate_experimental_layer(self, pdf):
             else:
                 masked_output_layers.append(mask_layer(output_layer))
 
-        # Finally concatenate all observables (so that experiments are one single entitiy)
+        # 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)
@@ -170,8 +171,8 @@ def observable_generator(
 
         # Extract the masks that will end up in the observable wrappers...
         trmask = mask_array[:, offset:offset + dataset.ndata] if apply_masks else None
-        tr_mask_layers.append(Mask(trmask, axis=1) if apply_masks else None)
-        vl_mask_layers.append(Mask(~trmask, axis=1) if apply_masks else None)
+        tr_mask_layers.append(Mask(trmask, axis=1, c=1) if apply_masks else None)
+        vl_mask_layers.append(Mask(~trmask, axis=1, c=1) if apply_masks else None)
 
         # Now generate the observable layer, which takes the following information:
         # operation name

From a43bd5de99bc999c68c5c609e55f448bfa26f539 Mon Sep 17 00:00:00 2001
From: goord <G.vandenOord@gmail.com>
Date: Wed, 15 Feb 2023 22:15:21 +0100
Subject: [PATCH 04/18] Fixed issues with replica-specific training data and
 inverse covariance matrix

---
 n3fit/src/n3fit/layers/losses.py |  9 ++++++---
 n3fit/src/n3fit/model_gen.py     | 23 +++++++++++++++--------
 n3fit/src/n3fit/model_trainer.py | 24 +++++++++++++++++++-----
 3 files changed, 40 insertions(+), 16 deletions(-)

diff --git a/n3fit/src/n3fit/layers/losses.py b/n3fit/src/n3fit/layers/losses.py
index 4315919907..378d4e768b 100644
--- a/n3fit/src/n3fit/layers/losses.py
+++ b/n3fit/src/n3fit/layers/losses.py
@@ -52,12 +52,12 @@ def __init__(self, invcovmat, y_true, mask=None, covmat=None, **kwargs):
             self._mask = op.numpy_to_tensor(mask)
         super().__init__(**kwargs)
 
-    def build(self, input_shape):
+    def build(self):
         """Transform the inverse covmat and the mask into
         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:
@@ -88,7 +88,10 @@ def call(self, y_pred, **kwargs):
             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/model_gen.py b/n3fit/src/n3fit/model_gen.py
index 206642fe83..dd1fa81fbd 100644
--- a/n3fit/src/n3fit/model_gen.py
+++ b/n3fit/src/n3fit/model_gen.py
@@ -55,8 +55,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:
@@ -100,8 +99,14 @@ def __call__(self, pdf_layer, mask=None):
         return loss_f(experiment_prediction)
 
 
-def observable_generator(
-    spec_dict, mask_array=None, 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.
@@ -156,6 +161,7 @@ def observable_generator(
     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
@@ -171,6 +177,7 @@ def observable_generator(
 
         # Extract the masks that will end up in the observable wrappers...
         trmask = mask_array[:, offset:offset + dataset.ndata] if apply_masks else None
+        masks.append(trmask)
         tr_mask_layers.append(Mask(trmask, axis=1, c=1) if apply_masks else None)
         vl_mask_layers.append(Mask(~trmask, axis=1, c=1) if apply_masks else None)
 
@@ -245,8 +252,8 @@ def observable_generator(
         model_observables,
         tr_mask_layers,
         dataset_xsizes,
-        invcovmat=spec_dict["invcovmat"],
-        data=spec_dict["expdata"],
+        invcovmat=invcovmat_tr,
+        data=training_data,
         rotation=obsrot_tr,
     )
     out_vl = ObservableWrapper(
@@ -254,8 +261,8 @@ def observable_generator(
         model_observables,
         vl_mask_layers,
         dataset_xsizes,
-        invcovmat=spec_dict["invcovmat_vl"],
-        data=spec_dict["expdata_vl"],
+        invcovmat=invcovmat_vl,
+        data=validation_data,
         rotation=obsrot_vl,
     )
     out_exp = ObservableWrapper(
diff --git a/n3fit/src/n3fit/model_trainer.py b/n3fit/src/n3fit/model_trainer.py
index 3df06e4cb0..d93777958b 100644
--- a/n3fit/src/n3fit/model_trainer.py
+++ b/n3fit/src/n3fit/model_trainer.py
@@ -263,7 +263,6 @@ def _fill_the_dictionaries(self):
             - ``name``: names of the experiment
             - ``ndata``: number of experimental points
         """
-
         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"])
@@ -534,8 +533,17 @@ def _generate_observables(
 
             # Stacked tr-vl mask array for all replicas for this dataset
             replica_masks = numpy.stack([e[index]["trmask"] for e in self.exp_info])
-
-            exp_layer = model_gen.observable_generator(exp_dict, mask_array=replica_masks)
+            training_data = numpy.stack([e[index]["expdata"].flatten() for e in self.exp_info])
+            validation_data = numpy.stack([e[index]["expdata_vl"].flatten() for e in self.exp_info])
+            invcovmat = numpy.stack([e[index]["invcovmat"] for e in self.exp_info])
+            invcovmat_vl = numpy.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"])
@@ -556,8 +564,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 = numpy.stack([pos_dict["trmask"] for i in range(len(self.exp_info))])
+            training_data = numpy.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"])
 

From f5e9a45e92e64014f3a1a20ddb95d142fe6f2ca0 Mon Sep 17 00:00:00 2001
From: goord <G.vandenOord@gmail.com>
Date: Wed, 15 Feb 2023 22:18:50 +0100
Subject: [PATCH 05/18] Added input shape to build method args again

---
 n3fit/src/n3fit/layers/losses.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/n3fit/src/n3fit/layers/losses.py b/n3fit/src/n3fit/layers/losses.py
index 378d4e768b..7911acb974 100644
--- a/n3fit/src/n3fit/layers/losses.py
+++ b/n3fit/src/n3fit/layers/losses.py
@@ -52,7 +52,7 @@ def __init__(self, invcovmat, y_true, mask=None, covmat=None, **kwargs):
             self._mask = op.numpy_to_tensor(mask)
         super().__init__(**kwargs)
 
-    def build(self):
+    def build(self, input_shape):
         """Transform the inverse covmat and the mask into
         weights of the layers"""
         init = MetaLayer.init_constant(self._invcovmat)

From 2ca63f02dd1830f4614e8b3ce9088090c462ad78 Mon Sep 17 00:00:00 2001
From: goord <G.vandenOord@gmail.com>
Date: Thu, 16 Feb 2023 09:02:38 +0100
Subject: [PATCH 06/18] Fixed broken sequential replica fits

---
 n3fit/src/n3fit/performfit.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/n3fit/src/n3fit/performfit.py b/n3fit/src/n3fit/performfit.py
index d1e47f5ba9..3aae82f445 100644
--- a/n3fit/src/n3fit/performfit.py
+++ b/n3fit/src/n3fit/performfit.py
@@ -176,7 +176,7 @@ def performfit(
         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]))
+            replicas_info[i] = (tuple([replica_idxs]), [info_tuple[1]], tuple([nnseeds]))
 
     for replica_idxs, exp_info, nnseeds in replicas_info:
         log.info("Starting replica fit " + str(replica_idxs))

From eb819af28581c1baa8ab187286384b114b0ba0ca Mon Sep 17 00:00:00 2001
From: goord <G.vandenOord@gmail.com>
Date: Mon, 13 Mar 2023 14:39:04 +0100
Subject: [PATCH 07/18] Add single-point datasets to training set for all
 parallel replicas

---
 validphys2/src/validphys/n3fit_data.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/validphys2/src/validphys/n3fit_data.py b/validphys2/src/validphys/n3fit_data.py
index 75e5efdc12..9e65247cc7 100644
--- a/validphys2/src/validphys/n3fit_data.py
+++ b/validphys2/src/validphys/n3fit_data.py
@@ -70,7 +70,7 @@ def __iter__(self):
             yield m
 
 
-def tr_masks(data, replica_trvlseed):
+def tr_masks(data, replica_trvlseed, parallel_models):
     """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
@@ -78,6 +78,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
@@ -94,7 +97,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)]
         )

From 7269c1d8b2b4ab90dce079cc0fbad37ca272e023 Mon Sep 17 00:00:00 2001
From: goord <G.vandenOord@gmail.com>
Date: Mon, 13 Mar 2023 15:28:33 +0100
Subject: [PATCH 08/18] Give better names to mask layers

---
 n3fit/src/n3fit/model_gen.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/n3fit/src/n3fit/model_gen.py b/n3fit/src/n3fit/model_gen.py
index dd1fa81fbd..9af73b738b 100644
--- a/n3fit/src/n3fit/model_gen.py
+++ b/n3fit/src/n3fit/model_gen.py
@@ -178,8 +178,8 @@ def observable_generator(spec_dict,
         # Extract the masks that will end up in the observable wrappers...
         trmask = mask_array[:, offset:offset + dataset.ndata] if apply_masks else None
         masks.append(trmask)
-        tr_mask_layers.append(Mask(trmask, axis=1, c=1) if apply_masks else None)
-        vl_mask_layers.append(Mask(~trmask, axis=1, c=1) if apply_masks else None)
+        tr_mask_layers.append(Mask(trmask, axis=1, c=1, name=f"trmask_{dataset_name}") if apply_masks else None)
+        vl_mask_layers.append(Mask(~trmask, axis=1, c=1, name=f"vlmask_{dataset_name}") if apply_masks else None)
 
         # Now generate the observable layer, which takes the following information:
         # operation name

From 303024e7a1cade4150caa2154a36f78d19d04257 Mon Sep 17 00:00:00 2001
From: goord <G.vandenOord@gmail.com>
Date: Mon, 1 May 2023 23:15:59 +0200
Subject: [PATCH 09/18] Fix mask layer kernel shape.

---
 n3fit/src/n3fit/layers/mask.py | 12 +++++++-----
 n3fit/src/n3fit/model_gen.py   |  7 ++-----
 2 files changed, 9 insertions(+), 10 deletions(-)

diff --git a/n3fit/src/n3fit/layers/mask.py b/n3fit/src/n3fit/layers/mask.py
index 479afdeaff..60a0426c63 100644
--- a/n3fit/src/n3fit/layers/mask.py
+++ b/n3fit/src/n3fit/layers/mask.py
@@ -19,7 +19,8 @@ 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):
@@ -36,16 +37,17 @@ def __init__(self, bool_mask=None, c=None, axis=None, **kwargs):
     def build(self, input_shape):
         if self.c is not None:
             initializer = MetaLayer.init_constant(value=self.c)
-            output_shape = list(input_shape)
-            output_shape[-1] = self.last_dim
             self.kernel = self.builder_helper(
-                "mask", output_shape, initializer, trainable=False
+                "mask", (1,), initializer, trainable=False
             )
         super(Mask, self).build(input_shape)
 
     def call(self, ret):
         if self.mask is not None:
-            ret = op.reshape(op.boolean_mask(ret, self.mask, axis=self.axis), shape=(1, -1, self.last_dim))
+            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 9af73b738b..9464f54c44 100644
--- a/n3fit/src/n3fit/model_gen.py
+++ b/n3fit/src/n3fit/model_gen.py
@@ -11,7 +11,6 @@
 """
 from dataclasses import dataclass
 
-import numpy
 import numpy as np
 from n3fit.msr import msr_impose
 from n3fit.layers import DIS, DY, ObsRotation, losses
@@ -23,8 +22,6 @@
 from n3fit.backends import MetaLayer, Lambda
 from n3fit.backends import base_layer_selector, regularizer_selector
 
-from keras.layers import Masking
-
 @dataclass
 class ObservableWrapper:
     """Wraps many observables into an experimental layer once the PDF model is prepared
@@ -178,8 +175,8 @@ def observable_generator(spec_dict,
         # Extract the masks that will end up in the observable wrappers...
         trmask = mask_array[:, offset:offset + dataset.ndata] if apply_masks else None
         masks.append(trmask)
-        tr_mask_layers.append(Mask(trmask, axis=1, c=1, name=f"trmask_{dataset_name}") if apply_masks else None)
-        vl_mask_layers.append(Mask(~trmask, axis=1, c=1, name=f"vlmask_{dataset_name}") if apply_masks else None)
+        tr_mask_layers.append(Mask(trmask, axis=1, name=f"trmask_{dataset_name}") if apply_masks else None)
+        vl_mask_layers.append(Mask(~trmask, axis=1, name=f"vlmask_{dataset_name}") if apply_masks else None)
 
         # Now generate the observable layer, which takes the following information:
         # operation name

From bc568761540fa4a96e91963ad314041fa81dc782 Mon Sep 17 00:00:00 2001
From: goord <G.vandenOord@gmail.com>
Date: Mon, 8 May 2023 14:50:58 +0200
Subject: [PATCH 10/18] Either apply masks to all observables or not

---
 n3fit/src/n3fit/model_gen.py | 26 +++++++++++++-------------
 1 file changed, 13 insertions(+), 13 deletions(-)

diff --git a/n3fit/src/n3fit/model_gen.py b/n3fit/src/n3fit/model_gen.py
index 9464f54c44..c42c2b1d88 100644
--- a/n3fit/src/n3fit/model_gen.py
+++ b/n3fit/src/n3fit/model_gen.py
@@ -77,12 +77,11 @@ def _generate_experimental_layer(self, pdf):
             output_layers = [obs(pdf) for obs in self.observables]
 
         masked_output_layers = []
-
-        for output_layer, mask_layer in zip(output_layers, self.trvl_mask_layers):
-            if mask_layer is None:
-                masked_output_layers.append(output_layer)
-            else:
+        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)
@@ -173,10 +172,11 @@ def observable_generator(spec_dict,
         operation_name = dataset.operation
 
         # Extract the masks that will end up in the observable wrappers...
-        trmask = mask_array[:, offset:offset + dataset.ndata] if apply_masks else None
-        masks.append(trmask)
-        tr_mask_layers.append(Mask(trmask, axis=1, name=f"trmask_{dataset_name}") if apply_masks else None)
-        vl_mask_layers.append(Mask(~trmask, axis=1, name=f"vlmask_{dataset_name}") if apply_masks else None)
+        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
@@ -221,7 +221,7 @@ def observable_generator(spec_dict,
         out_positivity = ObservableWrapper(
             spec_name,
             model_observables,
-            tr_mask_layers,
+            tr_mask_layers if apply_masks else None,
             dataset_xsizes,
             multiplier=positivity_initial,
             positivity=not integrability,
@@ -247,7 +247,7 @@ def observable_generator(spec_dict,
     out_tr = ObservableWrapper(
         spec_name,
         model_observables,
-        tr_mask_layers,
+        tr_mask_layers if apply_masks else None,
         dataset_xsizes,
         invcovmat=invcovmat_tr,
         data=training_data,
@@ -256,7 +256,7 @@ def observable_generator(spec_dict,
     out_vl = ObservableWrapper(
         f"{spec_name}_val",
         model_observables,
-        vl_mask_layers,
+        vl_mask_layers if apply_masks else None,
         dataset_xsizes,
         invcovmat=invcovmat_vl,
         data=validation_data,
@@ -265,7 +265,7 @@ def observable_generator(spec_dict,
     out_exp = ObservableWrapper(
         f"{spec_name}_exp",
         model_observables,
-        [None] * len(model_observables),
+        None,
         dataset_xsizes,
         invcovmat=spec_dict["invcovmat_true"],
         covmat=spec_dict["covmat"],

From 267d9730df3f1f46d8e951788236581064b05f4c Mon Sep 17 00:00:00 2001
From: goord <G.vandenOord@gmail.com>
Date: Wed, 10 May 2023 13:22:12 +0200
Subject: [PATCH 11/18] Wrapped load_with_cuts in lru cache function

---
 validphys2/src/validphys/n3fit_data_utils.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/validphys2/src/validphys/n3fit_data_utils.py b/validphys2/src/validphys/n3fit_data_utils.py
index 5743439946..21106f2848 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
 from itertools import zip_longest
 import dataclasses
 import numpy as np
@@ -94,7 +95,12 @@ def validphys_group_extractor(datasets, tr_masks):
     # Use zip_longest since tr_mask can be (and it is fine) an empty list
     for dspec, mask in zip_longest(datasets, tr_masks):
         # Load all fktables with the appropiate cuts
-        fktables = [fk.load_with_cuts(dspec.cuts) for fk in dspec.fkspecs]
+        fktables = [load_cached_fk_tables(fk, dspec.cuts) for fk in dspec.fkspecs]
         # And now put them in a FittableDataSet object which
         loaded_obs.append(FittableDataSet(dspec.name, fktables, dspec.op, dspec.frac, mask))
     return loaded_obs
+
+@functools.lru_cache
+def load_cached_fk_tables(fk, cuts):
+    return fk.load_with_cuts(cuts)
+

From 7505f9728999674c361131a1494a38ad6b9f2c19 Mon Sep 17 00:00:00 2001
From: goord <G.vandenOord@gmail.com>
Date: Wed, 10 May 2023 14:35:37 +0200
Subject: [PATCH 12/18] Better way of wrapping load_with_cuts in lru cache
 function

---
 validphys2/src/validphys/core.py             | 3 ++-
 validphys2/src/validphys/n3fit_data_utils.py | 7 +------
 2 files changed, 3 insertions(+), 7 deletions(-)

diff --git a/validphys2/src/validphys/core.py b/validphys2/src/validphys/core.py
index fc903bf969..c6ae2ea37b 100644
--- a/validphys2/src/validphys/core.py
+++ b/validphys2/src/validphys/core.py
@@ -486,8 +486,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/n3fit_data_utils.py b/validphys2/src/validphys/n3fit_data_utils.py
index 21106f2848..dbdff298d9 100644
--- a/validphys2/src/validphys/n3fit_data_utils.py
+++ b/validphys2/src/validphys/n3fit_data_utils.py
@@ -95,12 +95,7 @@ def validphys_group_extractor(datasets, tr_masks):
     # Use zip_longest since tr_mask can be (and it is fine) an empty list
     for dspec, mask in zip_longest(datasets, tr_masks):
         # Load all fktables with the appropiate cuts
-        fktables = [load_cached_fk_tables(fk, dspec.cuts) for fk in dspec.fkspecs]
+        fktables = [fk.load_with_cuts(dspec.cuts) for fk in dspec.fkspecs]
         # And now put them in a FittableDataSet object which
         loaded_obs.append(FittableDataSet(dspec.name, fktables, dspec.op, dspec.frac, mask))
     return loaded_obs
-
-@functools.lru_cache
-def load_cached_fk_tables(fk, cuts):
-    return fk.load_with_cuts(cuts)
-

From 653f580b187f133eda658d155ae85023afb678df Mon Sep 17 00:00:00 2001
From: goord <G.vandenOord@gmail.com>
Date: Mon, 15 May 2023 12:10:24 +0200
Subject: [PATCH 13/18] Fixed bug when running single replica in parallel

---
 n3fit/src/n3fit/layers/losses.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/n3fit/src/n3fit/layers/losses.py b/n3fit/src/n3fit/layers/losses.py
index 7911acb974..cf8befe4f3 100644
--- a/n3fit/src/n3fit/layers/losses.py
+++ b/n3fit/src/n3fit/layers/losses.py
@@ -85,8 +85,12 @@ 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:
             if len(self.kernel.shape) == 3:
                 res = op.einsum("bri, rij, brj -> r", tmp, self.kernel, tmp)

From 0409cf35dbc27d6792ed00e5e9f55aa140bdb8f4 Mon Sep 17 00:00:00 2001
From: goord <G.vandenOord@gmail.com>
Date: Mon, 5 Jun 2023 16:45:01 +0200
Subject: [PATCH 14/18] Fixed incorrect merge in model_trainer.py

---
 n3fit/src/n3fit/model_trainer.py | 28 +++++++++++++++++++++++-----
 1 file changed, 23 insertions(+), 5 deletions(-)

diff --git a/n3fit/src/n3fit/model_trainer.py b/n3fit/src/n3fit/model_trainer.py
index 66a70269c6..be1f642326 100644
--- a/n3fit/src/n3fit/model_trainer.py
+++ b/n3fit/src/n3fit/model_trainer.py
@@ -495,7 +495,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                #
@@ -536,11 +536,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"])
@@ -561,8 +573,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"])
 

From 623a2d7e4a684d879b4104da355c3dd00c3804d5 Mon Sep 17 00:00:00 2001
From: goord <G.vandenOord@gmail.com>
Date: Wed, 7 Jun 2023 11:21:10 +0200
Subject: [PATCH 15/18] Attempt to fix tests with default argument in validphys

---
 validphys2/src/validphys/n3fit_data.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/validphys2/src/validphys/n3fit_data.py b/validphys2/src/validphys/n3fit_data.py
index a47ed6c81c..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, parallel_models):
+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

From 7edad328537ad58b78557bf41d53c531944f15a9 Mon Sep 17 00:00:00 2001
From: goord <G.vandenOord@gmail.com>
Date: Wed, 7 Jun 2023 13:28:43 +0200
Subject: [PATCH 16/18] Fix 1-D mask test

---
 .../reproduce_nnpdf40/NNPDF40_nnlo_as_01180_1000.yml       | 7 +++++--
 n3fit/src/n3fit/layers/mask.py                             | 5 ++++-
 validphys2/src/validphys/photon/compute.py                 | 2 +-
 3 files changed, 10 insertions(+), 4 deletions(-)

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..0f9c174e1c 100644
--- a/n3fit/runcards/reproduce_nnpdf40/NNPDF40_nnlo_as_01180_1000.yml
+++ b/n3fit/runcards/reproduce_nnpdf40/NNPDF40_nnlo_as_01180_1000.yml
@@ -112,7 +112,7 @@ parameters: # This defines the parameter dictionary that is passed to the Model
     clipnorm: 6.073e-6
     learning_rate: 2.621e-3
     optimizer_name: Nadam
-  epochs: 17000
+  epochs: 100
   positivity:
     initial: 184.8
     multiplier:
@@ -125,6 +125,7 @@ parameters: # This defines the parameter dictionary that is passed to the Model
   threshold_chi2: 3.5
 
 fitting:
+  savepseudodata: false
   fitbasis: EVOL  # EVOL (7), EVOLQED (8), etc.
   basis:
   - {fl: sng, trainable: false, smallx: [1.094, 1.118], largex: [1.46, 3.003]}
@@ -163,4 +164,6 @@ integrability:
 
 ############################################################
 debug: false
-maxcores: 4
+maxcores: 8
+parallel_models: true
+same_trvl_per_replica: false
diff --git a/n3fit/src/n3fit/layers/mask.py b/n3fit/src/n3fit/layers/mask.py
index 60a0426c63..cc2ad4e552 100644
--- a/n3fit/src/n3fit/layers/mask.py
+++ b/n3fit/src/n3fit/layers/mask.py
@@ -29,7 +29,10 @@ def __init__(self, bool_mask=None, c=None, axis=None, **kwargs):
             self.last_dim = -1
         else:
             self.mask = op.numpy_to_tensor(bool_mask, dtype=bool)
-            self.last_dim = count_nonzero(bool_mask[0,...])
+            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)
diff --git a/validphys2/src/validphys/photon/compute.py b/validphys2/src/validphys/photon/compute.py
index 4dde4d952c..3c8309ab88 100644
--- a/validphys2/src/validphys/photon/compute.py
+++ b/validphys2/src/validphys/photon/compute.py
@@ -8,7 +8,7 @@
 from scipy.interpolate import interp1d
 import yaml
 
-from eko.io import EKO
+#from eko.io import EKO
 from n3fit.io.writer import XGRID
 from validphys.n3fit_data import replica_luxseed
 

From 79240ad34b018e5e40d7a184104e4bd682e5c2ab Mon Sep 17 00:00:00 2001
From: goord <G.vandenOord@gmail.com>
Date: Wed, 7 Jun 2023 13:33:36 +0200
Subject: [PATCH 17/18] Reverted accidental commits

---
 .../reproduce_nnpdf40/NNPDF40_nnlo_as_01180_1000.yml       | 7 ++-----
 validphys2/src/validphys/photon/compute.py                 | 2 +-
 2 files changed, 3 insertions(+), 6 deletions(-)

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 0f9c174e1c..4e2295470f 100644
--- a/n3fit/runcards/reproduce_nnpdf40/NNPDF40_nnlo_as_01180_1000.yml
+++ b/n3fit/runcards/reproduce_nnpdf40/NNPDF40_nnlo_as_01180_1000.yml
@@ -112,7 +112,7 @@ parameters: # This defines the parameter dictionary that is passed to the Model
     clipnorm: 6.073e-6
     learning_rate: 2.621e-3
     optimizer_name: Nadam
-  epochs: 100
+  epochs: 17000
   positivity:
     initial: 184.8
     multiplier:
@@ -125,7 +125,6 @@ parameters: # This defines the parameter dictionary that is passed to the Model
   threshold_chi2: 3.5
 
 fitting:
-  savepseudodata: false
   fitbasis: EVOL  # EVOL (7), EVOLQED (8), etc.
   basis:
   - {fl: sng, trainable: false, smallx: [1.094, 1.118], largex: [1.46, 3.003]}
@@ -164,6 +163,4 @@ integrability:
 
 ############################################################
 debug: false
-maxcores: 8
-parallel_models: true
-same_trvl_per_replica: false
+maxcores: 4
\ No newline at end of file
diff --git a/validphys2/src/validphys/photon/compute.py b/validphys2/src/validphys/photon/compute.py
index 3c8309ab88..4dde4d952c 100644
--- a/validphys2/src/validphys/photon/compute.py
+++ b/validphys2/src/validphys/photon/compute.py
@@ -8,7 +8,7 @@
 from scipy.interpolate import interp1d
 import yaml
 
-#from eko.io import EKO
+from eko.io import EKO
 from n3fit.io.writer import XGRID
 from validphys.n3fit_data import replica_luxseed
 

From 1aab479584a6bf17be3ce6a149200abdf3fb7eeb Mon Sep 17 00:00:00 2001
From: goord <G.vandenOord@gmail.com>
Date: Mon, 24 Jul 2023 22:23:21 +0200
Subject: [PATCH 18/18] various caches speeding up multii-replica init

---
 validphys2/src/validphys/commondata.py |  3 ++-
 validphys2/src/validphys/config.py     |  8 ++++++--
 validphys2/src/validphys/covmats.py    |  2 ++
 validphys2/src/validphys/filters.py    |  4 ++++
 validphys2/src/validphys/utils.py      | 25 +++++++++++++++++++++++++
 5 files changed, 39 insertions(+), 3 deletions(-)

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 e24950f39d..8f19e43633 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/covmats.py b/validphys2/src/validphys/covmats.py
index cb0adf6e68..8eb8611aeb 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/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