Enable the option to switch between differentiating a closure, and a function with constant arguments

[penelopeysm]

Prior to 'fast' LDF (i.e. in DynamicPPL <= 0.38) there were two different functions that were differentiated through when performing LogDensityFunctions.logdensity_and_gradient: one was a multi-argument function f(x, c1, c2, c3, ...) where x is the active argument, and one was a callable struct F(x) where F closed over the inactive arguments c1, c2, c3. See #806 and #922 for previous history, and the old version of src/logdensityfunction.jl here.

This feature was not ported over to FastLDF in #1139, which is probably an example of laziness on my part. The benefits of doing this is mainly improved performance, but for Enzyme it also allows us to avoid setting function_annotation (see also #1048).

Here are some benchmarks: clos means use the closure, func means the original function. As one can see most of the AD backends are faster with the function rather than the closure, hence _use_closure is mostly set to false. Part of me wonders if sticking an @inline on the callable struct might remove this difference, but my industriousness does have a limit.

# Run this once with `_use_closure` set to true and once with it set to false.

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),
    "MooncakeFwd" => AutoMooncakeForward(),
    "MooncakeRvs" => AutoMooncake(),
)
MODELS = DynamicPPL.TestUtils.DEMO_MODELS

RESULTS = []
for model in MODELS
    for (adtype_name, adtype) in ADTYPES
        result = run_ad(model, adtype; benchmark=true)
        push!(RESULTS, (model="$(model.f)", adtype=adtype_name, time=(result.grad_time / result.primal_time)))
    end
end
RESULT_DF = DataFrame(RESULTS)

using CSV
CSV.write("results.csv", RESULT_DF)

Here time refers to the time for gradient divided by time for primal.

