jax
equinox
matplotlib
setuptools_scm>=8.1
diffrax

[all]
jax[cuda12]
ruff
pytest
coverage
pytest-cov
notebook
ipykernel
pytest-env
mkdocs
mkdocs-material
mkdocs-autorefs
mkdocstrings[python]
mkdocs-jupyter

[dev]
ruff
pytest
coverage
pytest-cov
pytest-env
mkdocs
mkdocs-material
mkdocs-autorefs
mkdocstrings[python]

[docs]
mkdocs
mkdocs-material
mkdocs-autorefs
mkdocstrings[python]
mkdocs-jupyter

[gpu]
jax[cuda12]

[notebooks]
notebook
ipykernel
