Fix TPU tests, removed test exclusion mechanism. (#22571)

- Remove the file-based exclusion mechanism as it is no longer needed.
- `jax2onnx` fixed the ONNX export on TPU.
- tf_save_model export works on TPU, simply, 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.
- Only conv transpose tests remain excluded.

Author: hertschuh

conftest.py

@@ -52,23 +52,11 @@ def pytest_collection_modifyitems(config, items):
                 line.strip() for line in openvino_skipped_tests if line.strip()
             ]
 
-    tpu_skipped_tests = []
     if backend() == "jax":
         import jax
 
         has_multiple_devices = jax.device_count() > 1
 
-        if jax.default_backend() == "tpu":
-            with open(
-                "keras/src/backend/jax/excluded_tpu_tests.txt", "r"
-            ) as file:
-                tpu_skipped_tests = file.readlines()
-                # it is necessary to check if stripped line is not empty
-                # and exclude such lines
-                tpu_skipped_tests = [
-                    line.strip() for line in tpu_skipped_tests if line.strip()
-                ]
-
     requires_trainable_backend = pytest.mark.skipif(
         backend() in ["numpy", "openvino"],
         reason="Trainer not implemented for NumPy and OpenVINO backend.",
@@ -96,14 +84,6 @@ def pytest_collection_modifyitems(config, items):
                         "Not supported operation by openvino backend",
                     )
                 )
-        # also, skip concrete tests for TPU when using JAX backend
-        for skipped_test in tpu_skipped_tests:
-            if skipped_test in item.nodeid:
-                item.add_marker(
-                    pytest.mark.skip(
-                        reason="Known TPU test failure",
-                    )
-                )
 
 
 def skip_if_backend(given_backend, reason):

keras/src/backend/jax/excluded_tpu_tests.txt

@@ -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(