Guard tests that fail in libtpu==0.0.37, but succeed at nightly libtpu==0.0.38.dev20260315.

Preparing for JAX 0.9.2 release.

- 0.0.37: https://github.com/jax-ml/jax/actions/runs/23158234447
- 0.0.38.dev20260315: https://github.com/jax-ml/jax/actions/runs/23162037857

PiperOrigin-RevId: 885189173

Author: danielsuo

tests/pallas/ops_test.py

@@ -962,6 +962,8 @@ def test_sign(self, dtype, value):
         self.skipTest("f16 load not supported on TPU")
       if dtype in (jnp.int16, jnp.uint16) and jtu.get_tpu_version() < 6:
         self.skipTest("requires TPU v6+")
+      if dtype == jnp.bfloat16 and not jtu.is_cloud_tpu_at_least(2026, 3, 15):
+        self.skipTest("Requires libtpu >= 2026.3.15")
 
     @functools.partial(
         self.pallas_call,
@@ -1174,6 +1176,12 @@ def test_elementwise_array(self, fn, dtype):
           and not jtu.is_cloud_tpu_at_least(2026, 3, 1)
       ):
         self.skipTest("requires a newer libTPU")
+      if (
+          fn == jnp.sign
+          and dtype == "bfloat16"
+          and not jtu.is_cloud_tpu_at_least(2026, 3, 15)
+      ):
+        self.skipTest("Requires libtpu >= 2026.3.15")
       # TODO(b/370578663): implement these lowerings on TPU
       if fn in (
           jnp.acos, jnp.acosh, jnp.asin, jnp.asinh, jnp.atan, jnp.atanh,

tests/pallas/tpu_ops_test.py

@@ -956,6 +956,8 @@ def kernel(x_ref, y_ref):
       jnp.bfloat16,
   )
   def test_sigmoid(self, dtype):
+    if not jtu.is_cloud_tpu_at_least(2026, 3, 15):
+      self.skipTest("requires a newer libTPU")
 
     shape = (32, 128)
     x = jax.random.normal(jax.random.key(42), shape, dtype=dtype)

tests/pallas/tpu_pallas_state_test.py

@@ -353,6 +353,8 @@ def kernel(x_ref, out_ref, tmp_ref):
       core_type=[*pltpu.CoreType], use_tc_tiling_on_sc=[True, False]
   )
   def test_capture_scalar(self, core_type, use_tc_tiling_on_sc):
+    if not jtu.is_cloud_tpu_at_least(2026, 3, 15):
+      self.skipTest("Requires libtpu >= 2026.3.15")
     match core_type:
       case pltpu.CoreType.TC:
         mesh = pltpu.create_tensorcore_mesh("x", num_cores=1)

tests/pjit_test.py

@@ -830,6 +830,8 @@ def testWithCustomPRNGKey(self):
     pjit(lambda x: x, in_shardings=None, out_shardings=None)(key)
 
   def test_lower_with_wrapper_error(self):
+    if not jtu.is_cloud_tpu_at_least(2026, 3, 15):
+      self.skipTest("Requires libtpu >= 2026.3.15")
     @jax.jit
     def f(x):
       return x
@@ -10285,6 +10287,8 @@ def f(x):
   )
   @jtu.with_explicit_mesh((2, 2), ('x', 'y'))
   def test_reduce_sum_unreduced_inp_multi_mesh(self, axes, out_s, eq_out_s, mesh):
+    if not jtu.is_cloud_tpu_at_least(2026, 3, 15):
+      self.skipTest("Requires libtpu >= 2026.3.15")
 
     inp1 = jax.device_put(np.arange(16).reshape(8, 2), P('x', 'y'))
     inp2 = jax.device_put(np.arange(8).reshape(2, 4), P('y', None))
@@ -10301,6 +10305,8 @@ def f(x):
 
   @jtu.with_explicit_mesh((2,), 'x')
   def test_split_reduced_concat_unreduced(self, mesh):
+    if not jtu.is_cloud_tpu_at_least(2026, 3, 15):
+      self.skipTest("Requires libtpu >= 2026.3.15")
 
     x = jax.device_put(np.arange(8.).reshape(2, 4), P('x'))
     w = jax.device_put(np.arange(64.).reshape(4, 16), P(reduced={'x'}))