jax
jaxlib
matplotlib
numpy
scipy
