nightly test fix

e2e/integration_tests/convert_predict.py

@@ -37,6 +37,7 @@
 
 import numpy as np
 import tensorflow as tf
+import tf_keras
 from tensorflow.python.eager import def_function
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
@@ -232,16 +233,16 @@ def _create_saved_model_with_conv2d(save_dir):
     save_dir: directory name of where the saved model will be stored.
   """
   layers = [
-      tf.keras.layers.Conv2D(
+      tf_keras.layers.Conv2D(
           16, [3, 3], padding='same', use_bias=False),
-      tf.keras.layers.BatchNormalization(),
-      tf.keras.layers.ReLU()
+      tf_keras.layers.BatchNormalization(),
+      tf_keras.layers.ReLU()
   ]
-  model = tf.keras.Sequential(layers)
+  model = tf_keras.Sequential(layers)
   result = model.predict(tf.ones((1, 24, 24, 3)))
   # set the learning phase to avoid keara learning placeholder, which
   # will cause error when saving.
-  tf.keras.backend.set_learning_phase(0)
+  #tf_keras.backend.set_learning_phase(0)
   tf.saved_model.save(model, save_dir)
   return {
       "async": False,
@@ -263,14 +264,14 @@ def _create_saved_model_with_prelu(save_dir):
   # set the bias and alpha intitialize to make them constant and ensure grappler
   # be able to fuse the op.
   layers = [
-      tf.keras.layers.Conv2D(
+      tf_keras.layers.Conv2D(
           16, [3, 3], padding='same', use_bias=True,
           bias_initializer=tf.initializers.constant(0.25)),
-      tf.keras.layers.PReLU(alpha_initializer=tf.initializers.constant(0.25))
+      tf_keras.layers.PReLU(alpha_initializer=tf.initializers.constant(0.25))
   ]
-  model = tf.keras.Sequential(layers)
+  model = tf_keras.Sequential(layers)
   result = model.predict(tf.ones((1, 24, 24, 3)))
-  tf.keras.backend.set_learning_phase(0)
+  #tf_keras.backend.set_learning_phase(0)
   tf.saved_model.save(model, save_dir)
   return {
       "async": False,
@@ -351,13 +352,13 @@ def _create_saved_model_v2_with_tensorlist_ops(save_dir):
   Args:
     save_dir: directory name of where the saved model will be stored.
   """
-  model = tf.keras.Sequential()
-  model.add(tf.keras.layers.Embedding(100, 20, input_shape=[10]))
-  model.add(tf.keras.layers.GRU(4))
+  model = tf_keras.Sequential()
+  model.add(tf_keras.layers.Embedding(100, 20, input_shape=[10]))
+  model.add(tf_keras.layers.GRU(4))
 
   result = model.predict(tf.ones([1, 10]))
 
-  tf.keras.backend.set_learning_phase(0)
+  #tf_keras.backend.set_learning_phase(0)
   tf.saved_model.save(model, save_dir)
 
   return {
@@ -469,7 +470,7 @@ def lookup(input):
   }
 
 def _layers_mobilenet():
-  model = tf.keras.applications.MobileNetV2()
+  model = tf_keras.applications.MobileNetV2()
   model_path = 'mobilenet'
   tfjs.converters.save_keras_model(model, os.path.join(
       _tmp_dir, model_path))

remote-execution/BUILD.bazel

