Can DifferentiationInterface be useful for Turing?

[gdalle]

Hi there!
@adrhill and I recently started https://github.com/gdalle/DifferentiationInterface.jl to provide a common interface for automatic differentiation in Julia. We're currently chatting with Lux.jl, Flux.jl and Optimization.jl to see how they can benefit from it, and so my mind went to Turing.jl as another AD power user :)
DifferentiationInterface.jl only guarantees support for functions of the type f(x) = y or f!(y, x) with standard numbers or arrays in and out. Within these restrictions, we are compatible with 13 different AD backends, including the cool kids like Enzyme.jl and even the hipsters like Tapir.jl. Do you think it could come in handy?
Ping @yebai @willtebbutt

[gdalle]

Some issues that seem related to AD support (not sure they are all relevant here):

• ADVI is not up-to-date on the AD side #2186
• Transfer essential/ad.jl to DynamicPPL #2141
• Should we support Tracker.jl? #2000
• ADVI example fails with Zygote AD backend #1988
• transform! not differentiable #1835
• Zygote's compilation scales badly with the number of ~ statements #1754
• Optimizing model with vector-valued observations produces error when data are a LinearAlgebra.Adjoint #1481
• Feature request: allow user to pass gradient tape into sample() function #1419

[yebai]

Turing's current interface to autodiff backends is based on LogDensityProblemsAD -- it is a fairly lightweight package for sharing glue code to various autodiff backends. You can take a look there to see how things might work.

[torfjelde]

Looks great!

Is there a summary of how this is different from AbstractDifferentiation.jl somewhere?:)

[gdalle]

I updated the summary in this issue: JuliaDiff/AbstractDifferentiation.jl#131

[Red-Portal]

Hi, I think DifferentiationInterface could be perfectly useful for AdvancedVI. However, for that, I think it would be useful to add an interface where auxiliary information can be passed to the target function. For example, like this. This should be very useful for any statistical application that uses minibatch subsampling. Apparently, Tapir supports passing auxiliary input to the target function. So it would be nice to have a unifying interface that provides access to that. (It would also be really nice if ReverseDiff precompiled tapes provided this feature, but unfortunately doesn't.)

[gdalle]

Sounds reasonable, I have opened this issue to keep track:

• Multiple arguments / activities? JuliaDiff/DifferentiationInterface.jl#311

[yebai]

Closing this in favour of tpapp/LogDensityProblemsAD.jl#29; Turing will automatically use DI when tpapp/LogDensityProblemsAD.jl#29 is merged.

[penelopeysm]

The PR above (tpapp/LogDensityProblemsAD.jl#29) was closed, so we still aren't using DI.

#2354 though will introduce a dependency on [email protected] through OptimizationBase.

[gdalle]

For the record I tried again with tpapp/LogDensityProblemsAD.jl#39, now that DI has all the features necessary, but @tpapp said he wants to redesign his library first

[tpapp]

If Turing maintainers need this I am happy to merge the above PR as is, and we can address those changes later.

[yebai]

Assuming no known performance penalty exists, I think unifying the interfaces for various AutoDiff backends using DI is beneficial. DI also allows some new AutoDiff libraries (e.g. Enzyme, Mooncake) to focus more on the core functionality and less on supporting multiple interfaces.

[gdalle]

For Enzyme, as stated in DI's README, there can be some performance penalties or even correctness issues that come from not DI allowing multiple arguments. Apart from that and some StaticArrays shenanigans, I believe that most performance penalties between native backends and DI should either not exist or be easily fixable. Please open an issue with an MWE whenever you spot one.

[yebai]

I'd suggest merging the DI PR on the understanding that @gdalle will help fix emerging (minor) performance issues, if any.

[gdalle]

Note that the PR mentioned above (tpapp/LogDensityProblemsAD.jl#39) only implements the use of DI for backends that were not supported by LogDensityProblemsAD. It will be up to that package to gradually remove their existing extensions (if they want) and let DI take over for them.