Implementing adtype for closely linked "external" samplers / adding metric warm-up

[SamuelBrand1]

Hi everyone,

Being able to add external samplers is a really great feature of Turing, however, it seems that its at odds with the shift towards setting the AD backend via selecting a type from ADTypes as a kwarg to the sampler?

For example, in the documentation you give an example of treating the AdvancedHMC.NUTS wrapper to HMC as an external sampler for the purpose of using Pathfinder.jl to pre-heat the metric https://turinglang.org/docs/tutorials/docs-16-using-turing-external-samplers/#going-beyond-the-turing-api

However, this is the same sampler that is used by Turing.NUTS, albeit without the syntactic sugar of having the adtype kwarg, which implements AD backend with the LogDensityProblemsAD.jl interface here (I think...).

Obviously, I could do this with full AdvancedHMC but wouldn't it be preferable to either include adtype as an option for AdvancedHMC.NUTS or include setting the metric (rather than just metric type) in Turing.NUTS. At the moment, its a bit clunky with how you implement whichever options.

[SamuelBrand1]

Actually I think externalsampler does this already e.g.

https://github.com/TuringLang/Turing.jl/blob/ebc36af633f103e4a25050eb068c08bc14d10de9/src/mcmc/Inference.jl#L126C1-L140C4

Maybe make that clearer in the docs?

[torfjelde]

Actually I think externalsampler does this already e.g.

Aye, this was a recent "fix" so haven't made it's way to the docs yet.

And more generally wrt. the "overarching" question: all the work on externalsampler and other aspects are parts of a larger push to completely remove the Turing.NUTS samplers, etc. in favour of the corresponding packages' own samplers. That way we can just focus on making the "true" sampler constructor nice to work with rather than having to deal with discrepancies like what you describe here. So for now, your best bet is just using AdvancedHMC with externalsampler (as you noticed, you can indeed set the adtype), and sometime in the not-too-far-off-future the clunkiness will be gone:)