from jax import jitx = np.random.rand(1000,1000)y = jnp.array(x)def f(x): for _ in range(10): x = 0.5*x + 0.1* jnp.sin(x) return xg = jit(f)%timeit -n 5 -r 5 f(y).block_until_ready()# 5 loops, best of 5: 10.8 ms per loop%timeit -n 5 -r 5 g(y).block_until_ready()# 5 loops, best of 5: 341 µs per loop