jax>=0.7.1
numpy

[dev]
pytest
black
pyright

[test]
pytest
