numpy>=1.24
jax>=0.4.20

[dev]
pytest>=7
pytest-cov
optax>=0.1.7
build

[dev:python_version < "3.11"]
tomli>=2.0

[docs]
sphinx
sphinx-rtd-theme
myst-parser
matplotlib

[train]
optax>=0.1.7
