jax
jaxlib
numpy
optax

[dev]
pytest
