yeah it looks exactly like that file but replace "import numpy as np" with "import jax.numpy as np" :)