Following the outline in https://github.com/blackjax-devs/blackjax/issues/171. The reason we want to do that is to be able to interleave sampling and warmup steps when not all variables are sampled with the HMC/NUTS sampler; the current API to the warmup just runs the warmup for a set number of steps.
Be more specific
This also means we could separate step functions and adaptation steps (that take chain state as an input)?
Following the outline in https://github.com/blackjax-devs/blackjax/issues/171. The reason we want to do that is to be able to interleave sampling and warmup steps when not all variables are sampled with the HMC/NUTS sampler; the current API to the warmup just runs the warmup for a set number of steps.
Be more specific
This also means we could separate step functions and adaptation steps (that take chain state as an input)?