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

Use jaxlib version guard for triton instead of xla_extension_version

Browse Source 
PiperOrigin-RevId: 534974834
This commit is contained in:
Sharad Vikram 2023-05-24 14:06:10 -07:00 committed by jax authors
parent 6a54ebd031
commit 4fb834b351
1 changed files with 7 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -116,8 +116,13 @@ import jaxlib.gpu_linalg as gpu_linalg # pytype: disable=import-error
xla_extension_version: int = getattr(xla_client, '_version', 0)
import jaxlib.gpu_rnn as gpu_rnn # pytype: disable=import-error
if xla_extension_version >= 154:
import jaxlib.gpu_triton as gpu_triton # pytype: disable=import-error
if jaxlib.version.__version_info__ >= (0, 4, 11):
# TODO(sharadmv): make this unconditional when minimum jaxlib version is
# bumped to 0.4.11
try:
import jaxlib.gpu_triton as gpu_triton # pytype: disable=import-error
except ModuleNotFoundError:
pass
# Version number for MLIR:Python APIs, provided by jaxlib.
mlir_api_version = xla_client.mlir_api_version