Save gradients

[JaimeRZP]

Hi @sethaxen!

This should do what you were suggesting if you use the internal sampling method of AHMC and set drop_warmup=false and keep_gradient=true.

let me know if it works!

[torfjelde]

Generally speaking, it's not great to have the return-value change based on a keyword argument , as it leads to type-instabilities :confused: In this scenario, it's probably the not the worst (sample is generally going to be the outermost caller anyways), but it's still not great style.

IMO what we should do here is to

1. Move AdvancedHMC.jl completely over to the AbstractMCMC.jl interface, i.e. we have step function that takes in a state containing all the necessary information.
2. Then the gradients can easily be extracted through a callback which just extracts this information from the state.

(1) will also just be a huge gain in general, as it will make things much more modular:)