jax>=0.4.20
equinox>=0.11.0
matplotlib
diffrax>=0.6.0

[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
