Choosing AD backend in Turing integration

[weinstockj]

Does Pathfinder interact with the AD backend chosen in Turing, ie.., Turing.setadbackend ? As far as I can tell from the source, there is no way to pass a specific AD backend when using the Turing integration - it seems like Forwarddiff is used as the default.

P.S. Thanks for developing Pathfinder.jl ! Great to have Pathfinder in the Julia ecosystem.

[sethaxen]

Does Pathfinder interact with the AD backend chosen in Turing, ie.., Turing.setadbackend ?

Good question! We use Turing's implementation of either optim_function or optim_problem (both are exported but not yet documented). The code for these seems somewhat complicated, but I don't think it uses Turing's AD backend. Instead, they seem to use Optimization.jl's own AD backend machinery. So in the Turing example, to use ReverseDiff, this should work:

using ReverseDiff, Optimization
fun = optim_function(model, MAP(); constrained=false, autoad=Optimization.AutoReverseDiff())
dim = length(fun.init())
pathfinder(fun.func; dim)

Currently there's not a way to call pathfinder(model; autoad), but that's a kwarg we could certainly add support for.

In addition to Turing's Optimization.jl integration, it has specifically Optim.jl integration for its mode estimation functionality, and this seems to use Turing's own AD backend. We could explore using this machinery when an Optim.jl optimizer is selected (which is the default, LBFGS, and the only non-experimental choice).

Thanks for developing Pathfinder.jl ! Great to have Pathfinder in the Julia ecosystem.

You're welcome!

[weinstockj]

Got it - thanks so much!

[SamuelBrand1]

I also want to say thanks for creating Pathfinder!

Just want to bump this, because as of (I think) last month Turing seems to have made it easier to interface with an AD choice.

For example, optim_problem now takes an adtype argument. This is where the PathfinderTuringExt interfaces with Turing I think here.

So now I think you can just add an adtype kwarg to give to optim_problem in pathfinder(mdl::Model,...).

[sethaxen]

Thanks @SamuelBrand1 for the bump! It seems LogDensityProblemsAD and Turing (as well as Optimization.jl) now use the ADTypes interface for specifying AD, so that's a good reason to adopt this as a keyword, both when calling Pathfinder on a Turing model but also when a user passes a LogDensityProblem.

[SamuelBrand1]

@sethaxen That sounds great. For models with lots of parameters this will be super handy.

[sethaxen]

@SamuelBrand1 #178 adds an adtype keyword to pathfinder and multipathfinder for use for Turing models, LogDensityProblems, and log-density functions. Still need to update the docstrings and add a few more tests, but would you like to test it out and make sure it works fine for your models?

[SamuelBrand1]

Thanks @sethaxen !

I'm currently working on Epi inference project/package-to-be with @seabbs and @zsusswein so this is a good chance to try it out.