tpu-info==0.4.0
yapf==0.43.0
pytest
pytest-mock
absl-py
numpy
google-cloud-storage
jax==0.7.2
jaxlib==0.7.2
libtpu==0.0.23
jaxtyping
flax==0.11.1
torchax==0.0.7
qwix==0.1.1
torchvision==0.23.0
pathwaysutils
parameterized
