[Mosaic TPU] Give more time for the changes to propagate to libtpu

PiperOrigin-RevId: 873987035

Author: apaszke

jax/_src/tpu_custom_call.py

@@ -65,11 +65,11 @@
 # We should also add a TODO to remove the conditional one month later.
 def get_ir_version(ctx: mlir.LoweringRuleContext) -> int | None:
   backend = ctx.module_context.get_backend(optional=True)
-  # TODO(apaszke): remove the forward compatibility check after 2025-3-21.
+  # TODO(apaszke): remove the forward compatibility check after 2025-4-1.
   if (
       ctx.is_forward_compat()
       or backend is None
-      or is_cloud_tpu_older_than(2026, 2, 21, backend)
+      or is_cloud_tpu_older_than(2026, 3, 1, backend)
   ):
     return 9
   return None

tests/pallas/ops_test.py

@@ -1127,7 +1127,7 @@ def test_elementwise(self, fn, dtype):
       if (
           fn in (jnp.sin, jnp.cos, jnp.tan)
           and dtype == "bfloat16"
-          and not jtu.is_cloud_tpu_at_least(2026, 2, 21)
+          and not jtu.is_cloud_tpu_at_least(2026, 3, 1)
       ):
         self.skipTest("requires a newer libTPU")
       # TODO(b/370578663): implement these lowerings on TPU

tests/pallas/tpu_ops_test.py

@@ -791,9 +791,9 @@ def test_pack_elementwise(self, config, shape):
     if not jtu.is_device_tpu_at_least(version=5):
       self.skipTest("Requires TPU v5+")
     if packed_dtype == jnp.int2:
-      if not jtu.is_cloud_tpu_at_least(2026, 2, 21):
+      if not jtu.is_cloud_tpu_at_least(2026, 3, 1):
         raise self.skipTest(
-            "int2 is only supported for tpu at least 02/21/2026"
+            "int2 is only supported for tpu at least 03/01/2026"
         )
       if (shape[-2] % (8 * 16)) or (shape[-1] % 128):
         raise self.skipTest(

tests/pallas/tpu_sparsecore_pallas_test.py

@@ -1007,6 +1007,9 @@ def kernel(x_ref, o_ref):
     np.testing.assert_array_equal(kernel(x), x + np.arange(self.num_lanes))
 
   def test_write_to_transformed_ref(self):
+    if not jtu.is_cloud_tpu_at_least(2026, 3, 1):
+      self.skipTest("Requires a newer libTPU")
+
     x = jnp.arange(2 * self.num_lanes)
 
     @self.vector_subcore_kernel(out_shape=x)