Fix tests with CPU backend

jax-ml · Jan 5, 2025 · 0260213 · 0260213
1 parent ed25ada
commit 0260213
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 1 deletion.
diff --git a/tests/triton_call_test.py b/tests/triton_call_test.py
@@ -29,6 +29,12 @@
 
 config.parse_flags_with_absl()
 
+try:
+  jt.get_compute_capability(0)
+except AttributeError:
+  # TODO(stephen-huan): add in jaxlib
+  jt.get_compute_capability = lambda _: np.inf
+
 
 def setUpModule():
   config.update("jax_enable_x64", True)

diff --git a/tests/triton_test.py b/tests/triton_test.py
@@ -20,7 +20,10 @@
 import numpy as np
 import triton
 import triton.language as tl
-from triton.language.extra.cuda import libdevice
+try:
+  from triton.language.extra.cuda import libdevice
+except ImportError:
+  from triton.language.extra.cpu import libdevice
 
 
 @triton.jit