allow any kind of initialiser for the generation of the prepro exponents

NNPDF · Dec 13, 2024 · 0109ff9 · 0109ff9
1 parent 9aaad26
commit 0109ff9
Showing 1 changed file with 6 additions and 6 deletions.
diff --git a/n3fit/src/n3fit/layers/preprocessing.py b/n3fit/src/n3fit/layers/preprocessing.py
@@ -76,15 +76,15 @@ def generate_weight(self, name: str, kind: str, dictionary: dict, set_to_zero: b
             single_replica_initializer = MetaLayer.init_constant(0.0)
             trainable = False
         else:
-            minval, maxval = dictionary[kind]
+            prepro_kwargs = dictionary[kind]
             trainable = dictionary.get("trainable", True)
             # Seeds will be set in the wrapper MultiInitializer
-            single_replica_initializer = MetaLayer.select_initializer(
-                "random_uniform", minval=minval, maxval=maxval
-            )
+            single_replica_initializer = MetaLayer.select_initializer(**prepro_kwargs)
             # If we are training, constrain the weights to be within the limits
-            if trainable:
-                constraint = constraints.MinMaxWeight(minval, maxval)
+            if trainable and prepro_kwargs["ini_name"] == "random_uniform":
+                constraint = constraints.MinMaxWeight(
+                    prepro_kwargs["minval"], prepro_kwargs["maxval"]
+                )
 
         initializer = MultiInitializer(single_replica_initializer, self._replica_seeds, base_seed=0)
         # increment seeds for the next coefficient