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

Fixed the signature of the fallback get_compute_capability

Browse Source 
PiperOrigin-RevId: 615338312
This commit is contained in:
Sergei Lebedev 2024-03-13 02:34:00 -07:00 committed by jax authors
parent 187b7aa8e6
commit a7964445e6
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
jax/_src/pallas/triton/__init__.py
View File

@ -22,7 +22,8 @@ try:
get_compute_capability = triton_kernel_call_lib.get_compute_capability
except AttributeError:
def get_compute_capability() -> int:
def get_compute_capability(device) -> int:
del device # Unused.
raise RuntimeError(
"get_compute_capability is not available. Try installing jaxlib with"
" GPU support following instructions in"