ReverseDiff + callable `LogDensityAt` struct doesn't like accumulators

[penelopeysm]

Accumulators: #885

LogDensityAt struct: #922

ReverseDiff doesn't like it (check out the breaking branch where these two PRs were merged):

using DynamicPPL, Distributions
using DynamicPPL.TestUtils.AD: run_ad
using ADTypes: AutoReverseDiff
import ReverseDiff

@model f() = x ~ Normal()

run_ad(f(), AutoReverseDiff())

errors with

MethodError: no method matching seeded_reverse_pass!(::DiffResults.MutableDiffResult{1, Float64, Tuple{Vector{Float64}}}, ::@NamedTuple{logprior::ReverseDiff.TrackedReal{Float64, Float64, Nothing}, loglikelihood::ReverseDiff.TrackedReal{Float64, Float64, Nothing}}, ::ReverseDiff.TrackedArray{Float64, Float64, 1, Vector{Float64}, Vector{Float64}}, ::ReverseDiff.GradientTape{DynamicPPL.LogDensityAt{Model{typeof(DynamicPPL.TestUtils.demo_assume_matrix_observe_matrix_index), (:x, Symbol("##arg#500")), (), (), Tuple{Transpose{Float64, Matrix{Float64}}, DynamicPPL.TypeWrap{Array{Float64}}}, Tuple{}, DefaultContext}, SimpleVarInfo{@NamedTuple{s::Matrix{Float64}, m::Vector{Float64}}, DynamicPPL.AccumulatorTuple{3, @NamedTuple{LogPrior::LogPriorAccumulator{Float64}, LogLikelihood::LogLikelihoodAccumulator{Float64}, NumProduce::NumProduceAccumulator{Int64}}}, DynamicPPL.DynamicTransformation}, DefaultContext}, ReverseDiff.TrackedArray{Float64, Float64, 1, Vector{Float64}, Vector{Float64}}, @NamedTuple{logprior::ReverseDiff.TrackedReal{Float64, Float64, Nothing}, loglikelihood::ReverseDiff.TrackedReal{Float64, Float64, Nothing}}})
  The function `seeded_reverse_pass!` exists, but no method is defined for this combination of argument types.

Replacing LogDensityAt with the old closure, x -> logdensity_at(x, ...) fixes the error.

Another solution is to make ReverseDiff not use the closure, i.e. declare use_closure(::AutoReverseDiff) = false. (In fact that makes use_closure useless, because ReverseDiff was the only one that was using the closure)

These are the most direct solutions, but I'm not sure if there's a better one (I need to benchmark).

And in any case this seems like an upstream issue too (although my hopes of getting it fixed are ... slim, so we had better find a workaround here).

[penelopeysm]

Benchmarks

Code

using DynamicPPL, Distributions
using DynamicPPL.TestUtils.AD: run_ad
using ADTypes: AutoReverseDiff
import ReverseDiff
@model function gdemo()
    s ~ InverseGamma(2, 3)
    m ~ Normal(0, sqrt(s))
    1.5 ~ Normal(m, sqrt(s))
    2.0 ~ Normal(m, sqrt(s))
end
res = run_ad(gdemo(), AutoReverseDiff(); benchmark=true)
println(res.time_vs_primal)

solution (1a): use x -> logdensity_at(...) + disable DI strict preparation -- 23.167792525204522

solution (1b): use x -> logdensity_at(...) + store the closure with the LDF -- 30.59077966101695

solution (2): don't use closure (i.e. use DI.Constant) -- 25.35534728033473

For comparison, current main branch (which has the callable struct, but no accumulators) - 24.83905676855895

Unfortunately, we can't use DI.Constant because that kills performance for compiled ReverseDiff -- if there are any DI contexts, it doesn't compile the tape.

[penelopeysm]

Turns out it's because I needed to replace getlogp (which now returns a NamedTuple) with getlogjoint. So it's fixed properly now.

I took the time to re-benchmark use_closure for our main AD backends. Here, false or true refers to what use_closure returns, i.e. the false column uses DI.Constant and the true column uses the LogDensityAt callable struct.

