Add shape information to try to fix bug with python 3.11

NNPDF · Jan 10, 2024 · d77af4d · d77af4d
1 parent 997451b
commit d77af4d
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/n3fit/src/n3fit/layers/x_operations.py b/n3fit/src/n3fit/layers/x_operations.py
@@ -76,3 +76,6 @@ def __init__(self, grid_weights, x_axis=2, **kwargs):
 
     def call(self, pdf):
         return op.tensor_product(pdf, self.grid_weights, axes=[self.x_axis, 0])
+
+    def compute_output_shape(self, input_shape):
+        return input_shape[: self.x_axis] + input_shape[self.x_axis + 1 :]
diff --git a/n3fit/src/n3fit/msr.py b/n3fit/src/n3fit/msr.py
@@ -85,7 +85,7 @@ def generate_msr_model_and_grid(
     )([x_divided, pdf_xgrid_integration])
 
     # 4. Integrate the pdf
-    pdf_integrated = xIntegrator(weights_array, input_shape=(nx,))(pdf_integrand)
+    pdf_integrated = xIntegrator(weights_array, input_shape=(1, replicas, nx))(pdf_integrand)
 
     # 5. THe input for the photon integral, will be set to 0 if no photons
     photon_integral = Input(shape=(replicas, 1), batch_size=1, name='photon_integral')