weird behavior in step size trace plot during warmup

[ismael-mendoza]

I was running multilple the warmup (window_adaptation) of multiple NUTS chains at once using vmap with different random seeds. I observe this behavior with the step sizes during warmup

And I also see that that at these points there are divergences for all the chains. Is this expected behavior due to the default schedule or is there some weird correlation between my chains? Thank you!

[junpenglao]

IIRC this is intended behavior because the step size at the begining of each window will be reset to default.

[junpenglao]

To be more precise, the step size will be reset to default at each slow adaptation window:

blackjax/blackjax/adaptation/window_adaptation.py

Line 182 in 8a9b546

new_ss_state = da_init(da_final(warmup_state.ss_state))

[ismael-mendoza]

Thanks so much!