aesara-devs / aehmc

An HMC/NUTS implementation in Aesara

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

The warmup should be exposed as a kernel

rlouf opened this issue · comments

Following the outline in blackjax-devs/blackjax#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)?