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

Fix jax 0.3.11 GPU breakge when used with jaxlib 0.3.10.

Browse Source
This commit is contained in:
Peter Hawkins 2022-05-16 00:24:04 +00:00
parent 1381afc37f
commit 337ec47d13
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
jax/_src/lib/__init__.py
View File

@ -183,7 +183,7 @@ except ImportError:
hip_linalg = None
try:
import jaxlib.cuda_linalg as gpu_linalg # pytype: disable=import-error
import jaxlib.gpu_linalg as gpu_linalg # pytype: disable=import-error
except ImportError:
gpu_linalg = None