jax
jaxlib
chex
matplotlib
