Bump DifferentiationInterface to 0.7 #922
Conversation
Benchmark Report for Commit 7d75aa0Computer InformationBenchmark Results |
Codecov ReportAll modified and coverable lines are covered by tests ✅
Additional details and impacted files@@ Coverage Diff @@
## main #922 +/- ##
=======================================
Coverage 85.10% 85.11%
=======================================
Files 36 36
Lines 3954 3956 +2
=======================================
+ Hits 3365 3367 +2
Misses 589 589 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
Pull Request Test Coverage Report for Build 15112459001Details
💛 - Coveralls |
|
DynamicPPL.jl documentation for PR #922 is available at: |
|
Nice catch! Hope the debugging wasn't too much trouble, I really think this is a welcome addition to DI |
|
If it is any easier, an alternative would be to replace the closure with a callable struct, where you control what is closed over. With only one argument this would be |
|
@gdalle Thanks for the suggestion! I benchmarked it and the callable struct is actually a bit faster, so that's a win. (Results) results_closure.csv (To reproduce, 6f670c4 is the commit which introduces the struct - run this before and after this commit) using DynamicPPL, Distributions, ForwardDiff, ReverseDiff, Mooncake, Enzyme, ADTypes, DataFrames
using DynamicPPL.TestUtils.AD: run_ad, ADResult
ADTYPES = Dict(
"ForwardDiff" => AutoForwardDiff(),
"ReverseDiff" => AutoReverseDiff(; compile=false),
"ReverseDiffCompiled" => AutoReverseDiff(; compile=true),
"Mooncake" => AutoMooncake(; config=nothing),
"EnzymeForward" => AutoEnzyme(; mode=set_runtime_activity(Forward, true)),
"EnzymeReverse" => AutoEnzyme(; mode=set_runtime_activity(Reverse, true)),
)
MODELS = DynamicPPL.TestUtils.DEMO_MODELS
RESULTS = []
for model in MODELS
for (adtype_name, adtype) in ADTYPES
result = run_ad(model, adtype; benchmark=true)
@show result.time_vs_primal
push!(RESULTS, (model="$(model.f)", adtype=adtype_name, time=result.time_vs_primal))
end
end
RESULT_DF = DataFrame(RESULTS)
using CSV
CSV.write("results.csv", RESULT_DF) |
|
Do the type parameters increase performance even more, by removing type instabilities? |
|
Nope, it's pretty much the same as before (t(struct with parameters) / t(struct without parameters) = 1.004 ± 0.094) |
DifferentiationInterface 0.7 enables a strictness check by default (see JuliaDiff/DifferentiationInterface.jl#799). This introduces a requirement that the function being differentiated must have the same signature at the preparation and execution stage (i.e.,
prepare_gradientandvalue_and_gradientare called with the same function signature).This broke some code inside DynamicPPL, specifically, when the function being differentiated was a closure.
To illustrate, when
use_closure(adtype)istrue, the function being prepared here is an anonymous function of the formDynamicPPL.var"#1#2"{Model, VarInfo, Context}.DynamicPPL.jl/src/logdensityfunction.jl
Lines 126 to 129 in cdeb657
This preparation later gets used here, and there are two problems:
DynamicPPL.jl/src/logdensityfunction.jl
Lines 202 to 214 in cdeb657
Firstly, this closure must close over the
LogDensityFunctionin the argument tologdensity_and_gradient, so it gets a type that looks likeDynamicPPL.var"#3#4"{LogDensityFunction}. (This can be seen in the failing CI runs.)But secondly, even if you were to extricate and construct this closure outside of
logdensity_and_gradient, it would still get a different gensym'd name, so its type would be something likeDynamicPPL.var"#5#6"{Model, VarInfo, Context}, which still does not match the type at preparation.This PR resolves this issue by constructing the closure itself at preparation time, and then storing it inside the LogDensityFunction, so that it doesn't need to be reconstructed. This also enables DynamicPPL to work with DI 0.7, without having to disable the strictness check.