How to extract warmup samples?

[Madhav-Kanda]

Is it possible to extract warmup samples in Blackjax?

[Madhav-Kanda]

It was simple to get it. Mentioning it here in case someone is not able to figure it out:

Refer to the example provided on the blackjax site for window adaptation based on stan. To obtain the adaptation info (warmup samples) modify the code as follows:

warmup = blackjax.window_adaptation(blackjax.nuts, logdensity)
(state, parameters), warmup_samples = warmup.run(rng_key, initial_position, num_steps=1000)

In the above code, the variable named warmup_samples will contain the warmup samples.