tpu-info==0.7.1
yapf==0.43.0
pytest
pytest-mock
absl-py
numpy
google-cloud-storage
jax[tpu]==0.8.0
jaxlib==0.8.0
jaxtyping
flax==0.11.1
torchax==0.0.10
qwix==0.1.1
torchvision==0.24.0
pathwaysutils
parameterized
numba==0.62.1
runai-model-streamer[gcs,s3]==0.15.0
