From 2a49b6052c857fc706fb79d2d04dd79d4725e164 Mon Sep 17 00:00:00 2001
From: Aron <aronpjansen@gmail.com>
Date: Fri, 8 Dec 2023 12:36:04 +0100
Subject: [PATCH 01/30] Move replica loop into generate_nn function

---
 n3fit/src/n3fit/model_gen.py | 70 ++++++++++++++++++------------------
 1 file changed, 34 insertions(+), 36 deletions(-)

diff --git a/n3fit/src/n3fit/model_gen.py b/n3fit/src/n3fit/model_gen.py
index ae6ed01182..6ca13d463b 100644
--- a/n3fit/src/n3fit/model_gen.py
+++ b/n3fit/src/n3fit/model_gen.py
@@ -670,7 +670,6 @@ def pdfNN_layer_generator(
         sumrule_layer = lambda x: x
 
     # Only these layers change from replica to replica:
-    nn_replicas = []
     preprocessing_factor_replicas = []
     for i_replica, replica_seed in enumerate(seed):
         preprocessing_factor_replicas.append(
@@ -682,21 +681,19 @@ def pdfNN_layer_generator(
                 large_x=not subtract_one,
             )
         )
-        nn_replicas.append(
-            generate_nn(
-                layer_type=layer_type,
-                input_dimensions=nn_input_dimensions,
-                nodes=nodes,
-                activations=activations,
-                initializer_name=initializer_name,
-                replica_seed=replica_seed,
-                dropout=dropout,
-                regularizer=regularizer,
-                regularizer_args=regularizer_args,
-                last_layer_nodes=last_layer_nodes,
-                name=f"NN_{i_replica}",
-            )
-        )
+
+    nn_replicas = generate_nn(
+        layer_type=layer_type,
+        input_dimensions=nn_input_dimensions,
+        nodes=nodes,
+        activations=activations,
+        initializer_name=initializer_name,
+        replica_seeds=seed,
+        dropout=dropout,
+        regularizer=regularizer,
+        regularizer_args=regularizer_args,
+        last_layer_nodes=last_layer_nodes,
+    )
 
     # Apply NN layers for all replicas to a given input grid
     def neural_network_replicas(x, postfix=""):
@@ -784,40 +781,41 @@ def generate_nn(
     nodes: List[int],
     activations: List[str],
     initializer_name: str,
-    replica_seed: int,
+    replica_seeds: List[int],
     dropout: float,
     regularizer: str,
     regularizer_args: dict,
     last_layer_nodes: int,
-    name: str,
 ) -> MetaModel:
     """
     Create the part of the model that contains all of the actual neural network
     layers.
     """
+    x = Input(shape=(None, input_dimensions), batch_size=1, name='xgrids_processed')
+
     common_args = {
         'nodes_in': input_dimensions,
         'nodes': nodes,
         'activations': activations,
         'initializer_name': initializer_name,
-        'seed': replica_seed,
     }
-    if layer_type == "dense":
-        reg = regularizer_selector(regularizer, **regularizer_args)
-        list_of_pdf_layers = generate_dense_network(
-            **common_args, dropout_rate=dropout, regularizer=reg
-        )
-    elif layer_type == "dense_per_flavour":
-        list_of_pdf_layers = generate_dense_per_flavour_network(
-            **common_args, basis_size=last_layer_nodes
-        )
 
-    # Note: using a Sequential model would be more appropriate, but it would require
-    # creating a MetaSequential model.
-    x = Input(shape=(None, input_dimensions), batch_size=1, name='xgrids_processed')
-    pdf = x
-    for layer in list_of_pdf_layers:
-        pdf = layer(pdf)
+    models = []
+    for i_replica, replica_seed in enumerate(replica_seeds):
+        if layer_type == "dense":
+            reg = regularizer_selector(regularizer, **regularizer_args)
+            list_of_pdf_layers = generate_dense_network(
+                **common_args, dropout_rate=dropout, regularizer=reg, seed=replica_seed
+            )
+        elif layer_type == "dense_per_flavour":
+            list_of_pdf_layers = generate_dense_per_flavour_network(
+                **common_args, basis_size=last_layer_nodes, seed=replica_seed
+            )
+
+        pdf = x
+        for layer in list_of_pdf_layers:
+            pdf = layer(pdf)
+
+        models.append(MetaModel({'NN_input': x}, pdf, name=f"NN_{i_replica}"))
 
-    model = MetaModel({'NN_input': x}, pdf, name=name)
-    return model
+    return models

From 9c23fa5c51d13969da5fca53a2665701a39761e3 Mon Sep 17 00:00:00 2001
From: Aron <aronpjansen@gmail.com>
Date: Fri, 8 Dec 2023 13:02:15 +0100
Subject: [PATCH 02/30] Simplify handling of dropout

---
 n3fit/src/n3fit/model_gen.py | 15 ++++++---------
 1 file changed, 6 insertions(+), 9 deletions(-)

diff --git a/n3fit/src/n3fit/model_gen.py b/n3fit/src/n3fit/model_gen.py
index 6ca13d463b..aef8c0ee78 100644
--- a/n3fit/src/n3fit/model_gen.py
+++ b/n3fit/src/n3fit/model_gen.py
@@ -317,16 +317,7 @@ def generate_dense_network(
     the output layer for the basis choice)
     """
     list_of_pdf_layers = []
-    number_of_layers = len(nodes)
-    if dropout_rate > 0:
-        dropout_layer = number_of_layers - 2
-    else:
-        dropout_layer = -1
     for i, (nodes_out, activation) in enumerate(zip(nodes, activations)):
-        # if we have dropout set up, add it to the list
-        if dropout_rate > 0 and i == dropout_layer:
-            list_of_pdf_layers.append(base_layer_selector("dropout", rate=dropout_rate))
-
         # select the initializer and move the seed
         init = MetaLayer.select_initializer(initializer_name, seed=seed + i)
 
@@ -343,6 +334,12 @@ def generate_dense_network(
 
         list_of_pdf_layers.append(layer)
         nodes_in = int(nodes_out)
+
+    # add dropout as second to last layer
+    if dropout_rate > 0:
+        dropout_layer = MetaLayer.base_layer_selector("dropout", rate=dropout_rate)
+        list_of_pdf_layers.insert(dropout_layer, -2)
+
     return list_of_pdf_layers
 
 

From 683e3540bc794cbc457b8a8a383bf881dbcf1e1e Mon Sep 17 00:00:00 2001
From: Aron <aronpjansen@gmail.com>
Date: Fri, 8 Dec 2023 13:15:07 +0100
Subject: [PATCH 03/30] Factor out layer_generator in generate_dense_network

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

diff --git a/n3fit/src/n3fit/model_gen.py b/n3fit/src/n3fit/model_gen.py
index aef8c0ee78..fbdafcb2af 100644
--- a/n3fit/src/n3fit/model_gen.py
+++ b/n3fit/src/n3fit/model_gen.py
@@ -316,12 +316,9 @@ def generate_dense_network(
     for the next to last layer (i.e., the last layer of the dense network before getting to
     the output layer for the basis choice)
     """
-    list_of_pdf_layers = []
-    for i, (nodes_out, activation) in enumerate(zip(nodes, activations)):
-        # select the initializer and move the seed
-        init = MetaLayer.select_initializer(initializer_name, seed=seed + i)
 
-        # set the arguments that will define the layer
+    def layer_generator(nodes_in, nodes_out, activation, seed):
+        init = MetaLayer.select_initializer(initializer_name, seed=seed)
         arguments = {
             "kernel_initializer": init,
             "units": int(nodes_out),
@@ -329,9 +326,11 @@ def generate_dense_network(
             "input_shape": (nodes_in,),
             "kernel_regularizer": regularizer,
         }
+        return base_layer_selector("dense", **arguments)
 
-        layer = base_layer_selector("dense", **arguments)
-
+    list_of_pdf_layers = []
+    for i, (nodes_out, activation) in enumerate(zip(nodes, activations)):
+        layer = layer_generator(nodes_in, nodes_out, activation, seed + i)
         list_of_pdf_layers.append(layer)
         nodes_in = int(nodes_out)
 

From d02e1185d7edb41d2c3ea0ae643aa98f5e6ecd30 Mon Sep 17 00:00:00 2001
From: Aron <aronpjansen@gmail.com>
Date: Fri, 8 Dec 2023 13:31:54 +0100
Subject: [PATCH 04/30] Refactor dense_per_flavor_network

---
 n3fit/src/n3fit/model_gen.py | 52 +++++++++++++++++-------------------
 1 file changed, 24 insertions(+), 28 deletions(-)

diff --git a/n3fit/src/n3fit/model_gen.py b/n3fit/src/n3fit/model_gen.py
index fbdafcb2af..06570c45b6 100644
--- a/n3fit/src/n3fit/model_gen.py
+++ b/n3fit/src/n3fit/model_gen.py
@@ -349,23 +349,17 @@ def generate_dense_per_flavour_network(
     For each flavour generates a dense network of the chosen size
 
     """
-    list_of_pdf_layers = []
-    number_of_layers = len(nodes)
-    current_seed = seed
-    for i, (nodes_out, activation) in enumerate(zip(nodes, activations)):
-        initializers = []
-        for _ in range(basis_size):
-            # select the initializer and move the seed
-            initializers.append(MetaLayer.select_initializer(initializer_name, seed=current_seed))
-            current_seed += 1
-
-        # set the arguments that will define the layer
-        # but careful, the last layer must be nodes = 1
-        # TODO the mismatch is due to the fact that basis_size
-        # is set to the number of nodes of the last layer when it should
-        # come from the runcard
-        if i == number_of_layers - 1:
-            nodes_out = 1
+    # set the arguments that will define the layer
+    # but careful, the last layer must be nodes = 1
+    # TODO the mismatch is due to the fact that basis_size
+    # is set to the number of nodes of the last layer when it should
+    # come from the runcard
+    nodes[-1] = 1
+
+    def layer_generator(nodes_in, nodes_out, activation, seed):
+        initializers = [
+            MetaLayer.select_initializer(initializer_name, seed=seed + b) for b in range(basis_size)
+        ]
         arguments = {
             "kernel_initializer": initializers,
             "units": nodes_out,
@@ -373,22 +367,24 @@ def generate_dense_per_flavour_network(
             "input_shape": (nodes_in,),
             "basis_size": basis_size,
         }
+        return base_layer_selector("dense_per_flavour", **arguments)
 
-        layer = base_layer_selector("dense_per_flavour", **arguments)
+    list_of_pdf_layers = []
+    for i, (nodes_out, activation) in enumerate(zip(nodes, activations)):
+        layer = layer_generator(nodes_in, nodes_out, activation, seed + i)
+        list_of_pdf_layers.append(layer)
+        nodes_in = int(nodes_out)
 
-        if i == number_of_layers - 1:
-            # For the last layer, apply concatenate
-            concat = base_layer_selector("concatenate")
+    # For the last layer, apply concatenate
+    last_layer = list_of_pdf_layers[-1]
+    concat = base_layer_selector("concatenate")
 
-            def output_layer(ilayer):
-                result = layer(ilayer)
-                return concat(result)
+    def concatenated_last_layer(inputs):
+        result = last_layer(inputs)
+        return concat(result)
 
-            list_of_pdf_layers.append(output_layer)
-        else:
-            list_of_pdf_layers.append(layer)
+    list_of_pdf_layers[-1] = concatenated_last_layer
 
-        nodes_in = int(nodes_out)
     return list_of_pdf_layers
 
 

From f907d7fab218f649d55ed50d08e55c5f0b5545a1 Mon Sep 17 00:00:00 2001
From: Aron <aronpjansen@gmail.com>
Date: Fri, 8 Dec 2023 13:39:16 +0100
Subject: [PATCH 05/30] Move setting of last nodes to generate_nn

---
 n3fit/src/n3fit/model_gen.py | 14 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

diff --git a/n3fit/src/n3fit/model_gen.py b/n3fit/src/n3fit/model_gen.py
index 06570c45b6..fb74890c2d 100644
--- a/n3fit/src/n3fit/model_gen.py
+++ b/n3fit/src/n3fit/model_gen.py
@@ -349,12 +349,6 @@ def generate_dense_per_flavour_network(
     For each flavour generates a dense network of the chosen size
 
     """
-    # set the arguments that will define the layer
-    # but careful, the last layer must be nodes = 1
-    # TODO the mismatch is due to the fact that basis_size
-    # is set to the number of nodes of the last layer when it should
-    # come from the runcard
-    nodes[-1] = 1
 
     def layer_generator(nodes_in, nodes_out, activation, seed):
         initializers = [
@@ -785,6 +779,14 @@ def generate_nn(
     """
     x = Input(shape=(None, input_dimensions), batch_size=1, name='xgrids_processed')
 
+    if layer_type == "dense_per_flavour":
+        # set the arguments that will define the layer
+        # but careful, the last layer must be nodes = 1
+        # TODO the mismatch is due to the fact that basis_size
+        # is set to the number of nodes of the last layer when it should
+        # come from the runcard
+        nodes[-1] = 1
+
     common_args = {
         'nodes_in': input_dimensions,
         'nodes': nodes,

From aa76bc77458d330e80cd8b1c127254c2563501d7 Mon Sep 17 00:00:00 2001
From: Aron <aronpjansen@gmail.com>
Date: Fri, 8 Dec 2023 13:44:43 +0100
Subject: [PATCH 06/30] Add constant arguments

---
 n3fit/src/n3fit/model_gen.py | 17 ++++++++++-------
 1 file changed, 10 insertions(+), 7 deletions(-)

diff --git a/n3fit/src/n3fit/model_gen.py b/n3fit/src/n3fit/model_gen.py
index fb74890c2d..8fc4a9374f 100644
--- a/n3fit/src/n3fit/model_gen.py
+++ b/n3fit/src/n3fit/model_gen.py
@@ -787,29 +787,32 @@ def generate_nn(
         # come from the runcard
         nodes[-1] = 1
 
-    common_args = {
+    constant_args = {
         'nodes_in': input_dimensions,
         'nodes': nodes,
         'activations': activations,
         'initializer_name': initializer_name,
     }
+    if layer_type == "dense_per_flavour":
+        constant_args['basis_size'] = last_layer_nodes
+    if layer_type == "dense":
+        constant_args['dropout_rate'] = dropout
+        reg = regularizer_selector(regularizer, **regularizer_args)
+        constant_args['regularizer'] = reg
 
     models = []
     for i_replica, replica_seed in enumerate(replica_seeds):
         if layer_type == "dense":
-            reg = regularizer_selector(regularizer, **regularizer_args)
-            list_of_pdf_layers = generate_dense_network(
-                **common_args, dropout_rate=dropout, regularizer=reg, seed=replica_seed
-            )
+            list_of_pdf_layers = generate_dense_network(**constant_args, seed=replica_seed)
         elif layer_type == "dense_per_flavour":
             list_of_pdf_layers = generate_dense_per_flavour_network(
-                **common_args, basis_size=last_layer_nodes, seed=replica_seed
+                **constant_args, seed=replica_seed
             )
 
+        # apply layers to create model
         pdf = x
         for layer in list_of_pdf_layers:
             pdf = layer(pdf)
-
         models.append(MetaModel({'NN_input': x}, pdf, name=f"NN_{i_replica}"))
 
     return models

From 1ad7960b9f56622afc3b6ccbd111890635af7036 Mon Sep 17 00:00:00 2001
From: Aron <aronpjansen@gmail.com>
Date: Fri, 8 Dec 2023 13:46:58 +0100
Subject: [PATCH 07/30] Add constant arguments

---
 n3fit/src/n3fit/model_gen.py | 9 +++------
 1 file changed, 3 insertions(+), 6 deletions(-)

diff --git a/n3fit/src/n3fit/model_gen.py b/n3fit/src/n3fit/model_gen.py
index 8fc4a9374f..a1f9f963d1 100644
--- a/n3fit/src/n3fit/model_gen.py
+++ b/n3fit/src/n3fit/model_gen.py
@@ -795,19 +795,16 @@ def generate_nn(
     }
     if layer_type == "dense_per_flavour":
         constant_args['basis_size'] = last_layer_nodes
+        layers_generator = generate_dense_per_flavour_network
     if layer_type == "dense":
         constant_args['dropout_rate'] = dropout
         reg = regularizer_selector(regularizer, **regularizer_args)
         constant_args['regularizer'] = reg
+        layers_generator = generate_dense_network
 
     models = []
     for i_replica, replica_seed in enumerate(replica_seeds):
-        if layer_type == "dense":
-            list_of_pdf_layers = generate_dense_network(**constant_args, seed=replica_seed)
-        elif layer_type == "dense_per_flavour":
-            list_of_pdf_layers = generate_dense_per_flavour_network(
-                **constant_args, seed=replica_seed
-            )
+        list_of_pdf_layers = layers_generator(**constant_args, seed=replica_seed)
 
         # apply layers to create model
         pdf = x

From 7758497da4b438543e6e9765d760f80c23601443 Mon Sep 17 00:00:00 2001
From: Aron <aronpjansen@gmail.com>
Date: Fri, 8 Dec 2023 13:50:43 +0100
Subject: [PATCH 08/30] Move dropout to generate_nn

---
 n3fit/src/n3fit/model_gen.py | 16 +++++-----------
 1 file changed, 5 insertions(+), 11 deletions(-)

diff --git a/n3fit/src/n3fit/model_gen.py b/n3fit/src/n3fit/model_gen.py
index a1f9f963d1..e8de611aa8 100644
--- a/n3fit/src/n3fit/model_gen.py
+++ b/n3fit/src/n3fit/model_gen.py
@@ -306,15 +306,10 @@ def generate_dense_network(
     activations: List[str],
     initializer_name: str = "glorot_normal",
     seed: int = 0,
-    dropout_rate: float = 0.0,
     regularizer: str = None,
 ):
     """
     Generates a dense network
-
-    the dropout rate, if selected, is set
-    for the next to last layer (i.e., the last layer of the dense network before getting to
-    the output layer for the basis choice)
     """
 
     def layer_generator(nodes_in, nodes_out, activation, seed):
@@ -334,11 +329,6 @@ def layer_generator(nodes_in, nodes_out, activation, seed):
         list_of_pdf_layers.append(layer)
         nodes_in = int(nodes_out)
 
-    # add dropout as second to last layer
-    if dropout_rate > 0:
-        dropout_layer = MetaLayer.base_layer_selector("dropout", rate=dropout_rate)
-        list_of_pdf_layers.insert(dropout_layer, -2)
-
     return list_of_pdf_layers
 
 
@@ -797,7 +787,6 @@ def generate_nn(
         constant_args['basis_size'] = last_layer_nodes
         layers_generator = generate_dense_per_flavour_network
     if layer_type == "dense":
-        constant_args['dropout_rate'] = dropout
         reg = regularizer_selector(regularizer, **regularizer_args)
         constant_args['regularizer'] = reg
         layers_generator = generate_dense_network
@@ -806,6 +795,11 @@ def generate_nn(
     for i_replica, replica_seed in enumerate(replica_seeds):
         list_of_pdf_layers = layers_generator(**constant_args, seed=replica_seed)
 
+        # add dropout as second to last layer
+        if dropout > 0:
+            dropout_layer = MetaLayer.base_layer_selector("dropout", rate=dropout)
+            list_of_pdf_layers.insert(dropout_layer, -2)
+
         # apply layers to create model
         pdf = x
         for layer in list_of_pdf_layers:

From 2d388d922d795f24b67a62ba2b6fd0108b7fe0f9 Mon Sep 17 00:00:00 2001
From: Aron <aronpjansen@gmail.com>
Date: Fri, 8 Dec 2023 13:55:42 +0100
Subject: [PATCH 09/30] Move concatenation of per_flavor layers into
 generate_nn

---
 n3fit/src/n3fit/model_gen.py | 21 +++++++++++----------
 1 file changed, 11 insertions(+), 10 deletions(-)

diff --git a/n3fit/src/n3fit/model_gen.py b/n3fit/src/n3fit/model_gen.py
index e8de611aa8..40beedb8b7 100644
--- a/n3fit/src/n3fit/model_gen.py
+++ b/n3fit/src/n3fit/model_gen.py
@@ -359,16 +359,6 @@ def layer_generator(nodes_in, nodes_out, activation, seed):
         list_of_pdf_layers.append(layer)
         nodes_in = int(nodes_out)
 
-    # For the last layer, apply concatenate
-    last_layer = list_of_pdf_layers[-1]
-    concat = base_layer_selector("concatenate")
-
-    def concatenated_last_layer(inputs):
-        result = last_layer(inputs)
-        return concat(result)
-
-    list_of_pdf_layers[-1] = concatenated_last_layer
-
     return list_of_pdf_layers
 
 
@@ -800,6 +790,17 @@ def generate_nn(
             dropout_layer = MetaLayer.base_layer_selector("dropout", rate=dropout)
             list_of_pdf_layers.insert(dropout_layer, -2)
 
+        # In case of per flavour network, concatenate at the last layer
+        if layer_type == "dense_per_flavour":
+            last_layer = list_of_pdf_layers[-1]
+            concat = base_layer_selector("concatenate")
+
+            def concatenated_last_layer(inputs):
+                result = last_layer(inputs)
+                return concat(result)
+
+            list_of_pdf_layers[-1] = concatenated_last_layer
+
         # apply layers to create model
         pdf = x
         for layer in list_of_pdf_layers:

From 1ef87dc1a300ace72181ec0e0421e19890b985ad Mon Sep 17 00:00:00 2001
From: Aron <aronpjansen@gmail.com>
Date: Fri, 8 Dec 2023 14:05:19 +0100
Subject: [PATCH 10/30] Make the two layer generators almost equal

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

diff --git a/n3fit/src/n3fit/model_gen.py b/n3fit/src/n3fit/model_gen.py
index 40beedb8b7..6e0252bd6c 100644
--- a/n3fit/src/n3fit/model_gen.py
+++ b/n3fit/src/n3fit/model_gen.py
@@ -311,21 +311,23 @@ def generate_dense_network(
     """
     Generates a dense network
     """
+    layer_type = "dense"
+    custom_args = {"kernel_regularizer": regularizer}
 
-    def layer_generator(nodes_in, nodes_out, activation, seed):
-        init = MetaLayer.select_initializer(initializer_name, seed=seed)
-        arguments = {
-            "kernel_initializer": init,
-            "units": int(nodes_out),
-            "activation": activation,
-            "input_shape": (nodes_in,),
-            "kernel_regularizer": regularizer,
-        }
-        return base_layer_selector("dense", **arguments)
+    def initializer_generator(initializer_name, seed):
+        return MetaLayer.select_initializer(initializer_name, seed=seed)
 
     list_of_pdf_layers = []
     for i, (nodes_out, activation) in enumerate(zip(nodes, activations)):
-        layer = layer_generator(nodes_in, nodes_out, activation, seed + i)
+        init = initializer_generator(initializer_name, seed + i)
+        layer = base_layer_selector(
+            layer_type,
+            kernel_initializer=init,
+            units=nodes_out,
+            activation=activation,
+            input_shape=(nodes_in,),
+            **custom_args,
+        )
         list_of_pdf_layers.append(layer)
         nodes_in = int(nodes_out)
 
@@ -339,23 +341,26 @@ def generate_dense_per_flavour_network(
     For each flavour generates a dense network of the chosen size
 
     """
+    layer_type = "dense_per_flavour"
+    custom_args = {"basis_size": basis_size}
 
-    def layer_generator(nodes_in, nodes_out, activation, seed):
+    def initializer_generator(initializer_name, seed):
         initializers = [
             MetaLayer.select_initializer(initializer_name, seed=seed + b) for b in range(basis_size)
         ]
-        arguments = {
-            "kernel_initializer": initializers,
-            "units": nodes_out,
-            "activation": activation,
-            "input_shape": (nodes_in,),
-            "basis_size": basis_size,
-        }
-        return base_layer_selector("dense_per_flavour", **arguments)
+        return initializers
 
     list_of_pdf_layers = []
     for i, (nodes_out, activation) in enumerate(zip(nodes, activations)):
-        layer = layer_generator(nodes_in, nodes_out, activation, seed + i)
+        init = initializer_generator(initializer_name, seed + i)
+        layer = base_layer_selector(
+            layer_type,
+            kernel_initializer=init,
+            units=nodes_out,
+            activation=activation,
+            input_shape=(nodes_in,),
+            **custom_args,
+        )
         list_of_pdf_layers.append(layer)
         nodes_in = int(nodes_out)
 

From 806e2c15686582c125300e2899567f7f55e71eb3 Mon Sep 17 00:00:00 2001
From: Aron <aronpjansen@gmail.com>
Date: Fri, 8 Dec 2023 14:19:07 +0100
Subject: [PATCH 11/30] remove separate dense and dense_per_flavor functions

---
 n3fit/src/n3fit/model_gen.py | 114 +++++++++--------------------------
 1 file changed, 30 insertions(+), 84 deletions(-)

diff --git a/n3fit/src/n3fit/model_gen.py b/n3fit/src/n3fit/model_gen.py
index 6e0252bd6c..3b9cdd8b19 100644
--- a/n3fit/src/n3fit/model_gen.py
+++ b/n3fit/src/n3fit/model_gen.py
@@ -299,74 +299,6 @@ def observable_generator(
     return layer_info
 
 
-# Network generation functions
-def generate_dense_network(
-    nodes_in: int,
-    nodes: int,
-    activations: List[str],
-    initializer_name: str = "glorot_normal",
-    seed: int = 0,
-    regularizer: str = None,
-):
-    """
-    Generates a dense network
-    """
-    layer_type = "dense"
-    custom_args = {"kernel_regularizer": regularizer}
-
-    def initializer_generator(initializer_name, seed):
-        return MetaLayer.select_initializer(initializer_name, seed=seed)
-
-    list_of_pdf_layers = []
-    for i, (nodes_out, activation) in enumerate(zip(nodes, activations)):
-        init = initializer_generator(initializer_name, seed + i)
-        layer = base_layer_selector(
-            layer_type,
-            kernel_initializer=init,
-            units=nodes_out,
-            activation=activation,
-            input_shape=(nodes_in,),
-            **custom_args,
-        )
-        list_of_pdf_layers.append(layer)
-        nodes_in = int(nodes_out)
-
-    return list_of_pdf_layers
-
-
-def generate_dense_per_flavour_network(
-    nodes_in, nodes, activations, initializer_name="glorot_normal", seed=0, basis_size=8
-):
-    """
-    For each flavour generates a dense network of the chosen size
-
-    """
-    layer_type = "dense_per_flavour"
-    custom_args = {"basis_size": basis_size}
-
-    def initializer_generator(initializer_name, seed):
-        initializers = [
-            MetaLayer.select_initializer(initializer_name, seed=seed + b) for b in range(basis_size)
-        ]
-        return initializers
-
-    list_of_pdf_layers = []
-    for i, (nodes_out, activation) in enumerate(zip(nodes, activations)):
-        init = initializer_generator(initializer_name, seed + i)
-        layer = base_layer_selector(
-            layer_type,
-            kernel_initializer=init,
-            units=nodes_out,
-            activation=activation,
-            input_shape=(nodes_in,),
-            **custom_args,
-        )
-        list_of_pdf_layers.append(layer)
-        nodes_in = int(nodes_out)
-
-    return list_of_pdf_layers
-
-
 def generate_pdf_model(
     nodes: List[int] = None,
     activations: List[str] = None,
@@ -655,7 +587,7 @@ def pdfNN_layer_generator(
 
     nn_replicas = generate_nn(
         layer_type=layer_type,
-        input_dimensions=nn_input_dimensions,
+        nodes_in=nn_input_dimensions,
         nodes=nodes,
         activations=activations,
         initializer_name=initializer_name,
@@ -748,7 +680,7 @@ def compute_unnormalized_pdf(x, postfix=""):
 
 def generate_nn(
     layer_type: str,
-    input_dimensions: int,
+    nodes_in: int,
     nodes: List[int],
     activations: List[str],
     initializer_name: str,
@@ -762,8 +694,9 @@ def generate_nn(
     Create the part of the model that contains all of the actual neural network
     layers.
     """
-    x = Input(shape=(None, input_dimensions), batch_size=1, name='xgrids_processed')
+    x = Input(shape=(None, nodes_in), batch_size=1, name='xgrids_processed')
 
+    custom_args = {}
     if layer_type == "dense_per_flavour":
         # set the arguments that will define the layer
         # but careful, the last layer must be nodes = 1
@@ -771,24 +704,37 @@ def generate_nn(
         # is set to the number of nodes of the last layer when it should
         # come from the runcard
         nodes[-1] = 1
+        custom_args['basis_size'] = last_layer_nodes
 
-    constant_args = {
-        'nodes_in': input_dimensions,
-        'nodes': nodes,
-        'activations': activations,
-        'initializer_name': initializer_name,
-    }
-    if layer_type == "dense_per_flavour":
-        constant_args['basis_size'] = last_layer_nodes
-        layers_generator = generate_dense_per_flavour_network
-    if layer_type == "dense":
+        def initializer_generator(initializer_name, seed):
+            initializers = [
+                MetaLayer.select_initializer(initializer_name, seed=seed + b)
+                for b in range(basis_size)
+            ]
+            return initializers
+
+    elif layer_type == "dense":
         reg = regularizer_selector(regularizer, **regularizer_args)
-        constant_args['regularizer'] = reg
-        layers_generator = generate_dense_network
+        custom_args['regularizer'] = reg
+
+        def initializer_generator(initializer_name, seed):
+            return MetaLayer.select_initializer(initializer_name, seed=seed)
 
     models = []
     for i_replica, replica_seed in enumerate(replica_seeds):
-        list_of_pdf_layers = layers_generator(**constant_args, seed=replica_seed)
+        list_of_pdf_layers = []
+        for i_layer, (nodes_out, activation) in enumerate(zip(nodes, activations)):
+            init = initializer_generator(initializer_name, replica_seed + i_layer)
+            layer = base_layer_selector(
+                layer_type,
+                kernel_initializer=init,
+                units=nodes_out,
+                activation=activation,
+                input_shape=(nodes_in,),
+                **custom_args,
+            )
+            list_of_pdf_layers.append(layer)
+            nodes_in = int(nodes_out)
 
         # add dropout as second to last layer
         if dropout > 0:

From bdbc3c36440ad5b6d117bc3d752211e3be6fe4d7 Mon Sep 17 00:00:00 2001
From: Aron <aronpjansen@gmail.com>
Date: Fri, 8 Dec 2023 14:24:10 +0100
Subject: [PATCH 12/30] Add documentation.

---
 n3fit/src/n3fit/model_gen.py | 30 +++++++++++++++++++++++++++++-
 1 file changed, 29 insertions(+), 1 deletion(-)

diff --git a/n3fit/src/n3fit/model_gen.py b/n3fit/src/n3fit/model_gen.py
index 3b9cdd8b19..349f8bd26c 100644
--- a/n3fit/src/n3fit/model_gen.py
+++ b/n3fit/src/n3fit/model_gen.py
@@ -692,7 +692,35 @@ def generate_nn(
 ) -> MetaModel:
     """
     Create the part of the model that contains all of the actual neural network
-    layers.
+    layers, for each replica.
+
+    Parameters
+    ----------
+        layer_type: str
+            Type of layer to use. Can be "dense" or "dense_per_flavour".
+        nodes_in: int
+            Number of nodes in the input layer.
+        nodes: List[int]
+            Number of nodes in each hidden layer.
+        activations: List[str]
+            Activation function to use in each hidden layer.
+        initializer_name: str
+            Name of the initializer to use.
+        replica_seeds: List[int]
+            List of seeds to use for each replica.
+        dropout: float
+            Dropout rate to use (if 0, no dropout is used).
+        regularizer: str
+            Name of the regularizer to use.
+        regularizer_args: dict
+            Arguments to pass to the regularizer.
+        last_layer_nodes: int
+            Number of nodes in the last layer.
+
+    Returns
+    -------
+        nn_replicas: List[MetaModel]
+            List of MetaModel objects, one for each replica.
     """
     x = Input(shape=(None, nodes_in), batch_size=1, name='xgrids_processed')
 

From e3f9f0c0f8c242b4c0cc74b5d2d555d52bcbf1e2 Mon Sep 17 00:00:00 2001
From: Aron <aronpjansen@gmail.com>
Date: Fri, 8 Dec 2023 14:37:07 +0100
Subject: [PATCH 13/30] Simplify per_flavor layer concatenation

---
 n3fit/src/n3fit/model_gen.py | 10 +++-------
 1 file changed, 3 insertions(+), 7 deletions(-)

diff --git a/n3fit/src/n3fit/model_gen.py b/n3fit/src/n3fit/model_gen.py
index 349f8bd26c..9eee37101e 100644
--- a/n3fit/src/n3fit/model_gen.py
+++ b/n3fit/src/n3fit/model_gen.py
@@ -771,14 +771,10 @@ def initializer_generator(initializer_name, seed):
 
         # In case of per flavour network, concatenate at the last layer
         if layer_type == "dense_per_flavour":
-            last_layer = list_of_pdf_layers[-1]
             concat = base_layer_selector("concatenate")
-
-            def concatenated_last_layer(inputs):
-                result = last_layer(inputs)
-                return concat(result)
-
-            list_of_pdf_layers[-1] = concatenated_last_layer
+            list_of_pdf_layers[-1] = [
+                lambda x: concat(layer(x)) for layer in list_of_pdf_layers[-1]
+            ]
 
         # apply layers to create model
         pdf = x

From 3d9070f3fc24d3cc686643c93ffcc22702c55c2c Mon Sep 17 00:00:00 2001
From: Aron <aronpjansen@gmail.com>
Date: Fri, 8 Dec 2023 14:52:44 +0100
Subject: [PATCH 14/30] Reverse order of loops over replicas and layers

---
 n3fit/src/n3fit/model_gen.py | 60 +++++++++++++++++++++---------------
 1 file changed, 35 insertions(+), 25 deletions(-)

diff --git a/n3fit/src/n3fit/model_gen.py b/n3fit/src/n3fit/model_gen.py
index 9eee37101e..7c39bc3cb1 100644
--- a/n3fit/src/n3fit/model_gen.py
+++ b/n3fit/src/n3fit/model_gen.py
@@ -748,12 +748,16 @@ def initializer_generator(initializer_name, seed):
         def initializer_generator(initializer_name, seed):
             return MetaLayer.select_initializer(initializer_name, seed=seed)
 
-    models = []
-    for i_replica, replica_seed in enumerate(replica_seeds):
-        list_of_pdf_layers = []
-        for i_layer, (nodes_out, activation) in enumerate(zip(nodes, activations)):
-            init = initializer_generator(initializer_name, replica_seed + i_layer)
-            layer = base_layer_selector(
+    # First create all the layers...
+    # list_of_pdf_layers[d][r] is the layer at depth d for replica r
+    list_of_pdf_layers = []
+    for i_layer, (nodes_out, activation) in enumerate(zip(nodes, activations)):
+        inits = [
+            initializer_generator(initializer_name, replica_seed + i_layer)
+            for replica_seed in replica_seeds
+        ]
+        layers = [
+            base_layer_selector(
                 layer_type,
                 kernel_initializer=init,
                 units=nodes_out,
@@ -761,25 +765,31 @@ def initializer_generator(initializer_name, seed):
                 input_shape=(nodes_in,),
                 **custom_args,
             )
-            list_of_pdf_layers.append(layer)
-            nodes_in = int(nodes_out)
-
-        # add dropout as second to last layer
-        if dropout > 0:
-            dropout_layer = MetaLayer.base_layer_selector("dropout", rate=dropout)
-            list_of_pdf_layers.insert(dropout_layer, -2)
-
-        # In case of per flavour network, concatenate at the last layer
-        if layer_type == "dense_per_flavour":
-            concat = base_layer_selector("concatenate")
-            list_of_pdf_layers[-1] = [
-                lambda x: concat(layer(x)) for layer in list_of_pdf_layers[-1]
-            ]
+            for init in inits
+        ]
+        list_of_pdf_layers.append(layers)
+        nodes_in = int(nodes_out)
+
+    # add dropout as second to last layer
+    if dropout > 0:
+        dropout_layer = MetaLayer.base_layer_selector("dropout", rate=dropout)
+        list_of_pdf_layers.insert(dropout_layer, -2)
+
+    # In case of per flavour network, concatenate at the last layer
+    if layer_type == "dense_per_flavour":
+        concat = base_layer_selector("concatenate")
+        list_of_pdf_layers[-1] = [lambda x: concat(layer(x)) for layer in list_of_pdf_layers[-1]]
+
+    # ... then apply them to the input to create the models
+    xs = [layer(x) for layer in list_of_pdf_layers[0]]
+    for layers in list_of_pdf_layers[1:]:
+        if type(layers) is list:
+            xs = [layer(x) for layer, x in zip(layers, xs)]
+        else:
+            xs = [layers(x) for x in xs]
 
-        # apply layers to create model
-        pdf = x
-        for layer in list_of_pdf_layers:
-            pdf = layer(pdf)
-        models.append(MetaModel({'NN_input': x}, pdf, name=f"NN_{i_replica}"))
+    models = [
+        MetaModel({'NN_input': x}, pdf, name=f"NN_{i_replica}") for i_replica, pdf in enumerate(xs)
+    ]
 
     return models

From c8300c872299a464ed0cf023ed92f2932981dcc4 Mon Sep 17 00:00:00 2001
From: Aron <aronpjansen@gmail.com>
Date: Fri, 8 Dec 2023 15:10:25 +0100
Subject: [PATCH 15/30] Fixes for dropout

---
 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 7c39bc3cb1..12b33e2738 100644
--- a/n3fit/src/n3fit/model_gen.py
+++ b/n3fit/src/n3fit/model_gen.py
@@ -772,8 +772,8 @@ def initializer_generator(initializer_name, seed):
 
     # add dropout as second to last layer
     if dropout > 0:
-        dropout_layer = MetaLayer.base_layer_selector("dropout", rate=dropout)
-        list_of_pdf_layers.insert(dropout_layer, -2)
+        dropout_layer = base_layer_selector("dropout", rate=dropout)
+        list_of_pdf_layers.insert(-2, dropout_layer)
 
     # In case of per flavour network, concatenate at the last layer
     if layer_type == "dense_per_flavour":

From 0cf23f294257f485183b8d054993e1b8786dff28 Mon Sep 17 00:00:00 2001
From: Aron <aronpjansen@gmail.com>
Date: Fri, 8 Dec 2023 15:26:59 +0100
Subject: [PATCH 16/30] Fixes for per_flavour

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

diff --git a/n3fit/src/n3fit/model_gen.py b/n3fit/src/n3fit/model_gen.py
index 12b33e2738..c4c9319ab0 100644
--- a/n3fit/src/n3fit/model_gen.py
+++ b/n3fit/src/n3fit/model_gen.py
@@ -732,7 +732,7 @@ def generate_nn(
         # is set to the number of nodes of the last layer when it should
         # come from the runcard
         nodes[-1] = 1
-        custom_args['basis_size'] = last_layer_nodes
+        basis_size = last_layer_nodes
 
         def initializer_generator(initializer_name, seed):
             initializers = [

From b0a8e3b2b14f6cb6dcb65234784e8bc81e2cb429 Mon Sep 17 00:00:00 2001
From: Aron <aronpjansen@gmail.com>
Date: Mon, 11 Dec 2023 07:50:08 +0100
Subject: [PATCH 17/30] Fix issue with copying over nodes for per_flavour layer

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

diff --git a/n3fit/src/n3fit/model_gen.py b/n3fit/src/n3fit/model_gen.py
index c4c9319ab0..3845fd7b22 100644
--- a/n3fit/src/n3fit/model_gen.py
+++ b/n3fit/src/n3fit/model_gen.py
@@ -722,6 +722,7 @@ def generate_nn(
         nn_replicas: List[MetaModel]
             List of MetaModel objects, one for each replica.
     """
+    nodes_local = nodes.copy()  # make a local copy so we can modify it
     x = Input(shape=(None, nodes_in), batch_size=1, name='xgrids_processed')
 
     custom_args = {}
@@ -731,7 +732,7 @@ def generate_nn(
         # TODO the mismatch is due to the fact that basis_size
         # is set to the number of nodes of the last layer when it should
         # come from the runcard
-        nodes[-1] = 1
+        nodes_local[-1] = 1
         basis_size = last_layer_nodes
 
         def initializer_generator(initializer_name, seed):
@@ -751,7 +752,7 @@ def initializer_generator(initializer_name, seed):
     # First create all the layers...
     # list_of_pdf_layers[d][r] is the layer at depth d for replica r
     list_of_pdf_layers = []
-    for i_layer, (nodes_out, activation) in enumerate(zip(nodes, activations)):
+    for i_layer, (nodes_out, activation) in enumerate(zip(nodes_local, activations)):
         inits = [
             initializer_generator(initializer_name, replica_seed + i_layer)
             for replica_seed in replica_seeds

From 97d2efef66cdb3c63192a140ea1a05018099de67 Mon Sep 17 00:00:00 2001
From: Aron <aronpjansen@gmail.com>
Date: Mon, 11 Dec 2023 08:08:04 +0100
Subject: [PATCH 18/30] Fix seeds in per_flavour layer

---
 n3fit/src/n3fit/model_gen.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/n3fit/src/n3fit/model_gen.py b/n3fit/src/n3fit/model_gen.py
index 3845fd7b22..feab58190e 100644
--- a/n3fit/src/n3fit/model_gen.py
+++ b/n3fit/src/n3fit/model_gen.py
@@ -735,7 +735,8 @@ def generate_nn(
         nodes_local[-1] = 1
         basis_size = last_layer_nodes
 
-        def initializer_generator(initializer_name, seed):
+        def initializer_generator(initializer_name, seed, i_layer):
+            seed += i_layer * basis_size
             initializers = [
                 MetaLayer.select_initializer(initializer_name, seed=seed + b)
                 for b in range(basis_size)
@@ -746,7 +747,8 @@ def initializer_generator(initializer_name, seed):
         reg = regularizer_selector(regularizer, **regularizer_args)
         custom_args['regularizer'] = reg
 
-        def initializer_generator(initializer_name, seed):
+        def initializer_generator(initializer_name, seed, i_layer):
+            seed += i_layer
             return MetaLayer.select_initializer(initializer_name, seed=seed)
 
     # First create all the layers...
@@ -754,7 +756,7 @@ def initializer_generator(initializer_name, seed):
     list_of_pdf_layers = []
     for i_layer, (nodes_out, activation) in enumerate(zip(nodes_local, activations)):
         inits = [
-            initializer_generator(initializer_name, replica_seed + i_layer)
+            initializer_generator(initializer_name, replica_seed, i_layer)
             for replica_seed in replica_seeds
         ]
         layers = [

From 4c4a2d5757209f14e08e685d3120770185054f8c Mon Sep 17 00:00:00 2001
From: Aron <aronpjansen@gmail.com>
Date: Mon, 11 Dec 2023 08:09:31 +0100
Subject: [PATCH 19/30] Add error for combination of dropout with per_flavour
 layers

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

diff --git a/n3fit/src/n3fit/model_gen.py b/n3fit/src/n3fit/model_gen.py
index feab58190e..427274e6e6 100644
--- a/n3fit/src/n3fit/model_gen.py
+++ b/n3fit/src/n3fit/model_gen.py
@@ -722,6 +722,9 @@ def generate_nn(
         nn_replicas: List[MetaModel]
             List of MetaModel objects, one for each replica.
     """
+    if dropout > 0 and layer_type == "dense_per_flavour":
+        raise ValueError("Dropout is not supported for dense_per_flavour layers")
+
     nodes_local = nodes.copy()  # make a local copy so we can modify it
     x = Input(shape=(None, nodes_in), batch_size=1, name='xgrids_processed')
 

From 2287194a27c32628a0b4d368fb693b220ce98a44 Mon Sep 17 00:00:00 2001
From: Aron <aronpjansen@gmail.com>
Date: Mon, 11 Dec 2023 08:50:00 +0100
Subject: [PATCH 20/30] Add basis_size argument to per_flavour layer

---
 n3fit/src/n3fit/model_gen.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/n3fit/src/n3fit/model_gen.py b/n3fit/src/n3fit/model_gen.py
index 427274e6e6..8b41f03959 100644
--- a/n3fit/src/n3fit/model_gen.py
+++ b/n3fit/src/n3fit/model_gen.py
@@ -737,6 +737,7 @@ def generate_nn(
         # come from the runcard
         nodes_local[-1] = 1
         basis_size = last_layer_nodes
+        custom_args['basis_size'] = basis_size
 
         def initializer_generator(initializer_name, seed, i_layer):
             seed += i_layer * basis_size

From 2f68e3d5d2d825e8ecf231949072f128bd98cbe8 Mon Sep 17 00:00:00 2001
From: Aron <aronpjansen@gmail.com>
Date: Mon, 11 Dec 2023 08:51:19 +0100
Subject: [PATCH 21/30] Fix model_gen tests to use new generate_nn in favor of
 now removed generate_dense and generate_dense_per_flavour

---
 n3fit/src/n3fit/tests/test_modelgen.py | 60 ++++++++++----------------
 1 file changed, 23 insertions(+), 37 deletions(-)

diff --git a/n3fit/src/n3fit/tests/test_modelgen.py b/n3fit/src/n3fit/tests/test_modelgen.py
index 4cec4574ee..f76e022810 100644
--- a/n3fit/src/n3fit/tests/test_modelgen.py
+++ b/n3fit/src/n3fit/tests/test_modelgen.py
@@ -5,62 +5,48 @@
     It checks that both the number of layers and the shape
     of the weights of the layers are what is expected
 """
-import numpy as np
-import n3fit.model_gen
-from n3fit.backends import MetaModel
-from n3fit.backends import operations as op
+from n3fit.model_gen import generate_nn
 
 INSIZE = 16
-OUT_SIZES = (4, 3)
+OUT_SIZES = [4, 3]
 BASIS_SIZE = 3
 
+COMMON_ARGS = {
+    "nodes_in": INSIZE,
+    "nodes": OUT_SIZES,
+    "activations": ["sigmoid", "tanh"],
+    "initializer_name": "glorot_uniform",
+    "replica_seeds": [0],
+    "dropout": 0.0,
+    "regularizer": None,
+    "regularizer_args": {},
+    "last_layer_nodes": BASIS_SIZE,
+}
+
 
 def test_generate_dense_network():
-    nodes_in = INSIZE
-    nodes_out = OUT_SIZES
-    activations = ["sigmoid", "tanh"]
-    layers = n3fit.model_gen.generate_dense_network(nodes_in, nodes_out, activations)
-    arr = np.random.rand(1, INSIZE)
-    input_layer = op.numpy_to_input(arr)
-    curr_layer = input_layer
-    for layer in layers:
-        curr_layer = layer(curr_layer)
-    modelito = MetaModel({"input": input_layer}, curr_layer)
+    nn = generate_nn("dense", **COMMON_ARGS)[0]
+
     # The number of layers should be input layer + len(OUT_SIZES)
-    assert len(modelito.layers) == len(OUT_SIZES) + 1
+    assert len(nn.layers) == len(OUT_SIZES) + 1
     # Check that the number of parameters is as expected
     # We expect 4 weights where the two first ones are
     # (INSIZE, OUT_SIZE[0]) (OUT_SIZE[0],)
     # and the second one
     # (OUT_SIZE[0], OUT_SIZE[1]) (OUT_SIZE[1],)
-    expected_sizes = [
-        (INSIZE, OUT_SIZES[0]),
-        (OUT_SIZES[0],),
-        OUT_SIZES,
-        (OUT_SIZES[1],),
-    ]
-    for weight, esize in zip(modelito.weights, expected_sizes):
+    expected_sizes = [(INSIZE, OUT_SIZES[0]), (OUT_SIZES[0],), OUT_SIZES, (OUT_SIZES[1],)]
+    for weight, esize in zip(nn.weights, expected_sizes):
         assert weight.shape == esize
 
 
 def test_generate_dense_per_flavour_network():
-    nodes_in = INSIZE
-    nodes_out = OUT_SIZES
-    activations = ["sigmoid", "tanh"]
-    layers = n3fit.model_gen.generate_dense_per_flavour_network(
-        nodes_in, nodes_out, activations, basis_size=BASIS_SIZE
-    )
-    arr = np.random.rand(1, INSIZE)
-    input_layer = op.numpy_to_input(arr)
-    curr_layer = input_layer
-    for layer in layers:
-        curr_layer = layer(curr_layer)
-    modelito = MetaModel({"input": input_layer}, curr_layer)
+    nn = generate_nn("dense_per_flavour", **COMMON_ARGS)[0]
+
     # The number of layers should be input + BASIS_SIZE*len(OUT_SIZES) + concatenate
-    assert len(modelito.layers) == BASIS_SIZE * len(OUT_SIZES) + 2
+    assert len(nn.layers) == BASIS_SIZE * len(OUT_SIZES) + 2
     # The shape for this network of denses for flavours will depend on the basis_size
     expected_sizes = []
     expected_sizes += BASIS_SIZE * [(INSIZE, OUT_SIZES[0]), (OUT_SIZES[0],)]
     expected_sizes += BASIS_SIZE * [(OUT_SIZES[0], 1), (1,)]
-    for weight, esize in zip(modelito.weights, expected_sizes):
+    for weight, esize in zip(nn.weights, expected_sizes):
         assert weight.shape == esize

From 4dd1649917380c7a1a56057c83a8e535724470ea Mon Sep 17 00:00:00 2001
From: Aron <aronpjansen@gmail.com>
Date: Mon, 11 Dec 2023 10:49:48 +0100
Subject: [PATCH 22/30] Allow for nodes to be a tuple

---
 n3fit/src/n3fit/model_gen.py           | 6 +++---
 n3fit/src/n3fit/tests/test_modelgen.py | 2 +-
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/n3fit/src/n3fit/model_gen.py b/n3fit/src/n3fit/model_gen.py
index 8b41f03959..2cd1ebea1c 100644
--- a/n3fit/src/n3fit/model_gen.py
+++ b/n3fit/src/n3fit/model_gen.py
@@ -725,7 +725,7 @@ def generate_nn(
     if dropout > 0 and layer_type == "dense_per_flavour":
         raise ValueError("Dropout is not supported for dense_per_flavour layers")
 
-    nodes_local = nodes.copy()  # make a local copy so we can modify it
+    nodes_list = list(nodes)  # so we can modify it
     x = Input(shape=(None, nodes_in), batch_size=1, name='xgrids_processed')
 
     custom_args = {}
@@ -735,7 +735,7 @@ def generate_nn(
         # TODO the mismatch is due to the fact that basis_size
         # is set to the number of nodes of the last layer when it should
         # come from the runcard
-        nodes_local[-1] = 1
+        nodes_list[-1] = 1
         basis_size = last_layer_nodes
         custom_args['basis_size'] = basis_size
 
@@ -758,7 +758,7 @@ def initializer_generator(initializer_name, seed, i_layer):
     # First create all the layers...
     # list_of_pdf_layers[d][r] is the layer at depth d for replica r
     list_of_pdf_layers = []
-    for i_layer, (nodes_out, activation) in enumerate(zip(nodes_local, activations)):
+    for i_layer, (nodes_out, activation) in enumerate(zip(nodes_list, activations)):
         inits = [
             initializer_generator(initializer_name, replica_seed, i_layer)
             for replica_seed in replica_seeds
diff --git a/n3fit/src/n3fit/tests/test_modelgen.py b/n3fit/src/n3fit/tests/test_modelgen.py
index f76e022810..d50b98728a 100644
--- a/n3fit/src/n3fit/tests/test_modelgen.py
+++ b/n3fit/src/n3fit/tests/test_modelgen.py
@@ -8,7 +8,7 @@
 from n3fit.model_gen import generate_nn
 
 INSIZE = 16
-OUT_SIZES = [4, 3]
+OUT_SIZES = (4, 3)
 BASIS_SIZE = 3
 
 COMMON_ARGS = {

From 6bd6466a8d0d37bbcbcc39f4dbea5cc9aa871136 Mon Sep 17 00:00:00 2001
From: Aron <aronpjansen@gmail.com>
Date: Mon, 11 Dec 2023 15:39:23 +0100
Subject: [PATCH 23/30] Move dropout, per_flavour check to checks

---
 n3fit/src/n3fit/checks.py    | 7 +++++++
 n3fit/src/n3fit/model_gen.py | 3 ---
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/n3fit/src/n3fit/checks.py b/n3fit/src/n3fit/checks.py
index 7b509f0065..31b1fae2cb 100644
--- a/n3fit/src/n3fit/checks.py
+++ b/n3fit/src/n3fit/checks.py
@@ -114,6 +114,13 @@ def check_dropout(parameters):
     if dropout is not None and not 0.0 <= dropout <= 1.0:
         raise CheckError(f"Dropout must be between 0 and 1, got: {dropout}")
 
+    layer_type = parameters.get("layer_type")
+    if dropout is not None and dropout > 0.0 and layer_type == "dense_per_flavour":
+        raise CheckError(
+            "Dropout is not compatible with the dense_per_flavour layer type, "
+            "please use instead the dense layer type"
+        )
+
 
 def check_tensorboard(tensorboard):
     """Check that the tensorbard callback can be enabled correctly"""
diff --git a/n3fit/src/n3fit/model_gen.py b/n3fit/src/n3fit/model_gen.py
index 2cd1ebea1c..064ddae271 100644
--- a/n3fit/src/n3fit/model_gen.py
+++ b/n3fit/src/n3fit/model_gen.py
@@ -722,9 +722,6 @@ def generate_nn(
         nn_replicas: List[MetaModel]
             List of MetaModel objects, one for each replica.
     """
-    if dropout > 0 and layer_type == "dense_per_flavour":
-        raise ValueError("Dropout is not supported for dense_per_flavour layers")
-
     nodes_list = list(nodes)  # so we can modify it
     x = Input(shape=(None, nodes_in), batch_size=1, name='xgrids_processed')
 

From 2cd9e52f555a147aba21c8daee01f76a09e0e676 Mon Sep 17 00:00:00 2001
From: Aron Jansen <APJansen@users.noreply.github.com>
Date: Thu, 14 Dec 2023 10:12:22 +0100
Subject: [PATCH 24/30] Clarify layer type check

Co-authored-by: Juan M. Cruz-Martinez <juacrumar@lairen.eu>
---
 n3fit/src/n3fit/checks.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/n3fit/src/n3fit/checks.py b/n3fit/src/n3fit/checks.py
index 31b1fae2cb..f33213d5a0 100644
--- a/n3fit/src/n3fit/checks.py
+++ b/n3fit/src/n3fit/checks.py
@@ -118,7 +118,7 @@ def check_dropout(parameters):
     if dropout is not None and dropout > 0.0 and layer_type == "dense_per_flavour":
         raise CheckError(
             "Dropout is not compatible with the dense_per_flavour layer type, "
-            "please use instead the dense layer type"
+            "please use instead `parameters::layer_type: dense`"
         )
 
 

From 1ae1b84a6b95b10bff0b80c49fd6738d8ed1bb45 Mon Sep 17 00:00:00 2001
From: Aron <aronpjansen@gmail.com>
Date: Thu, 14 Dec 2023 10:18:34 +0100
Subject: [PATCH 25/30] Clarify naming in nn_generator

---
 n3fit/src/n3fit/model_gen.py | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/n3fit/src/n3fit/model_gen.py b/n3fit/src/n3fit/model_gen.py
index 064ddae271..7e8a2d406a 100644
--- a/n3fit/src/n3fit/model_gen.py
+++ b/n3fit/src/n3fit/model_gen.py
@@ -723,7 +723,7 @@ def generate_nn(
             List of MetaModel objects, one for each replica.
     """
     nodes_list = list(nodes)  # so we can modify it
-    x = Input(shape=(None, nodes_in), batch_size=1, name='xgrids_processed')
+    x_input = Input(shape=(None, nodes_in), batch_size=1, name='xgrids_processed')
 
     custom_args = {}
     if layer_type == "dense_per_flavour":
@@ -785,15 +785,16 @@ def initializer_generator(initializer_name, seed, i_layer):
         list_of_pdf_layers[-1] = [lambda x: concat(layer(x)) for layer in list_of_pdf_layers[-1]]
 
     # ... then apply them to the input to create the models
-    xs = [layer(x) for layer in list_of_pdf_layers[0]]
+    pdfs = [layer(x_input) for layer in list_of_pdf_layers[0]]
     for layers in list_of_pdf_layers[1:]:
         if type(layers) is list:
-            xs = [layer(x) for layer, x in zip(layers, xs)]
+            pdfs = [layer(x) for layer, x in zip(layers, pdfs)]
         else:
-            xs = [layers(x) for x in xs]
+            pdfs = [layers(x) for x in pdfs]
 
     models = [
-        MetaModel({'NN_input': x}, pdf, name=f"NN_{i_replica}") for i_replica, pdf in enumerate(xs)
+        MetaModel({'NN_input': x_input}, pdf, name=f"NN_{i_replica}")
+        for i_replica, pdf in enumerate(pdfs)
     ]
 
     return models

From e7a7cb411d0185aee5c120e67c89ba6743fedcee Mon Sep 17 00:00:00 2001
From: Aron <aronpjansen@gmail.com>
Date: Thu, 14 Dec 2023 10:24:28 +0100
Subject: [PATCH 26/30] Remove initializer_name argument

---
 n3fit/src/n3fit/model_gen.py | 9 +++------
 1 file changed, 3 insertions(+), 6 deletions(-)

diff --git a/n3fit/src/n3fit/model_gen.py b/n3fit/src/n3fit/model_gen.py
index 7e8a2d406a..0aecdd00bd 100644
--- a/n3fit/src/n3fit/model_gen.py
+++ b/n3fit/src/n3fit/model_gen.py
@@ -736,7 +736,7 @@ def generate_nn(
         basis_size = last_layer_nodes
         custom_args['basis_size'] = basis_size
 
-        def initializer_generator(initializer_name, seed, i_layer):
+        def initializer_generator(seed, i_layer):
             seed += i_layer * basis_size
             initializers = [
                 MetaLayer.select_initializer(initializer_name, seed=seed + b)
@@ -748,7 +748,7 @@ def initializer_generator(initializer_name, seed, i_layer):
         reg = regularizer_selector(regularizer, **regularizer_args)
         custom_args['regularizer'] = reg
 
-        def initializer_generator(initializer_name, seed, i_layer):
+        def initializer_generator(seed, i_layer):
             seed += i_layer
             return MetaLayer.select_initializer(initializer_name, seed=seed)
 
@@ -756,10 +756,7 @@ def initializer_generator(initializer_name, seed, i_layer):
     # list_of_pdf_layers[d][r] is the layer at depth d for replica r
     list_of_pdf_layers = []
     for i_layer, (nodes_out, activation) in enumerate(zip(nodes_list, activations)):
-        inits = [
-            initializer_generator(initializer_name, replica_seed, i_layer)
-            for replica_seed in replica_seeds
-        ]
+        inits = [initializer_generator(replica_seed, i_layer) for replica_seed in replica_seeds]
         layers = [
             base_layer_selector(
                 layer_type,

From 07c1e7d9eaa63a122e115d0247de0ccd16434f98 Mon Sep 17 00:00:00 2001
From: Aron Jansen <APJansen@users.noreply.github.com>
Date: Thu, 14 Dec 2023 10:27:00 +0100
Subject: [PATCH 27/30] clarify comment

Co-authored-by: Juan M. Cruz-Martinez <juacrumar@lairen.eu>
---
 n3fit/src/n3fit/model_gen.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/n3fit/src/n3fit/model_gen.py b/n3fit/src/n3fit/model_gen.py
index 0aecdd00bd..747a46a2cb 100644
--- a/n3fit/src/n3fit/model_gen.py
+++ b/n3fit/src/n3fit/model_gen.py
@@ -781,7 +781,7 @@ def initializer_generator(seed, i_layer):
         concat = base_layer_selector("concatenate")
         list_of_pdf_layers[-1] = [lambda x: concat(layer(x)) for layer in list_of_pdf_layers[-1]]
 
-    # ... then apply them to the input to create the models
+    # Apply all layers to the input to create the models
     pdfs = [layer(x_input) for layer in list_of_pdf_layers[0]]
     for layers in list_of_pdf_layers[1:]:
         if type(layers) is list:

From 25b8308cb660dcb96d543021dc88f4256d7796f4 Mon Sep 17 00:00:00 2001
From: Aron <aronpjansen@gmail.com>
Date: Thu, 14 Dec 2023 10:29:49 +0100
Subject: [PATCH 28/30] Add comment on shared layers

---
 n3fit/src/n3fit/model_gen.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/n3fit/src/n3fit/model_gen.py b/n3fit/src/n3fit/model_gen.py
index 747a46a2cb..0c0e92a695 100644
--- a/n3fit/src/n3fit/model_gen.py
+++ b/n3fit/src/n3fit/model_gen.py
@@ -784,6 +784,7 @@ def initializer_generator(seed, i_layer):
     # Apply all layers to the input to create the models
     pdfs = [layer(x_input) for layer in list_of_pdf_layers[0]]
     for layers in list_of_pdf_layers[1:]:
+        # Since some layers (dropout) are shared, we have to treat them separately
         if type(layers) is list:
             pdfs = [layer(x) for layer, x in zip(layers, pdfs)]
         else:

From 692014b55082e8198e0d860cc52db3e397d44612 Mon Sep 17 00:00:00 2001
From: Aron <aronpjansen@gmail.com>
Date: Thu, 14 Dec 2023 10:45:05 +0100
Subject: [PATCH 29/30] Rewrite comprehension over replica seeds

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

diff --git a/n3fit/src/n3fit/model_gen.py b/n3fit/src/n3fit/model_gen.py
index 0c0e92a695..09f337dbca 100644
--- a/n3fit/src/n3fit/model_gen.py
+++ b/n3fit/src/n3fit/model_gen.py
@@ -756,17 +756,16 @@ def initializer_generator(seed, i_layer):
     # list_of_pdf_layers[d][r] is the layer at depth d for replica r
     list_of_pdf_layers = []
     for i_layer, (nodes_out, activation) in enumerate(zip(nodes_list, activations)):
-        inits = [initializer_generator(replica_seed, i_layer) for replica_seed in replica_seeds]
         layers = [
             base_layer_selector(
                 layer_type,
-                kernel_initializer=init,
+                kernel_initializer=initializer_generator(replica_seed, i_layer),
                 units=nodes_out,
                 activation=activation,
                 input_shape=(nodes_in,),
                 **custom_args,
             )
-            for init in inits
+            for replica_seed in replica_seeds
         ]
         list_of_pdf_layers.append(layers)
         nodes_in = int(nodes_out)

From 903c75b7024a926dc2d6e4c316ab7f6dcb4b164b Mon Sep 17 00:00:00 2001
From: Aron <aronpjansen@gmail.com>
Date: Fri, 15 Dec 2023 09:44:18 +0100
Subject: [PATCH 30/30] Add check on layer type

---
 n3fit/src/n3fit/checks.py | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/n3fit/src/n3fit/checks.py b/n3fit/src/n3fit/checks.py
index f33213d5a0..885785a268 100644
--- a/n3fit/src/n3fit/checks.py
+++ b/n3fit/src/n3fit/checks.py
@@ -108,6 +108,16 @@ def check_initializer(initializer):
         raise CheckError(f"Initializer {initializer} not accepted by {MetaLayer}")
 
 
+def check_layer_type_implemented(parameters):
+    """Checks whether the layer_type is implemented"""
+    layer_type = parameters.get("layer_type")
+    implemented_types = ["dense", "dense_per_flavour"]
+    if layer_type not in implemented_types:
+        raise CheckError(
+            f"Layer type {layer_type} not implemented, must be one of {implemented_types}"
+        )
+
+
 def check_dropout(parameters):
     """Checks the dropout setup (positive and smaller than 1.0)"""
     dropout = parameters.get("dropout")
@@ -175,6 +185,7 @@ def wrapper_check_NN(basis, tensorboard, save, load, parameters):
     check_consistent_layers(parameters)
     check_basis_with_layers(basis, parameters)
     check_stopping(parameters)
+    check_layer_type_implemented(parameters)
     check_dropout(parameters)
     check_lagrange_multipliers(parameters, "integrability")
     check_lagrange_multipliers(parameters, "positivity")