tpu-info==0.4.0  # pinning to 0.4.0 as newer version leads to segmentation fault on TPU https://github.com/vllm-project/vllm/pull/23135
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.7
qwix==0.1.1
torchvision==0.23.0
pathwaysutils
parameterized
numba==0.62.1
