Fix some TPU tests.

- `jax2onnx` fixed the ONNX export on TPU.

Author: hertschuh

keras/src/backend/jax/excluded_tpu_tests.txt

@@ -1,79 +1 @@
-ConvTransposeBasicTest
-ExportArchiveTest::test_jax_endpoint_registration_tf_function
-ExportArchiveTest::test_jax_multi_unknown_endpoint_registration
-ExportArchiveTest::test_layer_export
-ExportArchiveTest::test_low_level_model_export_functional
-ExportArchiveTest::test_low_level_model_export_sequential
-ExportArchiveTest::test_low_level_model_export_subclass
-ExportArchiveTest::test_low_level_model_export_with_alias
-ExportArchiveTest::test_low_level_model_export_with_dynamic_dims_functional
-ExportArchiveTest::test_low_level_model_export_with_dynamic_dims_sequential
-ExportArchiveTest::test_low_level_model_export_with_dynamic_dims_subclass
-ExportArchiveTest::test_low_level_model_export_with_jax2tf_kwargs
-ExportArchiveTest::test_low_level_model_export_with_jax2tf_polymorphic_shapes
-ExportArchiveTest::test_model_combined_with_tf_preprocessing
-ExportArchiveTest::test_model_export_method_functional
-ExportArchiveTest::test_model_export_method_sequential
-ExportArchiveTest::test_model_export_method_subclass
-ExportArchiveTest::test_multi_input_output_functional_model
-ExportArchiveTest::test_non_standard_layer_signature
-ExportArchiveTest::test_non_standard_layer_signature_with_kwargs
-ExportArchiveTest::test_track_multiple_layers
-ExportONNXTest::test_export_with_input_names
-ExportONNXTest::test_export_with_opset_version_18
-ExportONNXTest::test_export_with_opset_version_none
-ExportONNXTest::test_model_with_input_structure_array
-ExportONNXTest::test_model_with_input_structure_dict
-ExportONNXTest::test_model_with_input_structure_tuple
-ExportONNXTest::test_model_with_multiple_inputs
-ExportONNXTest::test_standard_model_export_functional
-ExportONNXTest::test_standard_model_export_lstm
-ExportONNXTest::test_standard_model_export_sequential
-ExportONNXTest::test_standard_model_export_subclass
-ExportOpenVINOTest::test_model_with_input_structure_array
-ExportOpenVINOTest::test_model_with_input_structure_dict
-ExportOpenVINOTest::test_model_with_input_structure_tuple
-ExportOpenVINOTest::test_model_with_multiple_inputs
-ExportOpenVINOTest::test_standard_model_export_functional
-ExportOpenVINOTest::test_standard_model_export_sequential
-ExportOpenVINOTest::test_standard_model_export_subclass
-ExportSavedModelTest::test_input_signature_functional_<kerastensor shape=(none, 10), dtype=float32, sparse=false, ragged=false, name=inputs>
-ExportSavedModelTest::test_input_signature_functional_backend_tensor
-ExportSavedModelTest::test_input_signature_functional_inputspec(dtype=float32, shape=(none, 10), ndim=2)
-ExportSavedModelTest::test_input_signature_functional_tensorspec(shape=(none, 10), dtype=tf.float32, name='inputs')
-ExportSavedModelTest::test_input_signature_sequential_<kerastensor shape=(none, 10), dtype=float32, sparse=false, ragged=false, name=inputs>
-ExportSavedModelTest::test_input_signature_sequential_backend_tensor
-ExportSavedModelTest::test_input_signature_sequential_inputspec(dtype=float32, shape=(none, 10), ndim=2)
-ExportSavedModelTest::test_input_signature_sequential_tensorspec(shape=(none, 10), dtype=tf.float32, name='inputs')
-ExportSavedModelTest::test_input_signature_subclass_<kerastensor shape=(none, 10), dtype=float32, sparse=false, ragged=false, name=inputs>
-ExportSavedModelTest::test_input_signature_subclass_backend_tensor
-ExportSavedModelTest::test_input_signature_subclass_inputspec(dtype=float32, shape=(none, 10), ndim=2)
-ExportSavedModelTest::test_input_signature_subclass_tensorspec(shape=(none, 10), dtype=tf.float32, name='inputs')
-ExportSavedModelTest::test_jax_specific_kwargs_functional_false_{'enable_xla': true, 'native_serialization': true}
-ExportSavedModelTest::test_jax_specific_kwargs_functional_false_none
-ExportSavedModelTest::test_jax_specific_kwargs_functional_true_{'enable_xla': true, 'native_serialization': true}
-ExportSavedModelTest::test_jax_specific_kwargs_functional_true_none
-ExportSavedModelTest::test_jax_specific_kwargs_sequential_false_{'enable_xla': true, 'native_serialization': true}
-ExportSavedModelTest::test_jax_specific_kwargs_sequential_false_none
-ExportSavedModelTest::test_jax_specific_kwargs_sequential_true_{'enable_xla': true, 'native_serialization': true}
-ExportSavedModelTest::test_jax_specific_kwargs_sequential_true_none
-ExportSavedModelTest::test_jax_specific_kwargs_subclass_false_{'enable_xla': true, 'native_serialization': true}
-ExportSavedModelTest::test_jax_specific_kwargs_subclass_false_none
-ExportSavedModelTest::test_jax_specific_kwargs_subclass_true_{'enable_xla': true, 'native_serialization': true}
-ExportSavedModelTest::test_jax_specific_kwargs_subclass_true_none
-ExportSavedModelTest::test_model_with_input_structure_array
-ExportSavedModelTest::test_model_with_input_structure_dict
-ExportSavedModelTest::test_model_with_input_structure_tuple
-ExportSavedModelTest::test_model_with_multiple_inputs
-ExportSavedModelTest::test_model_with_non_trainable_state_export_functional
-ExportSavedModelTest::test_model_with_non_trainable_state_export_sequential
-ExportSavedModelTest::test_model_with_non_trainable_state_export_subclass
-ExportSavedModelTest::test_model_with_rng_export_functional
-ExportSavedModelTest::test_model_with_rng_export_sequential
-ExportSavedModelTest::test_model_with_rng_export_subclass
-ExportSavedModelTest::test_model_with_tf_data_layer_functional
-ExportSavedModelTest::test_model_with_tf_data_layer_sequential
-ExportSavedModelTest::test_model_with_tf_data_layer_subclass
-ExportSavedModelTest::test_standard_model_export_functional
-ExportSavedModelTest::test_standard_model_export_sequential
-ExportSavedModelTest::test_standard_model_export_subclass
+ConvTransposeBasicTest

