Merge pull request #26509 from andportnoy:aportnoy/pallas-mosaic-gpu-…

…test-sm90a

PiperOrigin-RevId: 726624339

jax/_src/test_util.py

@@ -472,6 +472,19 @@ def is_cuda_compute_capability_equal(capability: str) -> bool:
   current = tuple(int(x) for x in d.compute_capability.split("."))
   return current == target
 
+
+class CudaArchSpecificTest:
+  """A mixin with methods allowing to skip arch specific tests."""
+
+  def skip_unless_sm90a(self):
+    if not is_cuda_compute_capability_equal("9.0"):
+      self.skipTest("Only works on GPU with capability sm90a")
+
+  def skip_unless_sm100a(self):
+    if not is_cuda_compute_capability_equal("10.0"):
+      self.skipTest("Only works on GPU with capability sm100a")
+
+
 def _get_device_tags():
   """returns a set of tags defined for the device under test"""
   if is_device_rocm():