jax>=0.4.0
jaxlib>=0.4.0
matplotlib
