Thanks for the suggestion. I think there is a more general issue of choosing better strategies for initialization, e.g. using ADVI, MAP etc., that we would like to explore at some point. The ability to control the initial trace to begin sampling from could be more generally useful. I added this to the issue. The option to provide the initial trace should be a small change, and I will add it soon. In the meantime, you could intercept the
initial_trace method in
HMC to unblock yourself.