Merge pull request #3355 from ROCm:fix-tpu-backend-skips

PiperOrigin-RevId: 881535168

Author: Google-ML-Automation

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

@@ -110,7 +110,8 @@ def compare_sharding_jsons(json1: dict, model1_name: str, json2: dict, model2_na
 
 
 # Requires JAX TPU support to generate the simulated TPU topology.
-@pytest.mark.tpu_only
+@pytest.mark.cpu_only
+@pytest.mark.tpu_backend
 @pytest.mark.parametrize("model_name, topology, num_slice", TEST_CASES)
 def test_sharding_dump_for_model(model_name: str, topology: str, num_slice: str) -> None:
   """
@@ -214,11 +215,12 @@ def abstract_state_and_shardings(request):
   return model_name, topology, num_slice, abstract_state, state_mesh_shardings, logical_shardings
 
 
+@pytest.mark.cpu_only
+@pytest.mark.tpu_backend
 class TestGetAbstractState:
   """Test class for get_abstract_state function and sharding comparison."""
 
   # Requires JAX TPU support to generate the simulated TPU topology.
-  @pytest.mark.tpu_only
   def test_get_abstract_state_sharding(self, abstract_state_and_shardings):  # pylint: disable=redefined-outer-name
     """Tests that get_abstract_state returns a state with the correct abstract structure and compares sharding."""
 

tests/unit/train_compile_test.py

@@ -29,6 +29,7 @@
 from tests.utils.test_helpers import get_test_config_path
 
 
+@pytest.mark.tpu_backend
 class TrainCompile(unittest.TestCase):
   """Tests for the Ahead of Time Compilation functionality, train_compile.py"""