keras/src/export/onnx_test.py

@@ -101,7 +101,12 @@ def test_standard_model_export(self, model_type):
         ort_inputs = {
             k.name: v for k, v in zip(ort_session.get_inputs(), [ref_input])
         }
-        self.assertAllClose(ref_output, ort_session.run(None, ort_inputs)[0])
+        self.assertAllClose(
+            ref_output,
+            ort_session.run(None, ort_inputs)[0],
+            tpu_atol=1e-3,
+            tpu_rtol=1e-2,
+        )
         # Test with a different batch size
         ort_inputs = {
             k.name: v
@@ -291,7 +296,12 @@ def test_export_with_input_names(self):
         ort_inputs = {
             k.name: v for k, v in zip(ort_session.get_inputs(), [ref_input])
         }
-        self.assertAllClose(ref_output, ort_session.run(None, ort_inputs)[0])
+        self.assertAllClose(
+            ref_output,
+            ort_session.run(None, ort_inputs)[0],
+            tpu_atol=1e-3,
+            tpu_rtol=1e-2,
+        )
 
     @parameterized.named_parameters(
         named_product(

keras/src/export/saved_model_test.py

@@ -64,6 +64,22 @@ def get_model(type="sequential", input_shape=(10,), layer_list=None):
     reason="Torch backend export (via torch_xla) is incompatible with np 2.0",
 )
 class ExportSavedModelTest(testing.TestCase):
+    def setUp(self):
+        super().setUp()
+        self.export_kwargs = {}
+        if testing.jax_uses_gpu():
+            self.export_kwargs = {
+                "jax2tf_kwargs": {
+                    "native_serialization_platforms": ("cpu", "cuda")
+                }
+            }
+        elif testing.jax_uses_tpu():
+            self.export_kwargs = {
+                "jax2tf_kwargs": {
+                    "native_serialization_platforms": ("cpu", "tpu")
+                }
+            }
+
     @parameterized.named_parameters(
         named_product(model_type=["sequential", "functional", "subclass"])
     )
@@ -74,7 +90,9 @@ def test_standard_model_export(self, model_type):
         ref_input = np.random.normal(size=(batch_size, 10)).astype("float32")
         ref_output = model(ref_input)
 
-        saved_model.export_saved_model(model, temp_filepath)
+        saved_model.export_saved_model(
+            model, temp_filepath, **self.export_kwargs
+        )
         revived_model = tf.saved_model.load(temp_filepath)
         self.assertAllClose(ref_output, revived_model.serve(ref_input))
         # Test with a different batch size
@@ -106,7 +124,9 @@ def call(self, inputs):
         ref_input = tf.random.normal((3, 10))
         ref_output = model(ref_input)
 
-        saved_model.export_saved_model(model, temp_filepath)
+        saved_model.export_saved_model(
+            model, temp_filepath, **self.export_kwargs
+        )
         revived_model = tf.saved_model.load(temp_filepath)
         self.assertEqual(ref_output.shape, revived_model.serve(ref_input).shape)
         # Test with a different batch size
@@ -142,7 +162,9 @@ def call(self, inputs):
         model = get_model(model_type, layer_list=[StateLayer()])
         model(tf.random.normal((3, 10)))
 
-        saved_model.export_saved_model(model, temp_filepath)
+        saved_model.export_saved_model(
+            model, temp_filepath, **self.export_kwargs
+        )
         revived_model = tf.saved_model.load(temp_filepath)
 
         # The non-trainable counter is expected to increment
@@ -164,7 +186,9 @@ def test_model_with_tf_data_layer(self, model_type):
         ref_input = np.random.normal(size=(batch_size, 10)).astype("float32")
         ref_output = model(ref_input)
 
-        saved_model.export_saved_model(model, temp_filepath)
+        saved_model.export_saved_model(
+            model, temp_filepath, **self.export_kwargs
+        )
         revived_model = tf.saved_model.load(temp_filepath)
         self.assertAllClose(ref_output, revived_model.serve(ref_input))
         # Test with a different batch size
@@ -206,7 +230,9 @@ def call(self, inputs):
         temp_filepath = os.path.join(self.get_temp_dir(), "exported_model")
         ref_output = model(tree.map_structure(ops.convert_to_tensor, ref_input))
 
-        saved_model.export_saved_model(model, temp_filepath)
+        saved_model.export_saved_model(
+            model, temp_filepath, **self.export_kwargs
+        )
         revived_model = tf.saved_model.load(temp_filepath)
         self.assertAllClose(ref_output, revived_model.serve(ref_input))
 
@@ -247,7 +273,9 @@ def build(self, y_shape, x_shape):
         ref_input_y = np.random.normal(size=(batch_size, 10)).astype("float32")
         ref_output = model(ref_input_x, ref_input_y)
 
-        saved_model.export_saved_model(model, temp_filepath)
+        saved_model.export_saved_model(
+            model, temp_filepath, **self.export_kwargs
+        )
         revived_model = tf.saved_model.load(temp_filepath)
         self.assertAllClose(
             ref_output, revived_model.serve(ref_input_x, ref_input_y)
@@ -282,7 +310,10 @@ def test_input_signature(self, model_type, input_signature):
         else:
             input_signature = (input_signature,)
         saved_model.export_saved_model(
-            model, temp_filepath, input_signature=input_signature
+            model,
+            temp_filepath,
+            input_signature=input_signature,
+            **self.export_kwargs,
         )
         revived_model = tf.saved_model.load(temp_filepath)
         self.assertAllClose(
@@ -318,11 +349,17 @@ def test_jax_specific_kwargs(self, model_type, is_static, jax2tf_kwargs):
         ref_input = ops.random.uniform((3, 10))
         ref_output = model(ref_input)
 
+        export_kwargs = self.export_kwargs.copy()
+        if "jax2tf_kwargs" in export_kwargs:
+            export_kwargs["jax2tf_kwargs"].update(jax2tf_kwargs)
+        else:
+            export_kwargs["jax2tf_kwargs"] = jax2tf_kwargs
+
         saved_model.export_saved_model(
             model,
             temp_filepath,
             is_static=is_static,
-            jax2tf_kwargs=jax2tf_kwargs,
+            **export_kwargs,
         )
         revived_model = tf.saved_model.load(temp_filepath)
         self.assertAllClose(ref_output, revived_model.serve(ref_input))
@@ -342,6 +379,22 @@ def test_jax_specific_kwargs(self, model_type, is_static, jax2tf_kwargs):
     testing.torch_uses_gpu(), reason="Leads to core dumps on CI"
 )
 class ExportArchiveTest(testing.TestCase):
+    def setUp(self):
+        super().setUp()
+        self.add_endpoint_kwargs = {}
+        if testing.jax_uses_gpu():
+            self.add_endpoint_kwargs = {
+                "jax2tf_kwargs": {
+                    "native_serialization_platforms": ("cpu", "cuda")
+                }
+            }
+        elif testing.jax_uses_tpu():
+            self.add_endpoint_kwargs = {
+                "jax2tf_kwargs": {
+                    "native_serialization_platforms": ("cpu", "tpu")
+                }
+            }
+
     @parameterized.named_parameters(
         named_product(model_type=["sequential", "functional", "subclass"])
     )
@@ -365,6 +418,7 @@ def test_low_level_model_export(self, model_type):
             "call",
             model.__call__,
             input_signature=[tf.TensorSpec(shape=(None, 10), dtype=tf.float32)],
+            **self.add_endpoint_kwargs,
         )
         export_archive.write_out(temp_filepath)
         revived_model = tf.saved_model.load(temp_filepath)
@@ -385,6 +439,7 @@ def test_low_level_model_export_with_alias(self):
             "call",
             model.__call__,
             input_signature=[tf.TensorSpec(shape=(None, 10), dtype=tf.float32)],
+            **self.add_endpoint_kwargs,
         )
         export_archive.write_out(
             temp_filepath,
@@ -431,6 +486,7 @@ def call(self, inputs):
                     tf.TensorSpec(shape=(None, None), dtype=tf.float32),
                 ]
             ],