model                                       adtype               time(clos)   time(func)   clos/func    mean(clos/func)  std(clos/func)
demo_dot_assume_observe                     ForwardDiff          1.525467599  1.274712187  1.196715317  1.194930597      0.142726317
demo_assume_index_observe                   ForwardDiff          1.44429541   1.367799257  1.055926448                   
demo_assume_multivariate_observe            ForwardDiff          1.679705404  1.458947368  1.151313228                   
demo_dot_assume_observe_index               ForwardDiff          2.003910968  1.5681335    1.277895644                   
demo_assume_dot_observe                     ForwardDiff          3.160890644  2.269975335  1.392477969                   
demo_assume_multivariate_observe_literal    ForwardDiff          1.526315789  1.463927463  1.042617089                   
demo_dot_assume_observe_index_literal       ForwardDiff          1.490406868  1.100337767  1.354499421                   
demo_assume_dot_observe_literal             ForwardDiff          2.172935945  2.269538142  0.957435306                   
demo_assume_observe_literal                 ForwardDiff          1.739594118  1.306940931  1.331042648                   
demo_assume_submodel_observe_index_literal  ForwardDiff          1.597721311  1.290655327  1.237914785                   
demo_dot_assume_observe_submodel            ForwardDiff          1.497499745  1.131985731  1.322896308                   
demo_dot_assume_observe_matrix_index        ForwardDiff          1.375612421  1.140178153  1.206489019                   
demo_assume_matrix_observe_matrix_index     ForwardDiff          1.403399928  1.393818007  1.006874585                   
demo_dot_assume_observe                     MooncakeFwd          13.71284271  12.28016878  1.116665655  1.117348426      0.138978196
demo_assume_index_observe                   MooncakeFwd          9.806866033  10.16455696  0.964809983                   
demo_assume_multivariate_observe            MooncakeFwd          14.32729477  13.44351464  1.065740258                   
demo_dot_assume_observe_index               MooncakeFwd          12.99918633  10.98645643  1.183201009                   
demo_assume_dot_observe                     MooncakeFwd          8.9453125    6.336938553  1.411614209                   
demo_assume_multivariate_observe_literal    MooncakeFwd          14.36744828  13.78410311  1.042320139                   
demo_dot_assume_observe_index_literal       MooncakeFwd          12.39047587  11.19519899  1.106766917                   
demo_assume_dot_observe_literal             MooncakeFwd          7.90678468   8.648456658  0.914242274                   
demo_assume_observe_literal                 MooncakeFwd          8.14370993   6.063852475  1.342992753                   
demo_assume_submodel_observe_index_literal  MooncakeFwd          12.37424316  11.59245604  1.0674393                     
demo_dot_assume_observe_submodel            MooncakeFwd          13.58812441  11.94641598  1.137422674                   
demo_dot_assume_observe_matrix_index        MooncakeFwd          14.33099908  12.31012     1.164164044                   
demo_assume_matrix_observe_matrix_index     MooncakeFwd          14.38441882  14.26812892  1.008150326                   
demo_dot_assume_observe                     MooncakeRvs          4.975314511  3.923043303  1.268228293  1.161985086      0.1449571
demo_assume_index_observe                   MooncakeRvs          3.276771572  3.252789532  1.007372761                   
demo_assume_multivariate_observe            MooncakeRvs          5.260503186  4.572033898  1.150582717                   
demo_dot_assume_observe_index               MooncakeRvs          4.566469869  3.630591631  1.257775683                   
demo_assume_dot_observe                     MooncakeRvs          5.616897585  4.911655364  1.143585445                   
demo_assume_multivariate_observe_literal    MooncakeRvs          4.942171437  4.519455416  1.093532513                   
demo_dot_assume_observe_index_literal       MooncakeRvs          4.740758124  3.331779452  1.422890738                   
demo_assume_dot_observe_literal             MooncakeRvs          4.985933104  5.047140831  0.987872792                   
demo_assume_observe_literal                 MooncakeRvs          4.938183576  4.422866152  1.1165121                     
demo_assume_submodel_observe_index_literal  MooncakeRvs          4.572583408  3.873515664  1.18047371                    
demo_dot_assume_observe_submodel            MooncakeRvs          4.389248431  3.8731542    1.133249079                   
demo_dot_assume_observe_matrix_index        MooncakeRvs          5.60766652   4.015870574  1.396376307                   
demo_assume_matrix_observe_matrix_index     MooncakeRvs          4.727606079  4.990326923  0.947353981                   
demo_dot_assume_observe                     ReverseDiff          67.84729114  52.74588235  1.28630498   1.178560687      0.130455839
demo_assume_index_observe                   ReverseDiff          58.17253605  54.37883925  1.069764211                   
demo_assume_multivariate_observe            ReverseDiff          66.15684785  59.30163603  1.115599034                   
demo_dot_assume_observe_index               ReverseDiff          71.56960413  63.23089475  1.131877137                   
demo_assume_dot_observe                     ReverseDiff          157.984494   110.0992542  1.434927922                   
demo_assume_multivariate_observe_literal    ReverseDiff          55.93370698  52.77155172  1.059921589                   
demo_dot_assume_observe_index_literal       ReverseDiff          72.94905858  55.06082648  1.324881286                   
demo_assume_dot_observe_literal             ReverseDiff          103.8512397  96.80268542  1.072813623                   
demo_assume_observe_literal                 ReverseDiff          125.2098441  104.3174019  1.200277632                   
demo_assume_submodel_observe_index_literal  ReverseDiff          67.71928065  57.96516872  1.1682754                     
demo_dot_assume_observe_submodel            ReverseDiff          61.50066487  52.23269152  1.177436258                   
demo_dot_assume_observe_matrix_index        ReverseDiff          60.39138037  46.0713928   1.310821677                   
demo_assume_matrix_observe_matrix_index     ReverseDiff          45.24193548  46.71880176  0.968388182                   
demo_dot_assume_observe                     ReverseDiffCompiled  20.90625     51.28954839  0.407612285  0.367591864      0.052660854
demo_assume_index_observe                   ReverseDiffCompiled  19.2457378   53.7984      0.357738107                   
demo_assume_multivariate_observe            ReverseDiffCompiled  18.49863146  58.35986883  0.316975206                   
demo_dot_assume_observe_index               ReverseDiffCompiled  22.21232631  59.22746781  0.375034205                   
demo_assume_dot_observe                     ReverseDiffCompiled  48.54647089  105.2648943  0.461183866                   
demo_assume_multivariate_observe_literal    ReverseDiffCompiled  15.95453127  53.94781274  0.295740095                   
demo_dot_assume_observe_index_literal       ReverseDiffCompiled  22.12678788  56.35138968  0.39265736                    
demo_assume_dot_observe_literal             ReverseDiffCompiled  36.52301255  96.41645676  0.378804758                   
demo_assume_observe_literal                 ReverseDiffCompiled  42.81327191  95.05131878  0.450422703                   
demo_assume_submodel_observe_index_literal  ReverseDiffCompiled  20.31091181  59.84733874  0.339378696                   
demo_dot_assume_observe_submodel            ReverseDiffCompiled  17.81327801  53.81377064  0.331017095                   
demo_dot_assume_observe_matrix_index        ReverseDiffCompiled  17.93991416  47.62331839  0.376704412                   
demo_assume_matrix_observe_matrix_index     ReverseDiffCompiled  13.98849123  47.35032558  0.29542545                    

