tpu-info==0.7.1
yapf==0.43.0
pytest
pytest-mock
absl-py
numpy
google-cloud-storage
jaxtyping
flax==0.12.4
torchax==0.0.11
qwix==0.1.2
torchvision==0.24.0
pathwaysutils
parameterized
numba==0.62.1
runai-model-streamer[s3,gcs]==0.15.4
gcsfs==2026.1.0

--find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html
--find-links https://us-python.pkg.dev/ml-oss-artifacts-published/jax-public-nightly-artifacts-registry/simple/jax
--find-links https://us-python.pkg.dev/ml-oss-artifacts-published/jax-public-nightly-artifacts-registry/simple/jaxlib
jax[tpu]==0.8.3
jaxlib==0.8.3
libtpu==0.0.37.dev20260304
