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

Update Colab TPU driver version

Browse Source
This commit is contained in:
Skye Wanderman-Milne 2022-10-11 17:49:10 -07:00
parent 6b459f5ebf
commit e2aa939147
2 changed files with 4 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
CHANGELOG.md
View File

@ -9,6 +9,8 @@ PLEASE REMEMBER TO CHANGE THE '..main' WITH AN ACTUAL TAG in GITHUB LINK.
--> -->
## jax 0.3.23 ## jax 0.3.23
* Changes
* Update Colab TPU driver version for new jaxlib release.
## jaxlib 0.3.23 ## jaxlib 0.3.23

5
jax/tools/colab_tpu.py
View File

@ -22,7 +22,7 @@ from jax.config import config
TPU_DRIVER_MODE = 0 TPU_DRIVER_MODE = 0
def setup_tpu(tpu_driver_version='tpu_driver-0.2'): def setup_tpu(tpu_driver_version='tpu_driver_20221011'):
"""Sets up Colab to run on TPU. """Sets up Colab to run on TPU.
Note: make sure the Colab Runtime is set to Accelerator: TPU. Note: make sure the Colab Runtime is set to Accelerator: TPU.
@ -30,8 +30,7 @@ def setup_tpu(tpu_driver_version='tpu_driver-0.2'):
Args Args
---- ----
tpu_driver_version : (str) specify the version identifier for the tpu driver. tpu_driver_version : (str) specify the version identifier for the tpu driver.
Defaults to "tpu_driver-0.2", which can be used with jaxlib 0.3.20. Set to Set to "tpu_driver_nightly" to use the nightly tpu driver build.
"tpu_driver_nightly" to use the nightly tpu driver build.
""" """
global TPU_DRIVER_MODE global TPU_DRIVER_MODE