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.1.tar.gz
(3.5 kB
view details)
Built Distribution
File details
Details for the file rail_tpu_utils-0.0.1.tar.gz
.
File metadata
- Download URL: rail_tpu_utils-0.0.1.tar.gz
- Upload date:
- Size: 3.5 kB
- Tags: Source
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/5.1.0 CPython/3.10.13
File hashes
Algorithm | Hash digest | |
---|---|---|
SHA256 | cb5afe98fdf9c449a5e19725ac1cc40b70f87e08eaec26f8ae906aae12f0f8be |
|
MD5 | 24eb94831a7de0ee21a3c5306cdba42e |
|
BLAKE2b-256 | fbd1947142c2beabc8618f26913ef66a02237dc702309da68b0372624ce9aa7e |
File details
Details for the file rail_tpu_utils-0.0.1-py3-none-any.whl
.
File metadata
- Download URL: rail_tpu_utils-0.0.1-py3-none-any.whl
- Upload date:
- Size: 4.0 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 | 21d2f6628cca7cc4b2608b808cfbddd3b9d0d04b1b58bb762f892d5639828100 |
|
MD5 | 1d0b4402bca80cf1f081c7055f4f3771 |
|
BLAKE2b-256 | e07b51445525b94281df7b6558047bab54babff9fbca0603d7ad43d07ce52edb |