The first benchmark is wrong. Numpy random by default uses float64, and jax uses float32.

    In [9]: x = np.random.randn(10000,10000).astype('f')
    In [10]: %timeit -n5 fn(x)
    623 ms ± 9.31 ms per loop (mean ± std. dev. of 7 runs, 5 loops each)

With Jax

    In [17]: %timeit -n5 jax_fn(x).block_until_ready()
    98.4 ms ± 1.55 ms per loop (mean ± std. dev. of 7 runs, 5 loops each)

It's still 6x improvement but not as large that the article claims. I am on the latest Intel Mac