Skip test_dot_precision() with TF32 for Navi architecture (#520)

RDNA3 / Navi3x / gfx11xx and RDNA4 / Navi4x / gfx12xx don't support TF32/XF32 hence relevant `tests/pallas/pallas_test.py::PallasCallInterpretTest::test_dot_precision*` tests needs to be skipped.

Author: Arech8

jax/_src/test_util.py

@@ -395,6 +395,20 @@ def supported_dtypes():
 def is_device_rocm():
   return 'rocm' in xla_bridge.get_backend().platform_version
 
+def any_rocm_device_has_gfx_prefix(gfx_prefixes: str | tuple[str])->bool:
+  """Returns true if any available device has gfx value that starts with any of
+  prefixes in gfx_prefixes.
+  Intended use is to skip tests that require certain features not present in
+  devices."""
+  if not is_device_rocm():
+    return False
+  if isinstance(gfx_prefixes, str):
+    gfx_prefixes = (gfx_prefixes,)
+  assert isinstance(gfx_prefixes, tuple), "argument must be a tuple of strings"
+  assert all([isinstance(e, str) for e in gfx_prefixes])
+  gfxs = frozenset([d.compute_capability for d in jax.devices()])
+  return any([g.startswith(gfx_prefixes) for g in gfxs])
+
 def get_rocm_version():
   rocm_path = os.environ.get("ROCM_PATH", "/opt/rocm")
   version_path = Path(rocm_path) / ".info" / "version"

tests/pallas/pallas_test.py

@@ -721,6 +721,12 @@ def test_dot_precision(self, dtype, precision):
     if not jtu.test_device_matches(["gpu"]):
       self.skipTest("`DotAlgorithmPreset` only supported on GPU.")
 
+    if precision in (
+        jax.lax.DotAlgorithmPreset.TF32_TF32_F32,
+        jax.lax.DotAlgorithmPreset.TF32_TF32_F32_X3,
+      ) and jtu.any_rocm_device_has_gfx_prefix(("gfx11","gfx12")):
+      self.skipTest("Navi3x and Navi4x doesn't have hardware support for TF32")
+
     @functools.partial(
         self.pallas_call,
         out_shape=jax.ShapeDtypeStruct((32, 64), jnp.float32),