@shizejin reported that there is a misleading comparison in the introduction to JAX:

It might relates to the timing magic we used. %time is sometimes influenced other factors such as garbage collections, workload etc:
https://stackoverflow.com/questions/17579357/time-time-vs-timeit-timeit

I think %timeit might be a more robust way to compare the results, which is also used in JAX documents:
https://docs.jax.dev/en/latest/quickstart.html

The result under %timeit can be found here: