fix converter tests (#7901)

* fix converter tests

* fix jax test

* fix version check tests

tfjs-converter/python/BUILD.bazel

@@ -66,14 +66,15 @@ py_wheel(
     license = "Apache 2.0",
     python_tag = "py3",
     requires = [
-        "flax>=0.6.2,<0.6.3",
+        "flax>=0.7.2",
         "importlib_resources>=5.9.0",
-        "jax>=0.3.16",
-        "tensorflow>=2.12.0,<3",
-        "tensorflow-decision-forests>=1.3.0",
-        "six>=1.12.0,<2",
-        "tensorflow-hub>=0.13.0",
-        "packaging~=20.9",
+        "jax>=0.4.13",
+        "jaxlib>=0.4.13",
+        "tensorflow>=2.13.0,<3",
+        "tensorflow-decision-forests>=1.5.0",
+        "six>=1.16.0,<2",
+        "tensorflow-hub>=0.14.0",
+        "packaging~=23.1",
     ],
     strip_path_prefixes = [
         "tfjs-converter/python",

tfjs-converter/python/requirements.txt

@@ -1,8 +1,9 @@
-flax>=0.6.2,<0.6.3
-jax>=0.3.16
+flax>=0.7.2
 importlib_resources>=5.9.0
-tensorflow>=2.12.0,<3
-tensorflow-decision-forests>=1.3.0
-six>=1.12.0,<2
-tensorflow-hub>=0.13.0; python_version >= "3"
-packaging~=20.9
+jax>=0.4.13
+jaxlib>=0.4.13
+tensorflow>=2.13.0,<3
+tensorflow-decision-forests>=1.5.0
+six>=1.16.0,<2
+tensorflow-hub>=0.14.0
+packaging~=23.1

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

@@ -24,6 +24,7 @@
 import shutil
 import tempfile
 import unittest
+import keras
 
 import numpy as np
 import tensorflow.compat.v2 as tf
@@ -98,9 +99,8 @@ def testConvertSavedKerasModelNoSplitByLayer(self):
     self.assertIsInstance(model_json['model_config'], dict)
     self.assertIsInstance(model_json['model_config']['config'], dict)
     self.assertIn('layers', model_json['model_config']['config'])
-
     # Check the loaded weights.
-    self.assertEqual(tf.keras.__version__, model_json['keras_version'])
+    self.assertEqual(keras.__version__, model_json['keras_version'])
     self.assertEqual('tensorflow', model_json['backend'])
     self.assertEqual(1, len(groups))
     self.assertEqual(3, len(groups[0]))
@@ -136,7 +136,7 @@ def testConvertSavedKerasModelSplitByLayer(self):
     self.assertIn('layers', model_json['model_config']['config'])
 
     # Check the loaded weights.
-    self.assertEqual(tf.keras.__version__, model_json['keras_version'])
+    self.assertEqual(keras.__version__, model_json['keras_version'])
     self.assertEqual('tensorflow', model_json['backend'])
     self.assertEqual(2, len(groups))
     self.assertEqual(2, len(groups[0]))
@@ -274,6 +274,7 @@ def testTensorflowjsToKerasConversionSucceeds(self):
           tf.keras.layers.Dense(
               1, use_bias=False, kernel_initializer='ones', name='Dense2')])
       h5_path = os.path.join(self._tmp_dir, 'SequentialModel.h5')
+      sequential_model.use_legacy_config = True
       sequential_model.save(h5_path)
       converter.dispatch_keras_h5_to_tfjs_layers_model_conversion(
           h5_path, output_dir=self._tmp_dir)
@@ -287,6 +288,7 @@ def testTensorflowjsToKerasConversionSucceeds(self):
     # Load the new H5 and compare the model JSONs.
     with tf.Graph().as_default(), tf.compat.v1.Session():
       new_model = tf.keras.models.load_model(new_h5_path)
+      new_model.use_legacy_config = True
       self.assertEqual(old_model_json, new_model.to_json())
 
   def testTensorflowjsToKerasConversionFailsOnDirInputPath(self):

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

@@ -97,7 +97,7 @@ def test_convert_poly(self):
   def test_convert_tf_poly_mismatch_raises(self):
     apply_fn = lambda params, x: jnp.sum(x) * params['w']
     with self.assertRaisesRegex(
-        ValueError, 'polymorphic shape.* must match .* for argument shape'):
+        ValueError, 'syntax error in polymorphic shape.* in dimension.* Parsed.*, remaining.*'):
       jax_conversion.convert_jax(
           apply_fn,
           {'w': 0.5},

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

@@ -25,6 +25,7 @@
 import tempfile
 import unittest
 import six
+import keras
 
 import h5py
 import numpy as np
@@ -136,7 +137,7 @@ def testConvertModelWithNestedLayerNames(self):
     # Check meta-data in the artifact JSON.
     self.assertEqual(model_json['format'], 'layers-model')
     self.assertEqual(model_json['generatedBy'],
-                     'keras v%s' % tf.keras.__version__)
+                     'keras v%s' % keras.__version__)
     self.assertEqual(
         model_json['convertedBy'],
         'TensorFlow.js Converter v%s' % version.version)
