Support tensorflow 2.13

TileDB-Inc · Jul 23, 2023 · 5a3bb71 · 5a3bb71
1 parent d7d67fd
commit 5a3bb71
Showing 1 changed file with 11 additions and 2 deletions.
diff --git a/tiledb/ml/models/tensorflow_keras.py b/tiledb/ml/models/tensorflow_keras.py
@@ -15,24 +15,33 @@
 
 from ._base import Meta, TileDBArtifact, Timestamp
 
-FunctionalOrSequential = (keras.models.Functional, keras.models.Sequential)
 keras_major, keras_minor, keras_patch = keras.__version__.split(".")
+FunctionalOrSequential = keras.models.Sequential
 # Handle keras <=v2.10
 if int(keras_major) <= 2 and int(keras_minor) <= 10:
+    FunctionalOrSequential = (keras.models.Functional, keras.models.Sequential)
     TFOptimizer = keras.optimizers.TFOptimizer
     get_json_type = keras.saving.saved_model.json_utils.get_json_type
     preprocess_weights_for_loading = (
         keras.saving.hdf5_format.preprocess_weights_for_loading
     )
     saving_utils = keras.saving.saving_utils
 # Handle keras >=v2.11
-else:
+elif int(keras_major) <= 2 and int(keras_minor) <= 12:
+    FunctionalOrSequential = (keras.models.Functional, keras.models.Sequential)
     TFOptimizer = tf.keras.optimizers.legacy.Optimizer
     get_json_type = keras.saving.legacy.saved_model.json_utils.get_json_type
     preprocess_weights_for_loading = (
         keras.saving.legacy.hdf5_format.preprocess_weights_for_loading
     )
     saving_utils = keras.saving.legacy.saving_utils
+else:
+    TFOptimizer = tf.keras.optimizers.legacy.Optimizer
+    get_json_type = keras.src.saving.legacy.saved_model.json_utils.get_json_type
+    preprocess_weights_for_loading = (
+        keras.src.saving.legacy.hdf5_format.preprocess_weights_for_loading
+    )
+    saving_utils = keras.src.saving.legacy.saving_utils
 
 
 class TensorflowKerasTileDBModel(TileDBArtifact[tf.keras.Model]):