add a test for pytorch

.github/workflows/pytorch_test.yml

@@ -0,0 +1,28 @@
+name: Test the python installation
+
+on: [push]
+
+jobs:
+  build:
+    strategy:
+      matrix:
+        os: [ubuntu-latest]
+        python-version: ["3.12"]
+        include:
+          - os: ubuntu-latest
+      fail-fast: false
+    runs-on: ${{ matrix.os }}
+    env:
+      NETRC_FILE: ${{ secrets.NETRC_FILE }}
+      NNPDF_SSH_KEY: ${{ secrets.NNPDF_SSH_KEY }}
+    steps:
+    - uses: actions/checkout@v4
+    - name: Install nnpdf without LHAPDF
+      shell: bash -l {0}
+      run: |
+        pip install .[nolha,torch]
+    - name: Test we can run one runcard
+      shell: bash -l {0}
+      run: |
+        cd n3fit/runcards/examples
+        n3fit Basic_runcard.yml 4

n3fit/src/n3fit/backends/keras_backend/MetaModel.py

@@ -8,12 +8,11 @@
 from pathlib import Path
 import re
 
-from keras import backend as K
+from keras import Variable
 from keras import ops as Kops
 from keras import optimizers as Kopt
 from keras.models import Model
 import numpy as np
-import tensorflow as tf
 
 import n3fit.backends.keras_backend.operations as op
 
@@ -501,9 +500,9 @@ def get_layer_replica_weights(layer, i_replica: int):
     """
     if is_stacked_single_replicas(layer):
         weights_ref = layer.get_layer(f"{NN_PREFIX}_{i_replica}").weights
-        weights = [tf.Variable(w, name=w.name) for w in weights_ref]
+        weights = [Variable(w, name=w.name) for w in weights_ref]
     else:
-        weights = [tf.Variable(w[i_replica : i_replica + 1], name=w.name) for w in layer.weights]
+        weights = [Variable(w[i_replica : i_replica + 1], name=w.name) for w in layer.weights]
 
     return weights
 

n3fit/src/n3fit/backends/keras_backend/callbacks.py

@@ -15,9 +15,11 @@
 import logging
 from time import time
 
+# Callbacks need tensorflow installed even if the backend is pytorch
 from keras.callbacks import Callback, TensorBoard
 import numpy as np
-import tensorflow as tf
+
+from .operations import decorator_compiler
 
 log = logging.getLogger(__name__)
 
@@ -171,6 +173,7 @@ def on_train_begin(self, logs=None):
             layer = self.model.get_layer(layer_name)
             self.updateable_weights.append(layer.weights)
 
+    @decorator_compiler
     def _update_weights(self):
         """Update all the weight with the corresponding multipliers
         Wrapped with tf.function to compensate the for loops as both weights variables

n3fit/src/n3fit/backends/keras_backend/internal_state.py

@@ -17,17 +17,34 @@
 import keras
 from keras import backend as K
 import numpy as np
-import tensorflow as tf
 
 log = logging.getLogger(__name__)
 
+# Prepare Keras-backend dependent functions
+if K.backend() == "torch":
 
-def set_eager(flag=True):
-    """Set eager mode on or off
-    for a very slow but fine grained debugging call this function as early as possible
-    ideally after the first tf import
-    """
-    tf.config.run_functions_eagerly(flag)
+    def set_eager(flag=True):
+        """Pytorch is eager by default"""
+        pass
+
+    def set_threading(threads, core):
+        """Not implemented"""
+        pass
+
+elif K.backend() == "tensorflow":
+    import tensorflow as tf
+
+    def set_eager(flag=True):
+        """Set eager mode on or off
+        for a very slow but fine grained debugging call this function as early as possible
+        ideally after the first tf import
+        """
+        tf.config.run_functions_eagerly(flag)
+
+    def set_threading(threads, cores):
+        """Set the Tensorflow inter and intra parallelism options"""
+        tf.config.threading.set_inter_op_parallelism_threads(threads)
+        tf.config.threading.set_intra_op_parallelism_threads(cores)
 
 
 def set_number_of_cores(max_cores=None, max_threads=None):
@@ -65,8 +82,7 @@ def set_number_of_cores(max_cores=None, max_threads=None):
 
     log.info("Setting the number of cores to: %d", cores)
     try:
-        tf.config.threading.set_inter_op_parallelism_threads(threads)
-        tf.config.threading.set_intra_op_parallelism_threads(cores)
+        set_threading(threads, cores)
     except RuntimeError:
         # If pdfflow is being used, tensorflow will already be initialized by pdfflow
         # maybe it would be good to drop completely pdfflow before starting the fit? (TODO ?)
@@ -130,7 +146,7 @@ def set_initial_state(debug=False, external_seed=None, max_cores=None, double_pr
     clear_backend_state()
 
     if double_precision:
-        tf.keras.backend.set_floatx('float64')
+        K.set_floatx('float64')
 
     # Set the number of cores depending on the user choice of max_cores
     # if debug mode and no number of cores set by the user, set to 1
@@ -143,7 +159,11 @@ def set_initial_state(debug=False, external_seed=None, max_cores=None, double_pr
 
     # Once again, if in debug mode or external_seed set, set also the TF seed
     if debug or external_seed:
-        tf.random.set_seed(use_seed)
+        if K.backend() == "tensorflow":
+            # if backend is tensorflow, keep the old seeding
+            tf.random.set_seed(use_seed)
+        else:
+            keras.utils.set_random_seed(use_seed)
 
 
 def get_physical_gpus():

n3fit/src/n3fit/backends/keras_backend/operations.py

@@ -122,7 +122,7 @@ def batchit(x, batch_dimension=0, **kwarg):
 
 
 # layer generation
-def numpy_to_input(numpy_array: np.typing.NDArray, name: Optional[str] = None):
+def numpy_to_input(numpy_array, name=None):
     """
     Takes a numpy array and generates an Input layer with the same shape,
     but with a batch dimension (of size 1) added.

n3fit/src/n3fit/msr.py

@@ -84,7 +84,9 @@ def generate_msr_model_and_grid(
 
     # 3. Prepare the pdf for integration by dividing by x
     pdf_integrand = Lambda(
-        lambda x_pdf: op.batchit(x_pdf[0], batch_dimension=1) * x_pdf[1], name="pdf_integrand"
+        lambda x_pdf: op.batchit(x_pdf[0], batch_dimension=1) * x_pdf[1],
+        name="pdf_integrand",
+        output_shape=pdf_xgrid_integration.shape[1:],
     )([x_divided, pdf_xgrid_integration])
 
     # 4. Integrate the pdf

pyproject.toml

@@ -97,13 +97,16 @@ fiatlux = {version = "*", optional = true}
 # without lhapdf
 pdfflow = {version = "^1.2.1", optional = true}
 lhapdf-management = {version = "^0.5", optional = true}
+# torch
+torch = {version = "*", optional = true}
 
 # Optional dependencies
 [tool.poetry.extras]
 tests = ["pytest", "pytest-mpl", "hypothesis"]
 docs = ["sphinxcontrib", "sphinx-rtd-theme", "sphinx", "tabulate"]
 qed = ["fiatlux"]
 nolha = ["pdfflow", "lhapdf-management"]
+torch = ["torch"]
 
 [tool.poetry-dynamic-versioning]
 enable = true