fix import in keras tests

sony · Jan 2, 2025 · cebef33 · cebef33
1 parent 20e25a0
commit cebef33
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 5 deletions.
diff --git a/tests/keras_tests/function_tests/test_sensitivity_metric_interest_points.py b/tests/keras_tests/function_tests/test_sensitivity_metric_interest_points.py
@@ -17,7 +17,11 @@
 import numpy as np
 from keras.applications.densenet import DenseNet121
 from keras.applications.mobilenet_v2 import MobileNetV2
-from keras.src.engine.input_layer import InputLayer
+
+if tf.__version__ >= "2.13":
+    from keras.src.engine.input_layer import InputLayer
+else:
+    from keras.engine.input_layer import InputLayer
 
 from packaging import version
 

diff --git a/tests/keras_tests/function_tests/test_set_layer_to_bitwidth.py b/tests/keras_tests/function_tests/test_set_layer_to_bitwidth.py
@@ -18,7 +18,11 @@
 
 from keras import Input
 from keras.layers import Conv2D
-from keras.src.engine.input_layer import InputLayer
+
+if tf.__version__ >= "2.13":
+    from keras.src.engine.input_layer import InputLayer
+else:
+    from keras.engine.input_layer import InputLayer
 
 from mct_quantizers import KerasActivationQuantizationHolder
 from model_compression_toolkit.core import QuantizationConfig

diff --git a/tests/keras_tests/function_tests/test_weights_activation_split_substitution.py b/tests/keras_tests/function_tests/test_weights_activation_split_substitution.py
@@ -13,13 +13,18 @@
 # limitations under the License.
 # ==============================================================================
 import copy
-
+import tensorflow as tf
 import keras
 import unittest
 
-from keras.layers import Conv2D, Conv2DTranspose, DepthwiseConv2D, Dense, BatchNormalization, ReLU, Input
+if tf.__version__ >= "2.13":
+    from keras.src.layers import Conv2D, Conv2DTranspose, DepthwiseConv2D, Dense, BatchNormalization, ReLU, Input
+    from keras.src.engine.input_layer import InputLayer
+else:
+    from keras.layers import Conv2D, Conv2DTranspose, DepthwiseConv2D, Dense, BatchNormalization, ReLU, Input
+    from keras.engine.input_layer import InputLayer
+
 import numpy as np
-from keras.src.engine.input_layer import InputLayer
 
 from model_compression_toolkit.core import QuantizationConfig
 from model_compression_toolkit.core.common.graph.virtual_activation_weights_node import VirtualSplitActivationNode, \