Reduce number of sampled tests and mark some more tests as thread unsafe.

PiperOrigin-RevId: 815059810

Author: danielsuo

jax/experimental/jax2tf/tests/call_tf_test.py

@@ -871,6 +871,7 @@ def _transfer_guard(guard_level):
       jax2tf.call_tf(tf_fun)(jax_array_on_gpu)
 
 
+@jtu.thread_unsafe_test_class()
 class RoundTripToJaxTest(tf_test_util.JaxToTfTestCase):
   """Reloading output of jax2tf into JAX with call_tf."""
 

jax/experimental/jax2tf/tests/shape_poly_test.py

@@ -232,6 +232,7 @@ def check_shape_poly(tst, f_jax: Callable, *,
   return h.run_test(tst)
 
 
+@jtu.thread_unsafe_test_class()
 class ShapePolyTest(tf_test_util.JaxToTfTestCase):
 
   def test_simple_unary(self):

tests/BUILD

@@ -1437,7 +1437,7 @@ jax_multiplatform_test(
     },
     # Use fewer cases to prevent timeouts.
     backend_variant_args = {
-        "cpu": ["--jax_num_generated_cases=30"],
+        "cpu": ["--jax_num_generated_cases=20"],
         "cpu_x32": ["--jax_num_generated_cases=30"],
         "gpu_p100": ["--jax_num_generated_cases=40"],
         "gpu_v100": ["--jax_num_generated_cases=40"],