numpy
matplotlib
qutip
jax[cpu]
diffrax
flax
tqdm
optax
jax-tqdm

[dev]
jupyterlab>=3.1.0
mypy
pylint
black
coverage

[docs]
mkdocs
mkdocs-material
mkdocs-literate-nav
mkdocs-section-index
mkdocs-gen-files
mkdocstrings-python
mkdocs-jupyter
pymdown-extensions

[gpu]
jax[cuda12]

[tests]
pytest
pytest-cov

[tpu]
jax[tpu]
