eliminate window_adaptation overhead?

[ismael-mendoza]

I have noticed that when I run the following code snippet twice in a row the amount of time that it takes is reduced by half.

warmup = blackjax.window_adaptation(
    blackjax.nuts, logdensity, progress_bar=False, is_mass_matrix_diagonal=False,
)
(init_states, tuned_params), adapt_info = warmup.run(warmup_key, init_positions, 1000)

Is this a known feature of the window_adaptation warmup? Is it the case that there is some jit compilation that happens under the scenes? Is it possible to eliminate this overhead by somehow jitting warmup.run and reusing it for future chains?

Any insights or suggestions appreciated and sorry if I missed it in the documentation. Thanks!

[junpenglao]

This is a property of JAX, where in the first run it is compiled+run. That's why the second time you run it, it is faster as it is already compiled.

Also, while we did not explicitly jit the warmup.run, we use JAX scan() that "compiles f, so while it can be combined with jit(), it’s usually unnecessary."

[ismael-mendoza]

I didn't know that about scan(), thank you! know I understand