#
# runtime deps
#
numpy~=2.2.0; python_version<="3.13"
numpy>=2.3.2; python_version>="3.14"
scipy>=1.15.0; python_version<="3.12"
scipy>=1.15.2; python_version>="3.13"

ml_dtypes>=0.5.3
etils[epath]
opt-einsum

# zstandard is available in the Python 3.14 standard library as
# compression.zstd, and we'll use that if available.
zstandard; python_version<"3.14"

tensorstore

# For jax2tf_test
# TODO(phawkins): re-enable on Python 3.14+ once TensorFlow supports Python 3.14+
tensorflow; python_version<"3.14"

# Needed to build wheels
build
setuptools
wheel

# JAX's own libraries. We include these in the requirements so you can
# bazel test without building jaxlib and without manually updating the
# the requirements files.
jax
jaxlib

jax-cuda12-plugin; sys_platform == "linux"
jax-cuda13-plugin
jax-cuda12-pjrt; sys_platform == "linux"
jax-cuda13-pjrt

# TPU dependencies
libtpu; sys_platform == "linux" and platform_machine == "x86_64"

# Colorama is only needed on windows. Pip-compile ignores platform-specific
# requirements therefore we include it unconditionally here so that the lockfile
# regeneration always includes it.
colorama>=0.4.4