Implement predict, returned, logjoint, ... with OnlyAccsVarInfo

[penelopeysm]

This does lead to some improvement in performance, but not as much as I had hoped:

using Turing, Random, LinearAlgebra
y = [28, 8, -3, 7, -1, 1, 18, 12]
sigma = [15, 10, 16, 11, 9, 11, 10, 18]
@model function eight_schools(y, sigma)
    mu ~ Normal(0, 5)
    tau ~ truncated(Cauchy(0, 5); lower=0)
    theta ~ MvNormal(fill(mu, length(y)), tau^2 * I)
    for i in eachindex(y)
        y[i] ~ Normal(theta[i], sigma[i])
    end
    return (mu=mu, tau=tau)
end
model = eight_schools(y, sigma)

chn = sample(Xoshiro(468), model, NUTS(), 10000; progress=false)

@time returned(model, chn);
# main   :  0.760821 seconds (3.99 M allocations: 177.151 MiB, 2.56% gc time)
# This PR:  0.690771 seconds (3.15 M allocations: 90.624 MiB, 1.75% gc time)

It appears to me that most of this time is spent faffing with MCMCChains. Every time you try to get the value of @varname(mu) you have to go through the varname_to_symbol dict, etc. Even more importantly, there's an issue with theta, because that's vector-valued and when you access it you have to reconstruct the vector (with getvalue(dict, vn, dist). So I believe we are hitting a natural plateau that is caused by the data structure. Probably the performance gains with FlexiChains would be larger.

Still, I suppose it's worth putting this in because it's basically free performance, so why not? It's also cleaner code.

Also, I haven't tested this, but I'm about 95% sure that this makes predict, returned, etc. completely thread-safe if TSVI is used. That's because the combination of TSVI{<:OAVI} is completely threadsafe to everything.

---

A complete aside

Wouldn't it be fun if we could inspect the model, realise that the return value only involves mu and tau, realise that tau is a bits type and thus tau^2 * I cannot mutate tau, and thus optimise away all the things to do with theta and y?

[github-actions]

Benchmark Report

• this PR's head: edf4e74dedda18eb772a17fa82dfeff8b2aa46c2
• base branch: 766f6635903c401a79d3c2427dc60225f0053dad

Computer Information

Julia Version 1.11.7
Commit f2b3dbda30a (2025-09-08 12:10 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 │    351.32 │    484.41 │    0.73 │  15.90 │    8.75 │    1.82 │   5586.23 │   4237.70 │    1.32 │
│                   LDA │    12 │ reversediff │             typed │   true │   2850.98 │   3217.55 │    0.89 │   4.46 │    4.30 │    1.04 │  12707.99 │  13828.49 │    0.92 │
│   Loop univariate 10k │ 10000 │    mooncake │             typed │   true │ 107739.10 │ 100809.07 │    1.07 │   3.91 │    5.39 │    0.73 │ 421206.16 │ 543357.60 │    0.78 │
├───────────────────────┼───────┼─────────────┼───────────────────┼────────┼───────────┼───────────┼─────────┼────────┼─────────┼─────────┼───────────┼───────────┼─────────┤
│    Loop univariate 1k │  1000 │    mooncake │             typed │   true │   8317.30 │   8875.76 │    0.94 │   4.17 │    4.40 │    0.95 │  34702.41 │  39050.19 │    0.89 │
│      Multivariate 10k │ 10000 │    mooncake │             typed │   true │  30311.00 │  44967.67 │    0.67 │  10.01 │   11.28 │    0.89 │ 303403.64 │ 507402.74 │    0.60 │
│       Multivariate 1k │  1000 │    mooncake │             typed │   true │   3276.02 │   4746.36 │    0.69 │   9.36 │    6.96 │    1.34 │  30653.57 │  33037.90 │    0.93 │
├───────────────────────┼───────┼─────────────┼───────────────────┼────────┼───────────┼───────────┼─────────┼────────┼─────────┼─────────┼───────────┼───────────┼─────────┤
│ Simple assume observe │     1 │ forwarddiff │             typed │  false │      3.45 │      3.82 │    0.90 │   2.72 │    2.83 │    0.96 │      9.39 │     10.79 │    0.87 │
│           Smorgasbord │   201 │ forwarddiff │             typed │  false │   1187.90 │   1389.71 │    0.85 │ 120.40 │   73.56 │    1.64 │ 143020.95 │ 102228.07 │    1.40 │
│           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 │       err │       err │     err │    err │     err │     err │       err │       err │     err │
│           Smorgasbord │   201 │    mooncake │             typed │   true │   1635.16 │   2331.06 │    0.70 │   5.63 │    3.73 │    1.51 │   9201.87 │   8688.16 │    1.06 │
├───────────────────────┼───────┼─────────────┼───────────────────┼────────┼───────────┼───────────┼─────────┼────────┼─────────┼─────────┼───────────┼───────────┼─────────┤
│           Smorgasbord │   201 │ reversediff │             typed │   true │   1976.29 │   1922.20 │    1.03 │  70.42 │   85.65 │    0.82 │ 139169.52 │ 164632.58 │    0.85 │
│           Smorgasbord │   201 │ forwarddiff │      typed_vector │   true │   1632.93 │   1935.44 │    0.84 │  55.69 │   64.15 │    0.87 │  90937.32 │ 124167.75 │    0.73 │
│           Smorgasbord │   201 │ forwarddiff │           untyped │   true │   1625.71 │   1925.05 │    0.84 │  59.76 │   65.30 │    0.92 │  97154.66 │ 125701.17 │    0.77 │
├───────────────────────┼───────┼─────────────┼───────────────────┼────────┼───────────┼───────────┼─────────┼────────┼─────────┼─────────┼───────────┼───────────┼─────────┤
│           Smorgasbord │   201 │ forwarddiff │    untyped_vector │   true │   1620.52 │   1934.44 │    0.84 │  58.22 │   67.10 │    0.87 │  94339.11 │ 129809.45 │    0.73 │
│              Submodel │     1 │    mooncake │             typed │   true │      8.00 │      8.92 │    0.90 │   4.53 │    4.34 │    1.04 │     36.27 │     38.70 │    0.94 │
└───────────────────────┴───────┴─────────────┴───────────────────┴────────┴───────────┴───────────┴─────────┴────────┴─────────┴─────────┴───────────┴───────────┴─────────┘

[codecov]

Codecov Report

❌ Patch coverage is 89.28571% with 3 lines in your changes missing coverage. Please review.
✅ Project coverage is 80.60%. Comparing base (766f663) to head (edf4e74).
⚠️ Report is 1 commits behind head on breaking.

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

Additional details and impacted files

@@             Coverage Diff              @@
##           breaking    #1130      +/-   ##
============================================
- Coverage     81.07%   80.60%   -0.48%     
============================================
  Files            41       41              
  Lines          3894     3861      -33     
============================================
- Hits           3157     3112      -45     
- Misses          737      749      +12     

☔ 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 #1130 is available at:
https://TuringLang.github.io/DynamicPPL.jl/previews/PR1130/

[penelopeysm]

Fundamentally, all the functions in this extension really just use this under the hood.

FWIW the FlexiChains extension has a very similar structure and I believe these can be unified pretty much immediately after this PR. To be specific, DynamicPPL.reevaluate_with_chain should be implemented by each chain type in the most performant manner (FlexiChains doesn't use InitFromParams), but once that's done, the definitions of returned, logjoint, ..., pointwise_logdensities, ... can be shared.

predict can't yet be shared unfortunately, because the include_all keyword argument forces custom MCMCChains / FlexiChains code. That would require an extension of the AbstractChains API to support a subset-like operation.

[sunxd3]

code is good, maye something to be added to changelog

[sunxd3]

fallback being nothing is a change in default behavior, right? Might worth noting in changelog if true

[sunxd3]

no objection on getting rid of these, not sure if we should add an entry in changelog though. the likelihood that any of these are used somewhere is low

[penelopeysm]

Oh, you're completely correct, and I just totally neglected the changelog in this PR. My bad.

[sunxd3]

on

A complete aside
Wouldn't it be fun if we could inspect the model, realise that the return value only involves mu and tau, realise that tau is a bits type and thus tau^2 * I cannot mutate tau, and thus optimise away all the things to do with theta and y?

maybe we can somehow check Julia's effect analysis pipeline? (https://aviatesk.github.io/posts/effects-analysis/index.html, https://docs.julialang.org/en/v1/base/base/#Base.@assume_effects, https://github.com/JuliaLang/julia/blob/master/Compiler/src/effects.jl)

[penelopeysm]

Hmmm, I thought that the only way to make it work would be to handle the IR (i.e. like Mooncake), but I don't know anything about the effect system!