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.2.tar.gz
(3.7 kB
view details)
Built Distribution
File details
Details for the file rail_tpu_utils-0.0.2.tar.gz
.
File metadata
- Download URL: rail_tpu_utils-0.0.2.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 | 72d86454b609ec479003bee920083df3328474df35b08b219c792c10256d9128 |
|
MD5 | 2966a934c19240db6956616a1a89ffd2 |
|
BLAKE2b-256 | da9ec136221a8a2ea8f00cf3ec02269c42b74c3e17c350b1a0a10313196b56f4 |
File details
Details for the file rail_tpu_utils-0.0.2-py3-none-any.whl
.
File metadata
- Download URL: rail_tpu_utils-0.0.2-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 | 7a01b947f58df0096b21e9ff2451b33a08e7241d5a81c06051b37d751986d602 |
|
MD5 | 46c9ece30089db4a12ab316423c5c9c2 |
|
BLAKE2b-256 | 08b9d4f79cb231192852df12cedabce77cc883ace5a68e5427fe3cb5317cc26a |