Metadata-Version: 2.1
Name: rail_tpu_utils
Version: 0.0.2
Summary: Basic utilities for TPUs!
Author-email: Dibya Ghosh <dibya.ghosh@berkeley.edu>
Project-URL: Homepage, https://github.com/dibyaghosh/tpu_utils
Classifier: Programming Language :: Python :: 3
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: OS Independent
Requires-Python: >=3.8
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: requests
Requires-Dist: google-cloud-storage

# TPU Utils


## TL;DR: Please stop cross-region traffic

At the very beginning of your launcher, please add this:

```python
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