It seems that for FD, non-compiled RD, and Mooncake, the performance is basically the same.
For compiled RD, we need to use the closure (otherwise DI won't do preparation).

So it seems like we could just always use the closure. The only other AD backend we might care about is Enzyme, but Enzyme is currently broken with accumulators, so it's not something to worry about right now.

                                                                 FALSE        TRUE         false/true   mean(false/true)
demo_dot_assume_observe                     ForwardDiff          1.266874492  1.254608767  1.009776534  1.003714199
demo_assume_index_observe                   ForwardDiff          1.459917673  1.367416496  1.067646673  
demo_assume_multivariate_observe            ForwardDiff          1.285862044  1.581378249  0.813127438  
demo_dot_assume_observe_index               ForwardDiff          1.27787643   1.184032984  1.07925746   
demo_assume_dot_observe                     ForwardDiff          1.524875053  1.399566706  1.089533672  
demo_assume_multivariate_observe_literal    ForwardDiff          1.258994998  1.404267475  0.896549283  
demo_dot_assume_observe_index_literal       ForwardDiff          1.278971555  1.247331773  1.025365972  
demo_assume_dot_observe_literal             ForwardDiff          1.480135851  1.438707779  1.028795335  
demo_assume_observe_literal                 ForwardDiff          1.478736785  1.426797195  1.036402924  
demo_assume_submodel_observe_index_literal  ForwardDiff          1.191640741  1.222563621  0.974706526  
demo_dot_assume_observe_submodel            ForwardDiff          1.346320781  1.304514421  1.032047449  
demo_dot_assume_observe_matrix_index        ForwardDiff          1.310380107  1.334056257  0.98225251   
demo_assume_matrix_observe_matrix_index     ForwardDiff          1.315570406  1.298914678  1.012822804  
demo_dot_assume_observe                     Mooncake             5.011004835  5.148836246  0.97323057   0.978081325
demo_assume_index_observe                   Mooncake             4.537008093  4.588557637  0.988765632  
demo_assume_multivariate_observe            Mooncake             4.758504321  5.69036962   0.836238178  
demo_dot_assume_observe_index               Mooncake             5.004883843  4.982416051  1.004509417  
demo_assume_dot_observe                     Mooncake             5.982891476  5.758294883  1.03900401   
demo_assume_multivariate_observe_literal    Mooncake             5.101830508  5.669230769  0.899915829  
demo_dot_assume_observe_index_literal       Mooncake             5.034893617  4.710705554  1.068819428  
demo_assume_dot_observe_literal             Mooncake             5.964078738  5.45895746   1.092530723  
demo_assume_observe_literal                 Mooncake             5.586726933  5.358756772  1.042541614  
demo_assume_submodel_observe_index_literal  Mooncake             3.862682531  3.981818182  0.970080088  
demo_dot_assume_observe_submodel            Mooncake             3.566400751  4.206733049  0.847783948  
demo_dot_assume_observe_matrix_index        Mooncake             5.039647577  5.246572197  0.960560036  
demo_assume_matrix_observe_matrix_index     Mooncake             5.320758701  5.368659227  0.991077749  
demo_dot_assume_observe                     ReverseDiff          21.57567653  21.5603936   1.000708843  1.019388649
demo_assume_index_observe                   ReverseDiff          20.42826291  20.4515226   0.998862692  
demo_assume_multivariate_observe            ReverseDiff          21.58153846  24.14689577  0.893760369  
demo_dot_assume_observe_index               ReverseDiff          22.60257391  22.65236052  0.997802145  
demo_assume_dot_observe                     ReverseDiff          26.04329004  24.18965517  1.076629239  
demo_assume_multivariate_observe_literal    ReverseDiff          22.08626087  22.64074667  0.975509386  
demo_dot_assume_observe_index_literal       ReverseDiff          23.01459414  21.8354754   1.05400014   
demo_assume_dot_observe_literal             ReverseDiff          24.15617651  22.60341531  1.068695867  
demo_assume_observe_literal                 ReverseDiff          25.16919671  24.00683784  1.048417825  
demo_assume_submodel_observe_index_literal  ReverseDiff          18.58572414  17.71515789  1.049142449  
demo_dot_assume_observe_submodel            ReverseDiff          16.94370157  16.63714286  1.018426163  
demo_dot_assume_observe_matrix_index        ReverseDiff          22.15536105  21.40032     1.035281764  
demo_assume_matrix_observe_matrix_index     ReverseDiff          22.49583362  21.73897912  1.034815549  
demo_dot_assume_observe                     ReverseDiffCompiled  21.42937588  5.868246835  3.651750937  3.524218504
demo_assume_index_observe                   ReverseDiffCompiled  20.45152838  6.017974528  3.39840727   
demo_assume_multivariate_observe            ReverseDiffCompiled  21.21399594  6.613756614  3.207556187  
demo_dot_assume_observe_index               ReverseDiffCompiled  23.667       6.852001149  3.454027442  
demo_assume_dot_observe                     ReverseDiffCompiled  25.55202716  7.896393969  3.235910881  
demo_assume_multivariate_observe_literal    ReverseDiffCompiled  21.59030837  6.19646251   3.484295812  
demo_dot_assume_observe_index_literal       ReverseDiffCompiled  23.33264859  6.174416758  3.778923501  
demo_assume_dot_observe_literal             ReverseDiffCompiled  24.23076923  6.574670202  3.685472957  
demo_assume_observe_literal                 ReverseDiffCompiled  24.94915254  7.249468808  3.441514572  
demo_assume_submodel_observe_index_literal  ReverseDiffCompiled  18.11171475  4.91184377   3.687355624  
demo_dot_assume_observe_submodel            ReverseDiffCompiled  15.2266      4.688728217  3.247490427  
demo_dot_assume_observe_matrix_index        ReverseDiffCompiled  21.92737488  6.019529412  3.64270583   
demo_assume_matrix_observe_matrix_index     ReverseDiffCompiled  22.31974492  5.723849372  3.899429119  

[penelopeysm]

Enzyme with closure now seems to need AutoEnzyme(; ..., function_annotation = Enzyme.Duplicated), or else it refuses to even begin differentiating. AutoEnzyme(; ..., function_annotation = Enzyme.DuplicatedNoNeed) also seems to work. Enzyme without closure works too.

Actually AutoEnzyme(; ..., function_annotation = Enzyme.Const) works as well.

(In this context, 'works' means that it hits the NamedTuple deepcopying error which I reported here EnzymeAD/Enzyme.jl#2429.)