jax>=0.6.1
optax
equinox
optimistix
lineax

[docs]
jupytext
jupyter_book<2
sphinx_autodoc_typehints
sphinx_math_dollar

[gpu]
jax[cuda12]>=0.6.1

[optional]
diffrax

[tests]
matplotlib
pytest
