setuptools
wheel
numpy
jax<=0.4.20,>=0.4.1
jaxlib<=0.4.20,>=0.4.1
flax<=0.6.11,>=0.6.4
mpi4py
h5py
PyYAML
matplotlib

[dev]
setuptools
wheel
numpy
jax<=0.4.20,>=0.4.1
jaxlib<=0.4.20,>=0.4.1
flax<=0.6.11,>=0.6.4
mpi4py
h5py
PyYAML
matplotlib
sphinx
mock
sphinx_rtd_theme
pytest
pytest-mpi