@@ -9,7 +9,7 @@ platform(
     ],
     exec_properties = {
         # We use the same docker image for remote builds as we use for CI testing.
-        "container-image": "docker://gcr.io/learnjs-174218/release@sha256:d85abab6146eaf1e01312bdb9e353a5efa0508b913dccf30fc5e505d009026ff",
+        "container-image": "docker://gcr.io/learnjs-174218/release:latest@sha256:f712eae902a364750727f1bc2e4bfc3f75be846e2277f4e8026f9c03752f00e4",
         # By default in Google Cloud Remote build execution, network access is disabled. We explicitly set the
         # property in the platform again in case the default ever changes. Network access is not desirable in
         # Bazel builds as it is potential source of flaky tests and therefore also breaks hermeticity.

tfjs-converter/python/BUILD.bazel

@@ -71,9 +71,10 @@ py_wheel(
         "jax>=0.4.13",
         "jaxlib>=0.4.13",
         "tensorflow>=2.13.0,<3",
+        "tf-keras>=2.13.0",
         "tensorflow-decision-forests>=1.5.0",
         "six>=1.16.0,<2",
-        "tensorflow-hub>=0.14.0",
+        "tensorflow-hub>=0.16.1",
         "packaging~=23.1",
     ],
     strip_path_prefixes = [

tfjs-converter/python/requirements.txt

@@ -3,7 +3,8 @@ importlib_resources>=5.9.0
 jax>=0.4.23
 jaxlib>=0.4.23
 tensorflow>=2.13.0,<3
-tensorflow-decision-forests>=1.5.0
+tf-keras>=2.16.0
+tensorflow-decision-forests>=1.9.0
 six>=1.16.0,<2
-tensorflow-hub>=0.14.0
+tensorflow-hub>=0.16.1
 packaging~=23.1

tfjs-converter/python/tensorflowjs/BUILD.bazel

@@ -116,6 +116,12 @@ py_library(
     deps = [requirement("debugpy")],
 )
 
+py_library(
+    name = "expect_tf_keras_installed",
+    # tf-keras is used to provide keras 2.0 support.
+    deps = [requirement("tf-keras")],
+)
+
 py_library(
     name = "quantization",
     srcs = ["quantization.py"],

tfjs-converter/python/tensorflowjs/converters/BUILD.bazel

@@ -59,6 +59,7 @@ py_test(
         "//tfjs-converter/python/tensorflowjs:expect_h5py_installed",
         "//tfjs-converter/python/tensorflowjs:expect_numpy_installed",
         "//tfjs-converter/python/tensorflowjs:expect_tensorflow_installed",
+        "//tfjs-converter/python/tensorflowjs:expect_tf_keras_installed",
         "//tfjs-converter/python/tensorflowjs:version",
     ],
 )
@@ -87,6 +88,7 @@ py_test(
         ":tf_module_mapper",
         "//tfjs-converter/python/tensorflowjs:expect_numpy_installed",
         "//tfjs-converter/python/tensorflowjs:expect_tensorflow_installed",
+        "//tfjs-converter/python/tensorflowjs:expect_tf_keras_installed",
     ],
 )
 
@@ -105,6 +107,7 @@ py_library(
         ":graph_rewrite_util",
         "//tfjs-converter/python/tensorflowjs:expect_numpy_installed",
         "//tfjs-converter/python/tensorflowjs:expect_tensorflow_installed",
+        "//tfjs-converter/python/tensorflowjs:expect_tf_keras_installed",
     ],
 )
 
@@ -154,6 +157,7 @@ py_library(
         ":graph_rewrite_util",
         "//tfjs-converter/python/tensorflowjs:expect_numpy_installed",
         "//tfjs-converter/python/tensorflowjs:expect_tensorflow_installed",
+        "//tfjs-converter/python/tensorflowjs:expect_tf_keras_installed",
     ],
 )
 
@@ -180,6 +184,7 @@ py_library(
         ":graph_rewrite_util",
         "//tfjs-converter/python/tensorflowjs:expect_numpy_installed",
         "//tfjs-converter/python/tensorflowjs:expect_tensorflow_installed",
+        "//tfjs-converter/python/tensorflowjs:expect_tf_keras_installed",
     ],
 )
 
@@ -229,6 +234,7 @@ py_library(
         "//tfjs-converter/python/tensorflowjs:expect_tensorflow_decision_forests_installed",
         "//tfjs-converter/python/tensorflowjs:expect_tensorflow_hub_installed",
         "//tfjs-converter/python/tensorflowjs:expect_tensorflow_installed",
+        "//tfjs-converter/python/tensorflowjs:expect_tf_keras_installed",
         "//tfjs-converter/python/tensorflowjs:resource_loader",
         "//tfjs-converter/python/tensorflowjs:version",
         "//tfjs-converter/python/tensorflowjs:write_weights",
@@ -283,6 +289,7 @@ py_binary(
         "//tfjs-converter/python/tensorflowjs:expect_PyInquirer_installed",
         "//tfjs-converter/python/tensorflowjs:expect_h5py_installed",
         "//tfjs-converter/python/tensorflowjs:expect_tensorflow_installed",
+        "//tfjs-converter/python/tensorflowjs:expect_tf_keras_installed",
     ],
 )
 
@@ -314,6 +321,7 @@ py_binary(
         ":tf_saved_model_conversion_v2",
         "//tfjs-converter/python/tensorflowjs:expect_h5py_installed",
         "//tfjs-converter/python/tensorflowjs:expect_tensorflow_installed",
+        "//tfjs-converter/python/tensorflowjs:expect_tf_keras_installed",
         "//tfjs-converter/python/tensorflowjs:version",
     ],
 )
@@ -325,6 +333,7 @@ py_binary(
     srcs_version = "PY3",
     deps = [
         "//tfjs-converter/python/tensorflowjs:expect_tensorflow_installed",
+        "//tfjs-converter/python/tensorflowjs:expect_tf_keras_installed",
     ],
 )
 

