jax>=0.6.0

[cuda12]
jax[cuda12]

[dev]
commitizen
pytest
pytest-cov
pre-commit

[docs]
sphinx
sphinx-book-theme
myst-parser
nbsphinx
sphinx-autodoc-typehints
sphinx-copybutton
sphinx-math-dollar
ipython
matplotlib
jupyter
