requirements.txt

# Tensorflow.
# Note: when the version of Tensorflow is changed, the version tf_keras must be
# changed in .github/workflows/actions.yml (pip install --no-deps tf_keras).
tensorflow-cpu~=2.18.1;sys_platform != 'darwin'
tensorflow~=2.18.1;sys_platform == 'darwin'
tf2onnx

# Torch.
--extra-index-url https://download.pytorch.org/whl/cpu
torch==2.6.0
torch-xla==2.6.0;sys_platform != 'darwin'

# Jax.
# Pinned to 0.5.0 on CPU. JAX 0.5.1 requires Tensorflow 2.19 for saved_model_test.
# Note that we test against the latest JAX on GPU.
jax[cpu]==0.5.0
flax

# Common deps.
-r requirements-common.txt