Fix some TPU tests.

`jax2onnx` fixed the ONNX export on TPU, but the `native_serialization_platforms` argument has to be passed to also export to CPU as the reloading happens with a CPU only version of TensorFlow.

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(ort_session.run(None, ort_inputs)[0], ref_output)
+        self.assertAllClose(
+            ort_session.run(None, ort_inputs)[0],
+            ref_output,
+            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(ort_session.run(None, ort_inputs)[0], ref_output)
+        self.assertAllClose(
+            ort_session.run(None, ort_inputs)[0],
+            ref_output,
+            tpu_atol=1e-3,
+            tpu_rtol=1e-2,
+        )
 
     @parameterized.named_parameters(
         named_product(