+            **self.add_endpoint_kwargs,
         )
         export_archive.write_out(temp_filepath)
         revived_model = tf.saved_model.load(temp_filepath)
@@ -445,29 +501,13 @@ def call(self, inputs):
         reason="This test is only for the JAX backend.",
     )
     def test_low_level_model_export_with_jax2tf_kwargs(self):
-        temp_filepath = os.path.join(self.get_temp_dir(), "exported_model")
-
-        model = get_model()
-        ref_input = tf.random.normal((3, 10))
-        ref_output = model(ref_input)
-
         export_archive = saved_model.ExportArchive()
-        export_archive.track(model)
-        export_archive.add_endpoint(
-            "call",
-            model.__call__,
-            input_signature=[tf.TensorSpec(shape=(None, 10), dtype=tf.float32)],
-            jax2tf_kwargs={
-                "native_serialization": True,
-                "native_serialization_platforms": ("cpu", "tpu"),
-            },
-        )
         with self.assertRaisesRegex(
             ValueError, "native_serialization_platforms.*bogus"
         ):
             export_archive.add_endpoint(
-                "call2",
-                model.__call__,
+                "call",
+                lambda x: x,
                 input_signature=[
                     tf.TensorSpec(shape=(None, 10), dtype=tf.float32)
                 ],
@@ -476,9 +516,6 @@ def test_low_level_model_export_with_jax2tf_kwargs(self):
                     "native_serialization_platforms": ("cpu", "bogus"),
                 },
             )
-        export_archive.write_out(temp_filepath)
-        revived_model = tf.saved_model.load(temp_filepath)
-        self.assertAllClose(ref_output, revived_model.call(ref_input))
 
     @pytest.mark.skipif(
         backend.backend() != "jax",
@@ -506,12 +543,13 @@ def call(self, inputs):
                 "call",
                 model.__call__,
                 input_signature=signature,
-                jax2tf_kwargs={},
+                **self.add_endpoint_kwargs,
             )
             export_archive.write_out(temp_filepath)
 
         export_archive = saved_model.ExportArchive()
         export_archive.track(model)
+        # TODO
         export_archive.add_endpoint(
             "call",
             model.__call__,
@@ -585,6 +623,7 @@ def model_call(x):
             model_call,
             native_serialization=native_jax_compatible,
             polymorphic_shapes=["(b, 10)"],
+            # TODO
         )
 
         # you can now build a TF inference function