absl-py
numpy
xprof >= 2.20
protobuf >= 6

jax-md; sys_platform == 'linux'
jaxpi @ https://github.com/vimarsh6739/jaxpi/archive/6c22cfa9ac17eac4993aebf454b3af5d074fe430.tar.gz

# maxtext can't be installed concurrently, but installing it fixes
# https://github.com/wsmoses/maxtext/archive/bc50722be7d89e4003bd830b80e4ac968be658eb.tar.gz; python_version < "3.12"
# maxtext; python_version < "3.13"

neuralgcm
dm-haiku

gcsfs; sys_platform == 'linux'
requests; sys_platform == 'linux'

# https://github.com/avik-pal/keras/archive/7344964561c15611dd2d734a2bdd6b872a07709d.tar.gz; sys_platform == 'linux'
# https://github.com/wsmoses/keras-benchmarks/archive/9caedf38085b7fd23cef19e75ab9b7ab31ae99b0.tar.gz; sys_platform == 'linux'

# tf_keras; sys_platform == 'linux'
# keras-hub; sys_platform == 'linux'
# keras-cv; sys_platform == 'linux'

# tensorflow-cpu-datasets; sys_platform == 'linux'

# -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
# libtpu-nightly == 0.1.dev20240729; sys_platform == 'linux'
