Fixed transformations

[penelopeysm]

Closes #1249. I'll write more here later, but you can see the changelog for a good overview of this PR.

Things to do

• Generalise LinkedVecTransformAccumulator to something that's more like FixedTransformAccumulator
• Add a convenience function for 'static'-fying an LDF (?)
• Add tests for new behaviour
• Add dev docs on FixedTransform

Possibly in a separate PR:

• Look at the Turing VI interface to see what it needs from bijector(model), and whether that can be removed (though, note also VI interface is too low-level Turing.jl#2783)
  • I have looked at this and it doesn't need bijector(model). It runs perfectly fine with dynamic transforms and can be made to work with fixed transforms as well, which would potentially provide perf benefits but otherwise not much difference is required in the API

---

Some scripts to use to benchmark:

using DynamicPPL, Distributions, LogDensityProblems, Random, Chairmarks, ForwardDiff, ADTypes, LinearAlgebra

# @model function esc(J, y, sigma)
#     mu ~ Normal(0, 5)
#     tau ~ truncated(Cauchy(0, 5); lower=0)
#     theta ~ MvNormal(fill(mu, J), tau^2 * I)
#     for i in 1:J
#         y[i] ~ Normal(theta[i], sigma[i])
#     end
# end
# J = 8
# y = [28, 8, -3, 7, -1, 1, 18, 12]
# sigma = [15, 10, 16, 11, 9, 11, 10, 18]
# m = esc(J, y, sigma)

@model function f()
    x ~ product_distribution([Beta(2,2), Uniform(4, 5), Normal()])
end
m = f()

ldf1 = LogDensityFunction(m, getlogjoint_internal, LinkAll(); adtype=AutoForwardDiff(), fix_transforms=true);
p = rand(Xoshiro(468), ldf1);
@b LogDensityProblems.logdensity(ldf1 ,p)
@b LogDensityProblems.logdensity_and_gradient(ldf1 ,p)

ldf2 = LogDensityFunction(m, getlogjoint_internal, LinkAll(); adtype=AutoForwardDiff());
p = rand(Xoshiro(468), ldf2);
@b LogDensityProblems.logdensity(ldf2 ,p)
@b LogDensityProblems.logdensity_and_gradient(ldf2 ,p)

[github-actions]

Benchmark Report

• this PR's head: aba492292d7d891744a0564a430d15dcc217e026
• base branch: e01fbd60d5a54a3b314b5bc5ea6654cefcba34a0

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 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 │ Linked │     base │  this PR │ speedup │   base │ this PR │ speedup │      base │   this PR │ speedup │
├───────────────────────┼───────┼─────────────┼────────┼──────────┼──────────┼─────────┼────────┼─────────┼─────────┼───────────┼───────────┼─────────┤
│               Dynamic │    10 │    mooncake │   true │   262.69 │   272.83 │    0.96 │   8.46 │    8.19 │    1.03 │   2223.34 │   2233.61 │    1.00 │
│                   LDA │    12 │ reversediff │   true │  2359.28 │  2606.68 │    0.91 │   2.52 │    2.06 │    1.23 │   5950.20 │   5361.53 │    1.11 │
│   Loop univariate 10k │ 10000 │    mooncake │   true │ 29273.85 │ 30459.81 │    0.96 │   6.91 │    7.16 │    0.97 │ 202319.53 │ 218006.69 │    0.93 │
├───────────────────────┼───────┼─────────────┼────────┼──────────┼──────────┼─────────┼────────┼─────────┼─────────┼───────────┼───────────┼─────────┤
│    Loop univariate 1k │  1000 │    mooncake │   true │  3718.48 │  3108.48 │    1.20 │   5.44 │    6.92 │    0.79 │  20243.17 │  21523.87 │    0.94 │
│      Multivariate 10k │ 10000 │    mooncake │   true │ 29763.37 │ 32148.10 │    0.93 │  10.38 │   10.21 │    1.02 │ 309077.57 │ 328330.23 │    0.94 │
│       Multivariate 1k │  1000 │    mooncake │   true │  3224.81 │  3438.04 │    0.94 │   9.59 │    9.61 │    1.00 │  30935.40 │  33026.56 │    0.94 │
├───────────────────────┼───────┼─────────────┼────────┼──────────┼──────────┼─────────┼────────┼─────────┼─────────┼───────────┼───────────┼─────────┤
│ Simple assume observe │     1 │ forwarddiff │  false │     0.81 │     0.88 │    0.92 │  10.57 │   10.25 │    1.03 │      8.59 │      9.05 │    0.95 │
│           Smorgasbord │   201 │ forwarddiff │  false │   887.33 │   958.72 │    0.93 │ 145.97 │  140.90 │    1.04 │ 129525.17 │ 135082.99 │    0.96 │
│           Smorgasbord │   201 │      enzyme │   true │  1211.30 │  1313.78 │    0.92 │   5.01 │    4.80 │    1.04 │   6065.31 │   6302.19 │    0.96 │
├───────────────────────┼───────┼─────────────┼────────┼──────────┼──────────┼─────────┼────────┼─────────┼─────────┼───────────┼───────────┼─────────┤
│           Smorgasbord │   201 │ forwarddiff │   true │  1195.61 │  1302.76 │    0.92 │  71.77 │   70.53 │    1.02 │  85814.58 │  91884.00 │    0.93 │
│           Smorgasbord │   201 │    mooncake │   true │  1196.99 │  1311.29 │    0.91 │   5.09 │    5.09 │    1.00 │   6093.00 │   6669.28 │    0.91 │
│           Smorgasbord │   201 │ reversediff │   true │  1192.79 │  1308.30 │    0.91 │ 128.69 │  127.52 │    1.01 │ 153505.37 │ 166840.23 │    0.92 │
├───────────────────────┼───────┼─────────────┼────────┼──────────┼──────────┼─────────┼────────┼─────────┼─────────┼───────────┼───────────┼─────────┤
│              Submodel │     1 │    mooncake │   true │     0.81 │     0.88 │    0.93 │  28.98 │   28.93 │    1.00 │     23.54 │     25.37 │    0.93 │
└───────────────────────┴───────┴─────────────┴────────┴──────────┴──────────┴─────────┴────────┴─────────┴─────────┴───────────┴───────────┴─────────┘

[codecov]

Codecov Report

❌ Patch coverage is 79.03226% with 52 lines in your changes missing coverage. Please review.
✅ Project coverage is 79.06%. Comparing base (61eda05) to head (aba4922).
⚠️ Report is 8 commits behind head on breaking.

Files with missing lines	Patch %	Lines
src/transformed_values.jl	52.54%	28 Missing ⚠️
src/varinfo.jl	87.27%	7 Missing ⚠️
src/abstract_varinfo.jl	33.33%	4 Missing ⚠️
src/accumulators/vector_values.jl	63.63%	4 Missing ⚠️
src/contexts/default.jl	40.00%	3 Missing ⚠️
src/accumulators/fixed_transforms.jl	85.71%	2 Missing ⚠️
src/contexts/init.jl	88.88%	2 Missing ⚠️
src/accumulators/raw_values.jl	93.33%	1 Missing ⚠️
src/test_utils/ad.jl	50.00%	1 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##           breaking    #1327      +/-   ##
============================================
+ Coverage     78.13%   79.06%   +0.93%     
============================================
  Files            50       50              
  Lines          3590     3559      -31     
============================================
+ Hits           2805     2814       +9     
+ Misses          785      745      -40     

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

[penelopeysm]

@sunxd3 would you be able to take a look? Here's the TLDR and my overall thoughts on this:

• src/transformed_values.jl -- There used to be UntransformedValue, VectorValue and LinkedVectorValue <: AbstractTransformedValue -- these are renamed to TransformedValue(val, NoTransform()), TransformedValue(val, Unlink()), and TransformedValue(val, DynamicLink()). A lot of the changes in this PR are just boilerplate resulting from this slightly different type hierarchy. Better names are also welcome, although renaming some of this might be breaking at this point.

• The new type hierarchy lets us extend it to TransformedValue(val, FixedTransform(f)). The new types are breaking, and we could avoid this by instead making a new subtype, FixedTransformedValue(val, f). I'm still a bit uncertain about which way to go on this: I went with the overhaul, but I don't actually think that it's necessary, and part of me thinks that it would just be better to not break the old API. If you agree, I'll try to change it back.

• The remaining question would be whether the general interface makes sense to you (we discussed this on Monday). I think the docs are the best place to read about the interface https://turinglang.org/DynamicPPL.jl/previews/PR1327/transforms_fixed/. The implementation of the interface is in src/accumulators/fixed_transforms.jl and src/logdensityfunction.jl, but that's all quite standard boilerplate, nothing particularly interesting.

[sunxd3]

I think the type hierarchy changes all makes sense. Although it is very breaking, even though I don't imagine a lot of people depend on the transformation interface, it might still worth giving some deprecation buffer (we can talk about this in the meeting).

The general interface also make sense to me.

Multivariate 10k has some performance regression on Mooncake, I wonder why.

[sunxd3]

The biggest break from this might be the removal of vi[@varname]syntax. It just existed for quite a long time

[penelopeysm]

Multivariate 10k has some performance regression on Mooncake, I wonder why

It's just CI flakiness, unfortunately I've found the numbers to vary by up to 50% even on a good day. I benchmarked locally to be sure, and there's no difference

Benchmarking code & results

code

using DynamicPPL, Distributions, ADTypes, Mooncake, LogDensityProblems, Chairmarks
@model function multivariate(num_dims, ::Type{TV}=Vector{Float64}) where {TV}
    a = TV(undef, num_dims)
    o = TV(undef, num_dims)
    a ~ product_distribution(fill(Normal(0, 1), num_dims))
    m = sum(a)
    o ~ product_distribution(fill(Normal(m, 1), num_dims))
    return (; a=a)
end
x = randn(10000);
model = multivariate(10000) | (; o = x);
ldf = LogDensityFunction(model, getlogjoint_internal, LinkAll(); adtype=AutoMooncake());
p = rand(ldf);
@be LogDensityProblems.logdensity($ldf, $p)
@be LogDensityProblems.logdensity_and_gradient($ldf, $p)

main

julia> @be LogDensityProblems.logdensity($ldf, $p)
Benchmark: 371 samples with 1 evaluation
 min    104.875 μs (28 allocs: 1.063 MiB)
 median 148.292 μs (28 allocs: 1.063 MiB)
 mean   256.189 μs (28 allocs: 1.063 MiB, 4.57% gc time)
 max    18.380 ms (28 allocs: 1.063 MiB, 98.90% gc time)

julia> @be LogDensityProblems.logdensity_and_gradient($ldf, $p)
Benchmark: 52 samples with 1 evaluation
 min    1.454 ms (59 allocs: 2.220 MiB)
 median 1.570 ms (59 allocs: 2.220 MiB)
 mean   1.947 ms (59 allocs: 2.220 MiB, 4.33% gc time)
 max    13.387 ms (59 allocs: 2.220 MiB, 87.35% gc time)

This PR

julia> @be LogDensityProblems.logdensity($ldf, $p)
Benchmark: 432 samples with 1 evaluation
 min    106.250 μs (28 allocs: 1.063 MiB)
 median 142.667 μs (28 allocs: 1.063 MiB)
 mean   183.366 μs (28 allocs: 1.063 MiB, 4.52% gc time)
 max    1.048 ms (28 allocs: 1.063 MiB, 86.68% gc time)

julia> @be LogDensityProblems.logdensity_and_gradient($ldf, $p)
Benchmark: 50 samples with 1 evaluation
 min    1.483 ms (59 allocs: 2.220 MiB)
 median 1.571 ms (59 allocs: 2.220 MiB)
 mean   2.020 ms (59 allocs: 2.220 MiB, 5.51% gc time)
 max    14.799 ms (59 allocs: 2.220 MiB, 88.62% gc time)

[penelopeysm]

The biggest break from this might be the removal of vi[@varname] syntax. It just existed for quite a long time

Yup, I think I haven't done a changelog for this PR but that definitely needs to be mentioned upfront. There is a changelog. I was more hardworking than I thought I was!

The change is intentional on my part though:

1. vi[vn] has never really been correct -- it's known that this breaks precisely in the case where the transforms are not constant.

2. VarInfo can store variable values in multiple places -- there is the vi.values part itself, or it could have values stored in accumulators. It's not immediately obvious to someone reading the code that vi[vn] should read from one of those and not the others.

3. More broadly speaking, something like e.g. unflatten!! changes the former but doesn't change the latter, and I don't like that it can very easily end up in a state that's inconsistent, and that we regularly abuse inconsistent VarInfos everywhere in Turing. Because of this, my medium-term intention is to remove VarInfo entirely anyway, so this is really just one step towards that.

[sunxd3]

I saw the changelog and think the are clear. I agree with your judgement on this — it is probably impossible to make correct at the moment and removing makes sense. I was just mentioning because it's a pattern that was used a lot internally and externally.

Maybe we could leave a deprecate or error message?

[penelopeysm]

I added this:

julia> vi[@varname(x)]
ERROR: ArgumentError: The method `varinfo[varname(s)]` has been removed in DynamicPPL v0.41: please see https://turinglang.org/DynamicPPL.jl/stable/migration/#Getting-parameter-values for information on how to update your code if needed.

[penelopeysm]

I think the remaining question is whether to change the type hierarchy. As I see it, the main benefit of switching it over is that basically that we don't duplicate type info for the transform + transformed value (previously they would be separate things, now the latter wraps the former). I was a bit concerned about breakage, but I did a GitHub search and I think it only breaks Turing so maybe it's not too bad?

[sunxd3]

My judgement is that the type hierarchy breaking is not a big deal — it was pretty new and somewhat of an "advanced" feature

[penelopeysm]

@sunxd3 Thanks! Are you happy with the PR, or are there more things you think need to be discussed?

[sunxd3]

happy happy