[github-actions]

Benchmark Report

• this PR's head: d3a8a9221e7cdcc27956310d12f3f7edf0b2e82d
• base branch: 408ddb1f87348f5ccb6c72c055c29f39c1a284ad

Computer Information

Julia Version 1.11.8
Commit cf1da5e20e3 (2025-11-06 17:49 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 4 × AMD EPYC 7763 64-Core Processor
  WORD_SIZE: 64
  LLVM: libLLVM-16.0.6 (ORCJIT, znver3)
Threads: 1 default, 0 interactive, 1 GC (on 4 virtual cores)

Benchmark Results

┌───────────────────────┬───────┬─────────────┬───────────────────┬────────┬─────────────────────────────────┬────────────────────────────┬─────────────────────────────────┐
│                       │       │             │                   │        │        t(eval) / t(ref)         │     t(grad) / t(eval)      │        t(grad) / t(ref)         │
│                       │       │             │                   │        │ ──────────┬───────────┬──────── │ ───────┬─────────┬──────── │ ──────────┬───────────┬──────── │
│                 Model │   Dim │  AD Backend │           VarInfo │ Linked │      base │   this PR │ speedup │   base │ this PR │ speedup │      base │   this PR │ speedup │
├───────────────────────┼───────┼─────────────┼───────────────────┼────────┼───────────┼───────────┼─────────┼────────┼─────────┼─────────┼───────────┼───────────┼─────────┤
│               Dynamic │    10 │    mooncake │             typed │   true │    396.15 │    386.39 │    1.03 │   8.86 │    9.74 │    0.91 │   3511.32 │   3764.98 │    0.93 │
│                   LDA │    12 │ reversediff │             typed │   true │   2649.50 │   2642.53 │    1.00 │   5.16 │    5.14 │    1.00 │  13663.80 │  13577.35 │    1.01 │
│   Loop univariate 10k │ 10000 │    mooncake │             typed │   true │ 105514.20 │ 104107.32 │    1.01 │   3.90 │    3.80 │    1.03 │ 411738.09 │ 395588.83 │    1.04 │
├───────────────────────┼───────┼─────────────┼───────────────────┼────────┼───────────┼───────────┼─────────┼────────┼─────────┼─────────┼───────────┼───────────┼─────────┤
│    Loop univariate 1k │  1000 │    mooncake │             typed │   true │   7943.61 │   8024.26 │    0.99 │   4.76 │    4.55 │    1.05 │  37825.49 │  36472.62 │    1.04 │
│      Multivariate 10k │ 10000 │    mooncake │             typed │   true │  33740.52 │  32563.49 │    1.04 │   9.97 │   11.58 │    0.86 │ 336357.09 │ 377149.59 │    0.89 │
│       Multivariate 1k │  1000 │    mooncake │             typed │   true │   6430.17 │   3712.61 │    1.73 │   5.29 │    9.01 │    0.59 │  34021.81 │  33457.63 │    1.02 │
├───────────────────────┼───────┼─────────────┼───────────────────┼────────┼───────────┼───────────┼─────────┼────────┼─────────┼─────────┼───────────┼───────────┼─────────┤
│ Simple assume observe │     1 │ forwarddiff │             typed │  false │      2.72 │      2.69 │    1.01 │   3.96 │    3.93 │    1.01 │     10.76 │     10.58 │    1.02 │
│           Smorgasbord │   201 │ forwarddiff │             typed │  false │   1232.81 │   1231.68 │    1.00 │ 122.80 │  121.91 │    1.01 │ 151391.08 │ 150153.23 │    1.01 │
│           Smorgasbord │   201 │ forwarddiff │       simple_dict │   true │       err │       err │     err │    err │     err │     err │       err │       err │     err │
├───────────────────────┼───────┼─────────────┼───────────────────┼────────┼───────────┼───────────┼─────────┼────────┼─────────┼─────────┼───────────┼───────────┼─────────┤
│           Smorgasbord │   201 │ forwarddiff │ simple_namedtuple │   true │       err │       err │     err │    err │     err │     err │       err │       err │     err │
│           Smorgasbord │   201 │      enzyme │             typed │   true │   1678.08 │   1693.81 │    0.99 │   6.74 │    6.42 │    1.05 │  11302.78 │  10876.20 │    1.04 │
│           Smorgasbord │   201 │    mooncake │             typed │   true │   1682.04 │   1681.07 │    1.00 │   5.32 │    5.23 │    1.02 │   8952.83 │   8794.82 │    1.02 │
├───────────────────────┼───────┼─────────────┼───────────────────┼────────┼───────────┼───────────┼─────────┼────────┼─────────┼─────────┼───────────┼───────────┼─────────┤
│           Smorgasbord │   201 │ reversediff │             typed │   true │   1696.27 │   1700.93 │    1.00 │  90.05 │   90.10 │    1.00 │ 152742.80 │ 153251.37 │    1.00 │
│           Smorgasbord │   201 │ forwarddiff │      typed_vector │   true │   1684.76 │   1688.27 │    1.00 │  61.45 │   58.61 │    1.05 │ 103536.72 │  98942.76 │    1.05 │
│           Smorgasbord │   201 │ forwarddiff │           untyped │   true │   1680.45 │   1692.94 │    0.99 │  61.40 │   58.95 │    1.04 │ 103178.48 │  99797.25 │    1.03 │
├───────────────────────┼───────┼─────────────┼───────────────────┼────────┼───────────┼───────────┼─────────┼────────┼─────────┼─────────┼───────────┼───────────┼─────────┤
│           Smorgasbord │   201 │ forwarddiff │    untyped_vector │   true │   1695.48 │   1674.74 │    1.01 │  63.28 │   55.88 │    1.13 │ 107289.07 │  93576.77 │    1.15 │
│              Submodel │     1 │    mooncake │             typed │   true │      7.05 │      7.21 │    0.98 │   5.39 │    5.16 │    1.04 │     38.00 │     37.21 │    1.02 │
└───────────────────────┴───────┴─────────────┴───────────────────┴────────┴───────────┴───────────┴─────────┴────────┴─────────┴─────────┴───────────┴───────────┴─────────┘

[codecov]

Codecov Report

❌ Patch coverage is 85.00000% with 3 lines in your changes missing coverage. Please review.
✅ Project coverage is 78.91%. Comparing base (408ddb1) to head (d3a8a92).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
src/logdensityfunction.jl	85.00%	3 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1172      +/-   ##
==========================================
- Coverage   78.92%   78.91%   -0.01%     
==========================================
  Files          41       41              
  Lines        3914     3927      +13     
==========================================
+ Hits         3089     3099      +10     
- Misses        825      828       +3     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[mhauru]

LGTM

[github-actions]

DynamicPPL.jl documentation for PR #1172 is available at:
https://TuringLang.github.io/DynamicPPL.jl/previews/PR1172/