Remove direct tensorflow imports from tests

To be fully agnostic, we prefer having utility functions that are
checking the backend use like now `is_in_gpu_mode()`.

tests/conftest.py

@@ -21,6 +21,12 @@
 from numpy.testing import assert_almost_equal
 
 from decomon.core import ForwardMode, Slope
+from decomon.keras_utils import (
+    BACKEND_JAX,
+    BACKEND_NUMPY,
+    BACKEND_PYTORCH,
+    BACKEND_TENSORFLOW,
+)
 from decomon.models.utils import ConvertMethod
 from decomon.types import Tensor
 
@@ -165,10 +171,24 @@ def __call__(self, inputs_: List[np.ndarray]):
 
 class Helpers:
     @staticmethod
-    def tensorflow_in_GPU_mode() -> bool:
-        import tensorflow
-
-        return len(tensorflow.config.list_physical_devices("GPU")) > 0
+    def in_GPU_mode() -> bool:
+        backend = keras.config.backend()
+        if backend == BACKEND_TENSORFLOW:
+            import tensorflow
+
+            return len(tensorflow.config.list_physical_devices("GPU")) > 0
+        elif backend == BACKEND_PYTORCH:
+            import torch
+
+            return torch.cuda.is_available()
+        elif backend == BACKEND_NUMPY:
+            return False
+        elif backend == BACKEND_JAX:
+            import jax
+
+            return jax.devices()[0].platform != "cpu"
+        else:
+            raise NotImplementedError(f"Not implemented for {backend} backend.")
 
     @staticmethod
     def is_method_mode_compatible(method, mode):

tests/test_backward_conv.py

@@ -4,14 +4,13 @@
 import keras.config as keras_config
 import numpy as np
 import pytest
-from tensorflow.python.keras.backend import _get_available_gpus
 
 from decomon.backward_layers.convert import to_backward
 from decomon.layers.decomon_layers import DecomonConv2D
 
 
 def test_Decomon_conv_box(data_format, padding, use_bias, mode, floatx, decimal, helpers):
-    if data_format == "channels_first" and not helpers.tensorflow_in_GPU_mode():
+    if data_format == "channels_first" and not helpers.in_GPU_mode():
         pytest.skip("data format 'channels first' is possible only in GPU mode")
 
     odd, m_0, m_1 = 0, 0, 1

tests/test_backward_layers.py

@@ -1,7 +1,6 @@
 import keras.config as keras_config
 import pytest
 from keras.layers import Layer, Reshape
-from tensorflow.python.keras.backend import _get_available_gpus
 
 from decomon.backward_layers.convert import to_backward
 from decomon.core import ForwardMode, Slope
@@ -105,7 +104,7 @@ def test_Backward_Activation_multiD_box(odd, activation, floatx, decimal, mode,
 
 
 def test_Backward_Flatten_multiD_box(odd, floatx, decimal, mode, data_format, helpers):
-    if data_format == "channels_first" and not helpers.tensorflow_in_GPU_mode():
+    if data_format == "channels_first" and not helpers.in_GPU_mode():
         pytest.skip("data format 'channels first' is possible only in GPU mode")
 
     dc_decomp = False

tests/test_backward_native_layers.py

@@ -1,7 +1,6 @@
 import keras.config as keras_config
 import pytest
 from keras.layers import Activation, Flatten, Reshape
-from tensorflow.python.keras.backend import _get_available_gpus
 
 from decomon.backward_layers.convert import to_backward
 
@@ -65,7 +64,7 @@ def test_Backward_NativeActivation_multiD_box(odd, activation, floatx, decimal,
 
 
 def test_Backward_NativeFlatten_multiD_box(odd, floatx, decimal, mode, data_format, helpers):
-    if data_format == "channels_first" and not helpers.tensorflow_in_GPU_mode():
+    if data_format == "channels_first" and not helpers.in_GPU_mode():
         pytest.skip("data format 'channels first' is possible only in GPU mode")
 
     dc_decomp = False

tests/test_conv.py

@@ -5,15 +5,14 @@
 import numpy as np
 import pytest
 from keras.layers import Conv2D
-from tensorflow.python.keras.backend import _get_available_gpus
 
 from decomon.core import ForwardMode, get_affine, get_ibp
 from decomon.layers.convert import to_decomon
 from decomon.layers.decomon_layers import DecomonConv2D
 
 
 def test_Decomon_conv_box(data_format, mode, dc_decomp, floatx, decimal, helpers):
-    if data_format == "channels_first" and not helpers.tensorflow_in_GPU_mode():
+    if data_format == "channels_first" and not helpers.in_GPU_mode():
         pytest.skip("data format 'channels first' is possible only in GPU mode")
 
     odd, m_0, m_1 = 0, 0, 1

tests/test_utils_conv.py

@@ -17,7 +17,7 @@ def test_toeplitz_from_Keras(channels, filter_size, strides, flatten, data_forma
     if floatx == 16:
         decimal = 0
 
-    if data_format == "channels_first" and not helpers.tensorflow_in_GPU_mode():
+    if data_format == "channels_first" and not helpers.in_GPU_mode():
         pytest.skip("data format 'channels first' is possible only in GPU mode")
 
     dc_decomp = False
@@ -67,7 +67,7 @@ def test_toeplitz_from_Keras(channels, filter_size, strides, flatten, data_forma
 def test_toeplitz_from_Decomon(
     floatx, decimal, mode, channels, filter_size, strides, flatten, data_format, padding, helpers
 ):
-    if data_format == "channels_first" and not helpers.tensorflow_in_GPU_mode():
+    if data_format == "channels_first" and not helpers.in_GPU_mode():
         pytest.skip("data format 'channels first' is possible only in GPU mode")
 
     odd, m_0, m_1 = 0, 0, 1