Fix tests with CPU backend

Author: stephen-huan

tests/triton_call_test.py

@@ -28,6 +28,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)

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