jax
matplotlib
orthax
diffrax