try-except for 3.11

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

@@ -477,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/callbacks.py

@@ -10,9 +10,10 @@
 
 import logging
 from time import time
+
 import numpy as np
 import tensorflow as tf
-from tensorflow.keras.callbacks import TensorBoard, Callback
+from tensorflow.keras.callbacks import Callback, TensorBoard
 
 log = logging.getLogger(__name__)
 
@@ -30,7 +31,7 @@ def __init__(self, count_range=100):
         self.last_time = 0
 
     def on_epoch_end(self, epoch, logs=None):
-        """ At the end of every epoch it checks the time """
+        """At the end of every epoch it checks the time"""
         new_time = time()
         if epoch == 0:
             # The first epoch is only useful for starting
@@ -45,13 +46,13 @@ def on_epoch_end(self, epoch, logs=None):
         self.last_time = new_time
 
     def on_train_end(self, logs=None):
-        """ Print the results """
+        """Print the results"""
         total_time = time() - self.starting_time
         n_times = len(self.all_times)
         # Skip the first 100 epochs to avoid fluctuations due to compilations of part of the code
         # by epoch 100 all parts of the code have usually been called so it's a good compromise
-        mean = np.mean(self.all_times[min(110, n_times-1):])
-        std = np.std(self.all_times[min(110, n_times-1):])
+        mean = np.mean(self.all_times[min(110, n_times - 1) :])
+        std = np.std(self.all_times[min(110, n_times - 1) :])
         log.info(f"> > Average time per epoch: {mean:.5} +- {std:.5} s")
         log.info(f"> > > Total time: {total_time/60:.5} min")
 
@@ -82,7 +83,7 @@ def on_epoch_begin(self, epoch, logs=None):
         self._current_loss = self.model.compute_losses()
 
     def on_epoch_end(self, epoch, logs=None):
-        """ Function to be called at the end of every epoch """
+        """Function to be called at the end of every epoch"""
         logs = self._current_loss
         print_stats = ((epoch + 1) % self.log_freq) == 0
         # Note that the input logs correspond to the fit before the weights are updated
@@ -123,7 +124,7 @@ def __init__(self, datasets, multipliers, update_freq=100):
         self.updateable_weights = []
 
     def on_train_begin(self, logs=None):
-        """ Save an instance of all relevant layers """
+        """Save an instance of all relevant layers"""
         for layer_name in self.datasets:
             layer = self.model.get_layer(layer_name)
             self.updateable_weights.append(layer.weights)
@@ -139,7 +140,7 @@ def _update_weights(self):
                 w.assign(w * multiplier)
 
     def on_epoch_end(self, epoch, logs=None):
-        """ Function to be called at the end of every epoch """
+        """Function to be called at the end of every epoch"""
         if (epoch + 1) % self.update_freq == 0:
             self._update_weights()
 

n3fit/src/n3fit/backends/keras_backend/constraints.py

@@ -3,8 +3,8 @@
 """
 
 import tensorflow as tf
-from tensorflow.keras.constraints import MinMaxNorm
 from tensorflow.keras import backend as K
+from tensorflow.keras.constraints import MinMaxNorm
 
 
 class MinMaxWeight(MinMaxNorm):
@@ -14,14 +14,11 @@ class MinMaxWeight(MinMaxNorm):
     """
 
     def __init__(self, min_value, max_value, **kwargs):
-        super(MinMaxWeight, self).__init__(
-            min_value=min_value, max_value=max_value, **kwargs
-        )
+        super(MinMaxWeight, self).__init__(min_value=min_value, max_value=max_value, **kwargs)
 
     def __call__(self, w):
         norms = K.sum(w, axis=self.axis, keepdims=True)
         desired = (
-            self.rate * K.clip(norms, self.min_value, self.max_value)
-            + (1 - self.rate) * norms
+            self.rate * K.clip(norms, self.min_value, self.max_value) + (1 - self.rate) * norms
         )
         return w * desired / (K.epsilon() + norms)

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

@@ -22,12 +22,13 @@
     Note that tensor operations can also be applied to layers as the output of a layer is a tensor
     equally operations are automatically converted to layers when used as such.
 """
+
 from typing import Optional
 
+import keras
 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
@@ -36,6 +37,12 @@
 
 from validphys.convolution import OP
 
+try:
+    # For tensorflow >= 2.16, Keras >= 3
+    concat_backend = keras.ops.concatenate
+except AttributeError:
+    concat_backend = tf.concat
+
 
 def evaluate(tensor):
     """Evaluate input tensor using the backend"""
@@ -250,11 +257,15 @@ 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 = keras.ops.concatenate(tensor_list, axis=axis)
-    if target_shape is not None:
-        return K.reshape(concatenated_tensor, target_shape)
-    else:
+    try:
+        # For tensorflow >= 2.16, Keras >= 3
+        concatenated_tensor = keras.ops.concatenate(tensor_list, axis=axis)
+    except AttributeError:
+        concatenated_tensor = tf.concat(tensor_list, axis=axis)
+
+    if target_shape is None:
         return concatenated_tensor
+    return K.reshape(concatenated_tensor, target_shape)
 
 
 # Mathematical operations

n3fit/src/n3fit/layers/observable.py

@@ -1,6 +1,8 @@
-from n3fit.backends import MetaLayer
+from abc import ABC, abstractmethod
+
 import numpy as np
-from abc import abstractmethod, ABC
+
+from n3fit.backends import MetaLayer
 from n3fit.backends import operations as op
 
 

n3fit/src/n3fit/model_gen.py

@@ -9,6 +9,7 @@
 
 
 """
+
 from dataclasses import dataclass
 from typing import Callable, List
 
@@ -73,7 +74,9 @@ def _generate_loss(self, mask=None):
         if self.invcovmat is not None:
             if self.rotation:
                 # If we have a matrix diagonal only, padd with 0s and hope it's not too heavy on memory
-                invcovmat_matrix = np.eye(self.invcovmat.shape[-1]) * self.invcovmat[..., np.newaxis]
+                invcovmat_matrix = (
+                    np.eye(self.invcovmat.shape[-1]) * self.invcovmat[..., np.newaxis]
+                )
                 if self.covmat is not None:
                     covmat_matrix = np.eye(self.covmat.shape[-1]) * self.covmat[..., np.newaxis]
                 else:
@@ -82,11 +85,7 @@ def _generate_loss(self, mask=None):
                 covmat_matrix = self.covmat
                 invcovmat_matrix = self.invcovmat
             loss = losses.LossInvcovmat(
-                invcovmat_matrix,
-                self.data,
-                mask,
-                covmat=covmat_matrix,
-                name=self.name
+                invcovmat_matrix, self.data, mask, covmat=covmat_matrix, name=self.name
             )
         elif self.positivity:
             loss = losses.LossPositivity(name=self.name, c=self.multiplier)