jax
matplotlib
numpy
