jax>=0.6.0

[dev]
commitizen
pytest
pytest-cov
pre-commit
