Use structural AbstractPPL AD prep for LogDensityFunction

[yebai]

Fix #1364

TODO

• use main as base, as otherwise benchmarking would fail, need to change back to Is _use_closure still necessary with native gradient API #1364 before merging.
• Julia 1.10’s Pkg does not reliably apply the root [sources] override during build/test, so min CI fails.

This change switches LogDensityFunction AD preparation to pass a structural LogDensityAt problem directly into AbstractPPL.prepare(...) instead of wrapping it in backend-specific anonymous closures. It also removes the old _use_closure machinery, since the structural problem object now provides the stable one-argument evaluator shape that AD backends need.

The accompanying test coverage adds a focused ReverseDiff check to confirm that AutoReverseDiff(; compile=true) still retains a compiled tape and can be reused across repeated logdensity_and_gradient calls.

Benchmarks

  +----------------------+------------+--------------+--------+--------------+
  | Backend              | PR grad μs | main grad μs |  Ratio | Summary      |
  +----------------------+------------+--------------+--------+--------------+
  | ForwardDiff          |      0.322 |        0.322 |  1.001x | 0.1% slower  |
  | Forward Mooncake     |      0.900 |        2.601 |  0.346x | 2.89x faster |
  | Forward Enzyme       |      0.789 |        1.051 |  0.751x | 24.9% faster |
  | ReverseDiff          |     14.239 |       15.225 |  0.935x | 6.5% faster  |
  | ReverseDiff compiled |      5.355 |        5.872 |  0.912x | 8.8% faster  |
  | Reverse Mooncake     |      0.892 |        0.928 |  0.961x | 3.9% faster  |
  | Reverse Enzyme       |      1.253 |        1.268 |  0.988x | 1.2% faster  |
  +----------------------+------------+--------------+--------+--------------+

[github-actions]

Benchmark Report

• this PR's head: 378f974319d1e92114cb1b8d4796854957c3b84a
• base branch: e1cac2fd41e6f23e42c2fd8cf3622efd72c9a843

Computer Information

Julia Version 1.11.9
Commit 53a02c0720c (2026-02-06 00:27 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 4 × AMD EPYC 9V74 80-Core Processor
  WORD_SIZE: 64
  LLVM: libLLVM-16.0.6 (ORCJIT, znver4)
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 │ Linked │     base │  this PR │ speedup │   base │ this PR │ speedup │      base │   this PR │ speedup │
├───────────────────────┼───────┼─────────────┼────────┼──────────┼──────────┼─────────┼────────┼─────────┼─────────┼───────────┼───────────┼─────────┤
│               Dynamic │    10 │    mooncake │   true │   284.04 │   259.11 │    1.10 │   7.72 │    8.43 │    0.92 │   2193.09 │   2184.43 │    1.00 │
│                   LDA │    12 │ reversediff │   true │  2647.27 │  2472.80 │    1.07 │   2.13 │    2.07 │    1.03 │   5649.76 │   5127.77 │    1.10 │
│   Loop univariate 10k │ 10000 │    mooncake │   true │ 29965.39 │ 27033.91 │    1.11 │   7.13 │    7.14 │    1.00 │ 213637.27 │ 192923.12 │    1.11 │
├───────────────────────┼───────┼─────────────┼────────┼──────────┼──────────┼─────────┼────────┼─────────┼─────────┼───────────┼───────────┼─────────┤
│    Loop univariate 1k │  1000 │    mooncake │   true │  3061.03 │  2762.25 │    1.11 │   6.93 │    6.89 │    1.01 │  21222.95 │  19023.45 │    1.12 │
│      Multivariate 10k │ 10000 │    mooncake │   true │ 33659.72 │ 30373.47 │    1.11 │  10.32 │   10.13 │    1.02 │ 347334.39 │ 307682.29 │    1.13 │
│       Multivariate 1k │  1000 │    mooncake │   true │  3662.68 │  3350.48 │    1.09 │   9.28 │    9.30 │    1.00 │  33991.76 │  31163.62 │    1.09 │
├───────────────────────┼───────┼─────────────┼────────┼──────────┼──────────┼─────────┼────────┼─────────┼─────────┼───────────┼───────────┼─────────┤
│ Simple assume observe │     1 │ forwarddiff │  false │     0.73 │     0.80 │    0.91 │  12.62 │   10.73 │    1.18 │      9.24 │      8.62 │    1.07 │
│           Smorgasbord │   201 │ forwarddiff │  false │   932.11 │   834.35 │    1.12 │ 161.30 │   80.94 │    1.99 │ 150347.28 │  67532.72 │    2.23 │
│           Smorgasbord │   201 │      enzyme │   true │  1175.67 │  1076.27 │    1.09 │   5.69 │    5.80 │    0.98 │   6686.45 │   6240.18 │    1.07 │
├───────────────────────┼───────┼─────────────┼────────┼──────────┼──────────┼─────────┼────────┼─────────┼─────────┼───────────┼───────────┼─────────┤
│           Smorgasbord │   201 │ forwarddiff │   true │  1170.39 │  1067.07 │    1.10 │  78.20 │   80.30 │    0.97 │  91525.58 │  85684.99 │    1.07 │
│           Smorgasbord │   201 │    mooncake │   true │  1175.19 │  1085.43 │    1.08 │   5.51 │    5.49 │    1.00 │   6480.29 │   5958.61 │    1.09 │
│           Smorgasbord │   201 │ reversediff │   true │  1171.90 │  1069.04 │    1.10 │ 144.13 │  146.23 │    0.99 │ 168909.43 │ 156325.58 │    1.08 │
├───────────────────────┼───────┼─────────────┼────────┼──────────┼──────────┼─────────┼────────┼─────────┼─────────┼───────────┼───────────┼─────────┤
│              Submodel │     1 │    mooncake │   true │     0.73 │     0.67 │    1.09 │  30.99 │   32.88 │    0.94 │     22.71 │     22.01 │    1.03 │
└───────────────────────┴───────┴─────────────┴────────┴──────────┴──────────┴─────────┴────────┴─────────┴─────────┴───────────┴───────────┴─────────┘

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 82.20%. Comparing base (e1cac2f) to head (378f974).
⚠️ Report is 1 commits behind head on adproblems-interface.

Additional details and impacted files

@@                   Coverage Diff                    @@
##           adproblems-interface    #1365      +/-   ##
========================================================
+ Coverage                 78.68%   82.20%   +3.52%     
========================================================
  Files                        49       49              
  Lines                      3490     3490              
========================================================
+ Hits                       2746     2869     +123     
+ Misses                      744      621     -123     

☔ 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.

[github-actions]

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

[yebai]

@penelopeysm, this should be ready for a look.

[penelopeysm]

I think the multivariate 10k benchmark should probably be looked at locally too, maybe tomorrow.

[yebai]

@penelopeysm, all the issues should have been fixed.

I think the multivariate 10k benchmark should probably be looked at locally too, maybe tomorrow.

Can you run this locally?