Jax is unrolling for-loops during compilation. I don't know about the hypothetical performance of canonical Jax code for this example, but the comparison as it is seems unfair.

https://jax.readthedocs.io/en/latest/faq.html#jit-decorated-...