jax
jaxlib
numpy

[dev]
pytest
