numpy
sympy
pyvista
scipy
tqdm

[dev]
pytest
pre-commit
nbsphinx
pydata-sphinx-theme
jupyterlab
nbval
black
isort

[jax]
jax[cuda]>=0.4.31
diffrax
optax

[torch]
torch
torchdiffeq
