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

Removed obsolete call to libtpu_module.configure_library_path().

Browse Source 
PiperOrigin-RevId: 619340619
This commit is contained in:
Jieying Luo 2024-03-26 16:05:50 -07:00 committed by jax authors
parent 6e0c95585a
commit 4a9c8d1a0a
1 changed files with 4 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
jax/_src/xla_bridge.py
View File

@ -129,17 +129,11 @@ def _get_tpu_library_path() -> str | None:
libtpu_module = maybe_import_libtpu()
if libtpu_module is not None:
if hasattr(libtpu_module, "get_library_path"):
if xla_extension_version < 212:
# xla_extension_version < 212 uses tpu_tracer which requires calling
# configure_library_path.
libtpu_module.configure_library_path()
return libtpu_module.get_library_path()
else:
# TODO(b/305803029): Remove this branch around 01/2024 after the oldest
# supported TPU has get_library_path.
if xla_extension_version < 212:
# xla_extension_version < 212 uses tpu_tracer which requires calling
# configure_library_path.
libtpu_module.configure_library_path()
return os.getenv("TPU_LIBRARY_PATH", None)
return libtpu_module.get_library_path()
return None