fix markers

Author: gulsumgudukbay

tests/unit/sharding_compare_test.py

@@ -33,8 +33,6 @@
 
 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."""
@@ -112,7 +110,7 @@ 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.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:
   """
@@ -216,11 +214,11 @@ def abstract_state_and_shardings(request):
   return model_name, topology, num_slice, abstract_state, state_mesh_shardings, logical_shardings
 
 
+@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

@@ -28,9 +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]
-
 
+@pytest.mark.tpu_backend
 class TrainCompile(unittest.TestCase):
   """Tests for the Ahead of Time Compilation functionality, train_compile.py"""