Basic utilities for TPUs!
Project description
TPU Utils
TL;DR: Please stop cross-region traffic
At the very beginning of your launcher, please add this:
from tpu_utils import prevent_cross_region
...
def main(_):
prevent_cross_region(FLAGS.save_dir, FLAGS.data_dir, ...) # Anything that may be cross region
Installation:
pip install rail_tpu_utils
pip install git+https://github.com/dibyaghosh/tpu_utils.git
Also has:
tpu_utils.initialize_compilation_cache
: Initializes JAX's compilation cache
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
rail_tpu_utils-0.0.3.tar.gz
(3.7 kB
view details)
Built Distribution
File details
Details for the file rail_tpu_utils-0.0.3.tar.gz
.
File metadata
- Download URL: rail_tpu_utils-0.0.3.tar.gz
- Upload date:
- Size: 3.7 kB
- Tags: Source
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/5.1.0 CPython/3.10.13
File hashes
Algorithm | Hash digest | |
---|---|---|
SHA256 | 96835328269ce8be415d8bb36d67ae8910ae802255e4e714a7ee0c61f7730712 |
|
MD5 | 26e75c69108db80d95bc9c40a66680e1 |
|
BLAKE2b-256 | bce74dfb313c80bd3ace00b7abfda8f1a33c78fe493fb02faf196aef596bae7e |
File details
Details for the file rail_tpu_utils-0.0.3-py3-none-any.whl
.
File metadata
- Download URL: rail_tpu_utils-0.0.3-py3-none-any.whl
- Upload date:
- Size: 4.3 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/5.1.0 CPython/3.10.13
File hashes
Algorithm | Hash digest | |
---|---|---|
SHA256 | b0c4828fe4f487c5f46137564143bc766519155ca5417c43ff23f0a0f21f746f |
|
MD5 | 34d7db4156b20a48d10284d45d4215ec |
|
BLAKE2b-256 | e6498919fcc92690dc57da4c36c818bd0b5606bd8dfc81890a9a137f10e11bf9 |