@@ -190,8 +191,7 @@ def testConvertMergedModelFromSimpleModelNoSplitByLayer(self):
     # Check the loaded weights.
     # By default, all weights of the model ought to be put in the same group.
     self.assertEqual(1, len(groups))
-
-    self.assertEqual(tf.keras.__version__, out['keras_version'])
+    self.assertEqual(keras.__version__, out['keras_version'])
     self.assertEqual('tensorflow', out['backend'])
     weight_group = groups[0]
     self.assertEqual(3, len(weight_group))
@@ -243,8 +243,7 @@ def testConvertMergedModelFromSimpleModelSplitByLayer(self):
     # Due to `split_by_layer=True`, there ought to be two weight groups,
     # because the model has two layers.
     self.assertEqual(2, len(groups))
-
-    self.assertEqual(tf.keras.__version__, out['keras_version'])
+    self.assertEqual(keras.__version__, out['keras_version'])
     self.assertEqual('tensorflow', out['backend'])
     self.assertEqual(2, len(groups[0]))
     kernel1 = groups[0][0]

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

@@ -46,6 +46,7 @@ def tearDown(self):
 
   def _saveKerasModelForTest(self, path):
     model = tf.keras.Sequential()
+    model.use_legacy_config = True
     model.add(tf.keras.layers.Dense(
         2, input_shape=[12], bias_initializer='random_normal', name='dense'))
     model.add(tf.keras.layers.Dense(
@@ -58,6 +59,7 @@ def _saveKerasModelForTest(self, path):
 
   def _saveRNNKerasModelForTest(self, path):
     model = tf.keras.Sequential()
+    model.use_legacy_config = True
     model.add(tf.keras.layers.Embedding(100, 20, input_shape=[10]))
     model.add(tf.keras.layers.SimpleRNN(4))
     keras_h5_conversion.save_keras_model(model, path)
@@ -75,7 +77,7 @@ def testLoadKerasModelAndWeights(self):
     with tf.Graph().as_default(), tf.compat.v1.Session():
       model2 = keras_tfjs_loader.load_keras_model(
           os.path.join(tfjs_path, 'model.json'))
-
+      model2.use_legacy_config = True
       # Verify the equality of all the weight values.
       model2_weight_values = model2.get_weights()
       self.assertEqual(len(model1_weight_values), len(model2_weight_values))
@@ -98,7 +100,7 @@ def testLoadKerasRNNModelAndWeights(self):
     with tf.Graph().as_default(), tf.compat.v1.Session():
       model2 = keras_tfjs_loader.load_keras_model(
           os.path.join(tfjs_path, 'model.json'))
-
+      model2.use_legacy_config = True
       # Verify the equality of all the weight values.
       model2_weight_values = model2.get_weights()
       self.assertEqual(len(model1_weight_values), len(model2_weight_values))
@@ -127,6 +129,7 @@ def testDeserializeKerasModelTopologyOnlyFromBytesIO(self):
 
     with tf.Graph().as_default(), tf.compat.v1.Session():
       model2 = keras_tfjs_loader.deserialize_keras_model(buff.read())
+      model2.use_legacy_config = True
 
       # The two model JSONs should match exactly.
       self.assertEqual(model1.to_json(), model2.to_json())
@@ -145,7 +148,7 @@ def testDeserializeKerasModelTopologyOnlyFromJSONDict(self):
 
     with tf.Graph().as_default(), tf.compat.v1.Session():
       model2 = keras_tfjs_loader.deserialize_keras_model(config_json)
-
+      model2.use_legacy_config = True
       # The two model JSONs should match exactly.
       self.assertEqual(model1.to_json(), model2.to_json())
 
@@ -171,6 +174,7 @@ def testDeserializeKerasModelTopologyAndWeightsFromBuffers(self):
     with tf.Graph().as_default(), tf.compat.v1.Session():
       model2 = keras_tfjs_loader.deserialize_keras_model(
           json_buff, weight_data=weight_buffers)
+      model2.use_legacy_config = True
 
       # Verify the equality of all the weight values.
       model2_weight_values = model2.get_weights()
@@ -200,6 +204,7 @@ def testDeserializeKerasModelTopologyAndWeightsFromFileObjects(self):
     with tf.Graph().as_default(), tf.compat.v1.Session():
       model2 = keras_tfjs_loader.deserialize_keras_model(
           json_file, weight_files)
+      model2.use_legacy_config = True
 
       # Verify the equality of all the weight values.
       model2_weight_values = model2.get_weights()
@@ -225,7 +230,7 @@ def testLoadKerasModelWithCurrentWorkingDirectoryRelativePath(self):
     with tf.Graph().as_default(), tf.compat.v1.Session():
       # Use a relative path under the current working directory.
       model2 = keras_tfjs_loader.load_keras_model('model.json')
-
+      model2.use_legacy_config = True
       # Verify the equality of all the weight values.
       model2_weight_values = model2.get_weights()
       self.assertEqual(len(model1_weight_values), len(model2_weight_values))
@@ -247,6 +252,7 @@ def testLoadKerasModelWithoutWeights(self):
     with tf.Graph().as_default(), tf.compat.v1.Session():
       model2 = keras_tfjs_loader.load_keras_model(
           os.path.join(tfjs_path, 'model.json'), load_weights=False)
+      model2.use_legacy_config = True
       model2_weight_values = model2.get_weights()
       self.assertEqual(len(model1_weight_values), len(model2_weight_values))
       for model1_weight_value, model2_weight_value in zip(
@@ -276,7 +282,7 @@ def testLoadKerasModelFromNonDefaultWeightsPathWorks(self):
     with tf.Graph().as_default(), tf.compat.v1.Session():
       model2 = keras_tfjs_loader.load_keras_model(
           new_model_json_path, weights_path_prefix=tfjs_path)
-
+      model2.use_legacy_config = True
       # Verify the equality of all the weight values.
       model2_weight_values = model2.get_weights()
       self.assertEqual(len(model1_weight_values), len(model2_weight_values))
@@ -329,7 +335,7 @@ def testLoadKerasModelFromDataBuffers(self):
       model2 = keras_tfjs_loader.load_keras_model(
           os.path.join(tfjs_path, 'model.json'),
           weights_data_buffers=data_buffers)
-
+      model2.use_legacy_config = True
       # Verify the equality of all the weight values.
       model2_weight_values = model2.get_weights()
       self.assertEqual(len(model1_weight_values), len(model2_weight_values))