JAX speed in iterations after the warm up

[Houman-HM]

Dear all,

I have vmapped a jitted function and would want to call it 100 times.
For some reason, after a few iterations, the computation time increases dramatically. As you can see in the timing reported below, after the warm-up, the computation time is in order of 1e-4 till iteration 16 and then it increases to the order of 1e-2.

I was wondering if somebody could explain why this happens and if there's any way to avoid this increase in computation time.

Best,
Houman

Warm-up time in seconds: 7.804678440093994
Iteration 1, time taken in seconds: 0.0008289813995361328
Iteration 2, time taken in seconds: 0.0006806850433349609
Iteration 3, time taken in seconds: 0.0006709098815917969
Iteration 4, time taken in seconds: 0.0006458759307861328
Iteration 5, time taken in seconds: 0.001438140869140625
Iteration 6, time taken in seconds: 0.0006430149078369141
Iteration 7, time taken in seconds: 0.0005834102630615234
Iteration 8, time taken in seconds: 0.0005719661712646484
Iteration 9, time taken in seconds: 0.0005524158477783203
Iteration 10, time taken in seconds: 0.0005459785461425781
Iteration 11, time taken in seconds: 0.0019834041595458984
Iteration 12, time taken in seconds: 0.0005753040313720703
Iteration 13, time taken in seconds: 0.0005559921264648438
Iteration 14, time taken in seconds: 0.0005142688751220703
Iteration 15, time taken in seconds: 0.0005085468292236328
Iteration 16, time taken in seconds: 0.0005071163177490234
Iteration 17, time taken in seconds: 0.08022618293762207
Iteration 18, time taken in seconds: 0.08780765533447266
Iteration 19, time taken in seconds: 0.08937644958496094
Iteration 20, time taken in seconds: 0.08553099632263184
Iteration 21, time taken in seconds: 0.08857059478759766
Iteration 22, time taken in seconds: 0.08785057067871094
Iteration 23, time taken in seconds: 0.08759570121765137
Iteration 24, time taken in seconds: 0.08722472190856934
Iteration 25, time taken in seconds: 0.08829665184020996
Iteration 26, time taken in seconds: 0.0888204574584961
Iteration 27, time taken in seconds: 0.08738422393798828
Iteration 28, time taken in seconds: 0.08774065971374512
Iteration 29, time taken in seconds: 0.08782577514648438

[jakevdp]

I suspect you're timing the iterations without calling .block_until_ready(), so the first few runs are just measuring the dispatch time rather than the actual computation time. See https://jax.readthedocs.io/en/latest/async_dispatch.html and https://jax.readthedocs.io/en/latest/faq.html#benchmarking-jax-code for more information.

[Houman-HM]

Thanks a lot Jake for the quick reply.
It was indeed the dispatch time.