From 97f2515ae2b20fd5a679da6eebac16c569074ee1 Mon Sep 17 00:00:00 2001
From: juacrumar <juacrumar@lairen.eu>
Date: Fri, 22 Nov 2024 19:52:53 +0100
Subject: [PATCH] fix for normalized distributions

---
 n3fit/src/n3fit/layers/DY.py         | 8 +++++---
 n3fit/src/n3fit/layers/observable.py | 7 +++++--
 n3fit/src/n3fit/model_gen.py         | 1 -
 3 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/n3fit/src/n3fit/layers/DY.py b/n3fit/src/n3fit/layers/DY.py
index f05416c5e4..94c982a391 100644
--- a/n3fit/src/n3fit/layers/DY.py
+++ b/n3fit/src/n3fit/layers/DY.py
@@ -86,7 +86,7 @@ def compute_dy_observable_many_replica(pdf, padded_fk):
     """
     pdfa = pdf[1]
     pdfb = pdf[0]
-    
+
     temp = op.einsum('nxfyg, bryg -> brnxf', padded_fk, pdfa)
     return op.einsum('brnxf, brxf -> brn', temp, pdfb)
 
@@ -96,11 +96,13 @@ def compute_dy_observable_one_replica(pdf, mask_and_fk):
     Same operations as above but a specialized implementation that is more efficient for 1 replica,
     masking the PDF rather than the fk table.
     """
+    # mask: (channels, flavs_b, flavs_a) Ffg
+    # fk: (npoints, channels, x_a, x_b) nFyx
     mask, fk = mask_and_fk
     # Retrieve the two PDFs (which may potentially be coming from different initial states)
     # Since this is the one-replica function, remove the batch and replica dimension
-    pdfb = pdf[0][0][0]  # xf
-    pdfa = pdf[1][0][0]  # yg
+    pdfb = pdf[0][0][0]  # (x_b, flavs_b) xf
+    pdfa = pdf[1][0][0]  # (x_a, flavs_a) yg
 
     # TODO: check which PDF must go first in case of different initial states!!!
     mask_x_pdf = op.tensor_product(mask, pdfa, axes=[(2,), (1,)])  # Ffg, yg -> Ffy
diff --git a/n3fit/src/n3fit/layers/observable.py b/n3fit/src/n3fit/layers/observable.py
index f06fa6e3ee..29c5754ecc 100644
--- a/n3fit/src/n3fit/layers/observable.py
+++ b/n3fit/src/n3fit/layers/observable.py
@@ -89,7 +89,7 @@ def __init__(
         operation_name="NULL",
         nfl=14,
         n_replicas=1,
-        **kwargs
+        **kwargs,
     ):
         super(MetaLayer, self).__init__(**kwargs)
 
@@ -178,7 +178,10 @@ def call(self, pdf):
                 rank 3 tensor (batchsize, replicas, ndata)
         """
         if self.splitting:
-            pdfs = op.split(pdf, self.splitting, axis=2)
+            splitter = op.tensor_splitter(
+                pdf.shape, self.splitting, axis=2, name=f"pdf_splitter_{self.name}"
+            )
+            pdfs = splitter(pdf)
         else:
             pdfs = [pdf] * len(self.padded_fk_tables)
 
diff --git a/n3fit/src/n3fit/model_gen.py b/n3fit/src/n3fit/model_gen.py
index 5c7113bdf7..f7fe7c5608 100644
--- a/n3fit/src/n3fit/model_gen.py
+++ b/n3fit/src/n3fit/model_gen.py
@@ -99,7 +99,6 @@ def _generate_experimental_layer(self, pdf):
         the input PDF is evaluated in all points that the experiment needs and needs to be split
         """
         if len(self.dataset_xsizes) > 1:
-
             sp_layer = op.tensor_splitter(
                 pdf.shape, self.dataset_xsizes, axis=2, name=f"{self.name}_split"
             )