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

Merge pull request #18108 from google:libtpu_import_fix

Browse Source 
PiperOrigin-RevId: 573302310
This commit is contained in:
jax authors 2023-10-13 13:14:56 -07:00
commit ce4d0d2b2d
1 changed files with 3 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
jax/_src/xla_bridge.py
View File

@ -128,6 +128,9 @@ def tpu_client_timer_callback(timer_secs: float) -> Optional[xla_client.Client]:
if xla_extension_version >= 205:
client = xla_client.make_tpu_client(_get_tpu_library_path()) # type: ignore
else:
libtpu_module = maybe_import_libtpu()
if libtpu_module is not None:
libtpu_module.configure_library_path()
client = xla_client.make_tpu_client() # type: ignore
finally:
t.cancel()