kicking down the can

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

@@ -7,7 +7,6 @@
 
 import re
 
-import h5py
 import numpy as np
 import tensorflow as tf
 from tensorflow.keras import optimizers as Kopt
@@ -16,12 +15,6 @@
 
 import n3fit.backends.keras_backend.operations as op
 
-# Check the TF version to check if legacy-mode is needed (TF < 2.2)
-tf_version = tf.__version__.split(".")
-if int(tf_version[0]) == 2 and int(tf_version[1]) < 2:
-    raise NotImplementedError("n3fit needs TF > 2.2 in order to work")
-
-
 # We need a function to transform tensors to numpy/python primitives
 # which is not part of the official TF interface and can change with the version
 if hasattr(tf_utils, "to_numpy_or_python_type"):
@@ -484,7 +477,7 @@ def set_layer_replica_weights(layer, weights, i_replica: int):
     if is_stacked_single_replicas(layer):
         layer.get_layer(f"{NN_PREFIX}_{i_replica}").set_weights(weights)
         return
-
+    
     full_weights = [w.numpy() for w in layer.weights]
     for w_old, w_new in zip(full_weights, weights):
         w_old[i_replica : i_replica + 1] = w_new

n3fit/src/n3fit/backends/keras_backend/multi_dense.py

@@ -99,8 +99,8 @@ def call(self, inputs):
         If the input already contains multiple replica outputs, it is equivalent
         to applying each replica to its corresponding input.
         """
-        if inputs.dtype.base_dtype != self._compute_dtype_object.base_dtype:
-            inputs = tf.cast(inputs, dtype=self._compute_dtype_object)
+        # cast always
+        inputs = tf.cast(inputs, dtype=self.compute_dtype)
 
         outputs = self.matmul(inputs)
 

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

@@ -27,6 +27,7 @@
 import numpy as np
 import numpy.typing as npt
 import tensorflow as tf
+import keras
 from tensorflow.keras import backend as K
 from tensorflow.keras.layers import Input
 from tensorflow.keras.layers import Lambda as keras_Lambda
@@ -249,8 +250,8 @@ def concatenate(tensor_list, axis=-1, target_shape=None, name=None):
     Concatenates a list of numbers or tensor into a bigger tensor
     If the target shape is given, the output is reshaped to said shape
     """
-    concatenated_tensor = tf.concat(tensor_list, axis, name=name)
-    if target_shape:
+    concatenated_tensor = keras.ops.concatenate(tensor_list, axis=axis)
+    if target_shape is not None:
         return K.reshape(concatenated_tensor, target_shape)
     else:
         return concatenated_tensor

n3fit/src/n3fit/layers/observable.py

@@ -68,9 +68,6 @@ def __init__(self, fktable_data, fktable_arr, operation_name, nfl=14, **kwargs):
         self.operation = op.c_to_py_fun(operation_name)
         self.output_dim = self.fktables[0].shape[0]
 
-    def compute_output_shape(self, input_shape):
-        return (self.output_dim, None)
-
     # Overridables
     @abstractmethod
     def gen_mask(self, basis):

n3fit/src/n3fit/model_gen.py

@@ -644,7 +644,7 @@ def compute_unnormalized_pdf(x):
             # add batch and flavor dimensions
             photon_integrals = op.batchit(op.batchit(photons.integral))
         else:
-            photon_integrals = np.zeros((1, num_replicas, 1))
+            photon_integrals = op.numpy_to_tensor(np.zeros((1, num_replicas, 1)))
 
         PDFs_normalized = sumrule_layer(
             {