numpy
scipy
jaxlib
jax>=0.2.6
