fix tpu_backend_support checking

Author: gulsumgudukbay

tests/conftest.py

@@ -54,15 +54,11 @@ def _has_tpu_backend_support() -> bool:
   except Exception:  # pragma: no cover  pylint: disable=broad-exception-caught
     return False
 
-  # Heuristic: TPU-enabled jaxlib exposes a TPU client/extension module.
-  # This check avoids initializing any backend at collection time.
-  for mod in ("jaxlib.tpu_client", "jaxlib._src.tpu_client", "jaxlib.tpu_extension"):
-    try:
-      if importlib.util.find_spec(mod) is not None:
-        return True
-    except Exception:  # pragma: no cover  pylint: disable=broad-exception-caught
-      continue
-  return False
+  # Heuristic: TPU backend support is provided via the `libtpu` package.
+  try:
+    return importlib.util.find_spec("libtpu") is not None
+  except Exception:  # pragma: no cover  pylint: disable=broad-exception-caught
+    return False
 
 
 def pytest_collection_modifyitems(config, items):

tests/unit/sharding_compare_test.py

@@ -33,6 +33,8 @@
 
 Transformer = models.transformer_as_linen
 
+pytestmark = [pytest.mark.cpu_only, pytest.mark.tpu_backend]
+
 
 def compute_checksum(d: dict) -> str:
   """Compute a checksum (SHA256) of a dictionary."""

tests/unit/train_compile_test.py

@@ -28,6 +28,8 @@
 from maxtext.trainers.pre_train.train_compile import main as train_compile_main
 from tests.utils.test_helpers import get_test_config_path
 
+pytestmark = [pytest.mark.tpu_backend]
+
 
 class TrainCompile(unittest.TestCase):
   """Tests for the Ahead of Time Compilation functionality, train_compile.py"""