tfjs-converter/python/tensorflowjs/converters/converter.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Artifact conversion to and from Python TensorFlow and tf.keras."""
+"""Artifact conversion to and from Python TensorFlow and tf_keras."""
 
 from __future__ import absolute_import
 from __future__ import division
@@ -28,6 +28,7 @@
 import h5py
 import tensorflow.compat.v1 as tf1
 import tensorflow.compat.v2 as tf
+import tf_keras
 
 from tensorflowjs import quantization
 from tensorflowjs import version
@@ -50,7 +51,7 @@ def dispatch_keras_h5_to_tfjs_layers_model_conversion(
     - A weights-only HDF5 (e.g., generated with Keras Model's `save_weights()`
       method),
     - A topology+weights combined HDF5 (e.g., generated with
-      `tf.keras.model.save_model`).
+      `tf_keras.model.save_model`).
 
   Args:
     h5_path: path to an HDF5 file containing keras model data as a `str`.
@@ -199,7 +200,7 @@ def dispatch_keras_h5_to_tfjs_graph_model_conversion(
 
   Args:
     h5_path: Path to the HDF5-format file that contains the model saved from
-      keras or tf.keras.
+      keras or tf_keras.
     output_dir: The destination to which the tfjs GraphModel artifacts will be
       written.
     quantization_dtype_map: A mapping from dtype (`uint8`, `uint16`, `float16`)
@@ -223,7 +224,7 @@ def dispatch_keras_h5_to_tfjs_graph_model_conversion(
         'directory: %s' % h5_path)
 
   temp_savedmodel_dir = tempfile.mktemp(suffix='.savedmodel')
-  model = tf.keras.models.load_model(h5_path, compile=False)
+  model = tf_keras.models.load_model(h5_path, compile=False)
   model.save(temp_savedmodel_dir, include_optimizer=False, save_format='tf')
 
   # NOTE(cais): This cannot use `tf.compat.v1` because
@@ -253,13 +254,13 @@ def dispatch_keras_saved_model_to_tensorflowjs_conversion(
   """Converts keras model saved in the SavedModel format to tfjs format.
 
   Note that the SavedModel format exists in keras, but not in
-  keras-team/tf.keras.
+  keras-team/tf_keras.
 
   Args:
     keras_saved_model_path: path to a folder in which the
       assets/saved_model.json can be found. This is usually a subfolder
       that is under the folder passed to
-      `tf.keras.models.save_model()` and has a Unix epoch time
+      `tf_keras.models.save_model()` and has a Unix epoch time
       as its name (e.g., 1542212752).
     output_dir: Output directory to which the TensorFlow.js-format model JSON
       file and weights files will be written. If the directory does not exist,
@@ -274,7 +275,7 @@ def dispatch_keras_saved_model_to_tensorflowjs_conversion(
     metadata: User defined metadata map.
   """
   with tf.Graph().as_default(), tf.compat.v1.Session():
-    model = tf.keras.models.load_model(keras_saved_model_path)
+    model = tf_keras.models.load_model(keras_saved_model_path)
 
     # Save model temporarily in HDF5 format.
     temp_h5_path = tempfile.mktemp(suffix='.h5')
@@ -363,12 +364,12 @@ def dispatch_tensorflowjs_to_keras_keras_conversion(config_json_path, v3_path):
           'but cannot read valid JSON content from %s.' % config_json_path)
 
   model = keras_tfjs_loader.load_keras_keras_model(config_json_path)
-  tf.keras.saving.save_model(model, v3_path, save_format="keras")
+  tf_keras.saving.save_model(model, v3_path, save_format="keras")
 
 
 def dispatch_tensorflowjs_to_keras_saved_model_conversion(
     config_json_path, keras_saved_model_path):
-  """Converts a TensorFlow.js Layers model format to a tf.keras SavedModel.
+  """Converts a TensorFlow.js Layers model format to a tf_keras SavedModel.
 
   Args:
     config_json_path: Path to the JSON file that includes the model's
@@ -397,7 +398,7 @@ def dispatch_tensorflowjs_to_keras_saved_model_conversion(
 
   with tf.Graph().as_default(), tf.compat.v1.Session():
     model = keras_tfjs_loader.load_keras_model(config_json_path)
-    tf.keras.models.save_model(
+    tf_keras.models.save_model(
         model, keras_saved_model_path, save_format='tf')
 
 
@@ -751,7 +752,7 @@ def get_arg_parser():
       help='Input format. '
       'For "keras", the input path can be one of the two following formats:\n'
       '  - A topology+weights combined HDF5 (e.g., generated with'
-      '    `tf.keras.model.save_model()` method).\n'
+      '    `tf_keras.model.save_model()` method).\n'
       '  - A weights-only HDF5 (e.g., generated with Keras Model\'s '
       '    `save_weights()` method). \n'
       'For "keras_saved_model", the input_path must point to a subfolder '
@@ -885,7 +886,7 @@ def convert(arguments):
   if args.show_version:
     print('\ntensorflowjs %s\n' % version.version)
     print('Dependency versions:')
-    print('  keras %s' % tf.keras.__version__)
+    print('  keras %s' % tf_keras.__version__)
     print('  tensorflow %s' % tf.__version__)
     return