jax
jaxlib
-r requirements-tpu.txt