mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 03:46:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Pallas/Mosaic GPU] Explicitly disable ops_test on Mosaic GPU pre-Hopper.

Browse Source 
PiperOrigin-RevId: 735744473
This commit is contained in:
Benjamin Chetioui 2025-03-11 07:10:22 -07:00 committed by jax authors
parent b6da46ecda
commit 7fd32ecc04
1 changed files with 3 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
tests/pallas/ops_test.py
View File

@ -283,6 +283,9 @@ class PallasBaseTest(jtu.JaxTestCase):
if (jtu.test_device_matches(["cuda"]) and
not jtu.is_cuda_compute_capability_at_least("8.0")):
self.skipTest("Only works on GPUs with capability >= sm80")
if (jtu.test_device_matches(["cuda"]) and use_mosaic_gpu and
not jtu.is_cuda_compute_capability_at_least("9.0")):
self.skipTest("Mosaic GPU requires capability >= sm90")
super().setUp()