From 22be408ffe4fd5985b3ab3c5b88c193c8dd5b37a Mon Sep 17 00:00:00 2001
From: Markus Hauru <mhauru@turing.ac.uk>
Date: Tue, 25 Jun 2024 11:11:09 +0100
Subject: [PATCH 01/17] Add support for Turing v0.33

---
 Project.toml               |  2 +-
 ext/PathfinderTuringExt.jl | 74 +++++++++++++++++++++++++++++++++-----
 2 files changed, 67 insertions(+), 9 deletions(-)

diff --git a/Project.toml b/Project.toml
index 4b6cc0e9b..3924007f1 100644
--- a/Project.toml
+++ b/Project.toml
@@ -58,7 +58,7 @@ SciMLBase = "1.61.2, 2"
 Statistics = "1.6"
 StatsBase = "0.33.1, 0.34"
 Transducers = "0.4.66"
-Turing = "0.24.2, 0.25, 0.26, 0.27, 0.28, 0.29, 0.30, 0.31, 0.32"
+Turing = "0.24.2, 0.25, 0.26, 0.27, 0.28, 0.29, 0.30, 0.31, 0.32, 0.33"
 UnPack = "1"
 julia = "1.6"
 
diff --git a/ext/PathfinderTuringExt.jl b/ext/PathfinderTuringExt.jl
index f6dd37f30..76c8e6978 100644
--- a/ext/PathfinderTuringExt.jl
+++ b/ext/PathfinderTuringExt.jl
@@ -4,6 +4,7 @@ if isdefined(Base, :get_extension)
     using Accessors: Accessors
     using DynamicPPL: DynamicPPL
     using MCMCChains: MCMCChains
+    using Optimization: Optimization
     using Pathfinder: Pathfinder
     using Random: Random
     using Turing: Turing
@@ -12,6 +13,7 @@ else  # using Requires
     using ..Accessors: Accessors
     using ..DynamicPPL: DynamicPPL
     using ..MCMCChains: MCMCChains
+    using ..Optimization: Optimization
     using ..Pathfinder: Pathfinder
     using ..Random: Random
     using ..Turing: Turing
@@ -101,6 +103,58 @@ function varnames_to_ranges(metadata::DynamicPPL.Metadata)
     return Dict(zip(metadata.vns, ranges))
 end
 
+"""
+    transform_to_constrained(
+        p::AbstractArray, vi::DynamicPPL.VarInfo, model::DynamicPPL.Model
+    )
+
+Transform a vector of parameters `p` from unconstrained to constrained space.
+"""
+function transform_to_constrained(
+    p::AbstractArray, vi::DynamicPPL.VarInfo, model::DynamicPPL.Model
+)
+    p = copy(p)
+    @assert DynamicPPL.istrans(vi)
+    vi = DynamicPPL.unflatten(vi, p)
+    p .= DynamicPPL.invlink!!(vi, model)[:]
+    # Restore the linking, since we mutated vi.
+    DynamicPPL.link!!(vi, model)
+    return p
+end
+
+"""
+    set_up_model_optimisation(model::DynamicPPL.Model, init)
+
+Create the necessary pieces for running optimisation on `model`.
+
+Returns
+* An `Optimization.OptimizationFunction` that evaluates the log density of the model and its
+gradient in the unconstrained space.
+* The initial value `init` transformed to unconstrained space.
+* A function `transform_result` that transforms the results back to constrained space. It
+takes a single vector argument.
+"""
+function set_up_model_optimisation(model::DynamicPPL.Model, init)
+    # The inner context deterimines whether we are solving MAP or MLE.
+    inner_context = DynamicPPL.DefaultContext()
+    ctx = Turing.Optimisation.OptimizationContext(inner_context)
+    log_density = Turing.Optimisation.OptimLogDensity(model, ctx)
+    # Initialise the varinfo with the initial value and then transform to unconstrained
+    # space.
+    Accessors.@set log_density.varinfo = DynamicPPL.unflatten(log_density.varinfo, init)
+    transformed_varinfo = DynamicPPL.link(log_density.varinfo, log_density.model)
+    log_density = Accessors.@set log_density.varinfo = transformed_varinfo
+    init = log_density.varinfo[:]
+    # Create a function that applies the appropriate inverse transformation to results, to
+    # bring them back to constrained space.
+    transform_result(p) = transform_to_constrained(p, log_density.varinfo, model)
+    f = Optimization.OptimizationFunction(
+        (x, _) -> log_density(x),;
+        grad = (G,x,p) -> log_density(nothing, G, x),
+    )
+    return f, init, transform_result
+end
+
 function Pathfinder.pathfinder(
     model::DynamicPPL.Model;
     rng=Random.GLOBAL_RNG,
@@ -110,10 +164,13 @@ function Pathfinder.pathfinder(
     kwargs...,
 )
     var_names = flattened_varnames_list(model)
-    prob = Turing.optim_problem(model, Turing.MAP(); constrained=false, init_theta=init)
-    init_sampler(rng, prob.prob.u0)
-    result = Pathfinder.pathfinder(prob.prob; rng, input=model, kwargs...)
-    draws = reduce(vcat, transpose.(prob.transform.(eachcol(result.draws))))
+    # If no initial value is provided, sample from prior.
+    init = init === nothing ? rand(Vector, model) : init
+    f, init, transform_result = set_up_model_optimisation(model, init)
+    prob = Optimization.OptimizationProblem(f, init)
+    init_sampler(rng, init)
+    result = Pathfinder.pathfinder(prob; rng, input=model, kwargs...)
+    draws = reduce(vcat, transpose.(transform_result.(eachcol(result.draws))))
     chns = MCMCChains.Chains(draws, var_names; info=(; pathfinder_result=result))
     result_new = Accessors.@set result.draws_transformed = chns
     return result_new
@@ -129,14 +186,15 @@ function Pathfinder.multipathfinder(
     kwargs...,
 )
     var_names = flattened_varnames_list(model)
-    fun = Turing.optim_function(model, Turing.MAP(); constrained=false)
-    init1 = fun.init()
+    # Sample from prior.
+    init1 = rand(Vector, model)
+    fun, init1, transform_result = set_up_model_optimisation(model, init1)
     init = [init_sampler(rng, init1)]
     for _ in 2:nruns
         push!(init, init_sampler(rng, deepcopy(init1)))
     end
-    result = Pathfinder.multipathfinder(fun.func, ndraws; rng, input=model, init, kwargs...)
-    draws = reduce(vcat, transpose.(fun.transform.(eachcol(result.draws))))
+    result = Pathfinder.multipathfinder(fun, ndraws; rng, input=model, init, kwargs...)
+    draws = reduce(vcat, transpose.(transform_result.(eachcol(result.draws))))
     chns = MCMCChains.Chains(draws, var_names; info=(; pathfinder_result=result))
     result_new = Accessors.@set result.draws_transformed = chns
     return result_new

From f3d4dd857c9ebd7b56a36938872b4e57c73fe486 Mon Sep 17 00:00:00 2001
From: Seth Axen <seth@sethaxen.com>
Date: Wed, 3 Jul 2024 09:54:49 +0200
Subject: [PATCH 02/17] Bump Turing compat for docs and tests

---
 docs/Project.toml                    | 2 +-
 test/integration/Turing/Project.toml | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/docs/Project.toml b/docs/Project.toml
index 226d7c1cf..2dec3213d 100644
--- a/docs/Project.toml
+++ b/docs/Project.toml
@@ -30,4 +30,4 @@ StatsFuns = "1"
 StatsPlots = "0.14.21, 0.15"
 TransformVariables = "0.6.2, 0.7, 0.8"
 TransformedLogDensities = "1.0.2"
-Turing = "0.30.5, 0.31, 0.32"
+Turing = "0.30.5, 0.31, 0.32, 0.33"
diff --git a/test/integration/Turing/Project.toml b/test/integration/Turing/Project.toml
index 5cdacb82e..6a77a274b 100644
--- a/test/integration/Turing/Project.toml
+++ b/test/integration/Turing/Project.toml
@@ -6,5 +6,5 @@ Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
 
 [compat]
 Pathfinder = "0.9"
-Turing = "0.30.5, 0.31, 0.32"
+Turing = "0.30.5, 0.31, 0.32 0.33"
 julia = "1.6"

From 25b42e6a3a6936993733e11dabc502519d8625a5 Mon Sep 17 00:00:00 2001
From: Seth Axen <seth@sethaxen.com>
Date: Wed, 3 Jul 2024 10:06:25 +0200
Subject: [PATCH 03/17] Fix Turing compat format

---
 test/integration/Turing/Project.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/integration/Turing/Project.toml b/test/integration/Turing/Project.toml
index 6a77a274b..c9fb875e3 100644
--- a/test/integration/Turing/Project.toml
+++ b/test/integration/Turing/Project.toml
@@ -6,5 +6,5 @@ Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
 
 [compat]
 Pathfinder = "0.9"
-Turing = "0.30.5, 0.31, 0.32 0.33"
+Turing = "0.30.5, 0.31, 0.32, 0.33"
 julia = "1.6"

From 172b4cecbb9d0a7a56f25437ac0042b4a06a9cc9 Mon Sep 17 00:00:00 2001
From: Seth Axen <seth@sethaxen.com>
Date: Thu, 4 Jul 2024 11:06:36 +0200
Subject: [PATCH 04/17] Add and test utility to create a LogDensityProblem

---
 ext/PathfinderTuringExt.jl           | 17 ++++++++++++++++
 test/integration/Turing/Project.toml |  2 ++
 test/integration/Turing/runtests.jl  | 30 +++++++++++++++++++++++++++-
 3 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/ext/PathfinderTuringExt.jl b/ext/PathfinderTuringExt.jl
index 258f9177f..51f2343e6 100644
--- a/ext/PathfinderTuringExt.jl
+++ b/ext/PathfinderTuringExt.jl
@@ -22,6 +22,23 @@ else  # using Requires
     import ..Pathfinder: flattened_varnames_list
 end
 
+"""
+    create_log_density_problem(model::DynamicPPL.Model)
+
+Create a log density problem from a `model`.
+
+The return value is an object implementing the LogDensityProblems API whose log-density is
+that of the `model` transformed to unconstrained space with the appropriate log-density
+adjustment due to change of variables.
+"""
+function create_log_density_problem(model::DynamicPPL.Model)
+    # create an unconstrained VarInfo
+    varinfo = DynamicPPL.link(DynamicPPL.VarInfo(model), model)
+    # DefaultContext ensures that the log-density adjustment is computed
+    prob = DynamicPPL.LogDensityFunction(varinfo, model, DynamicPPL.DefaultContext())
+    return prob
+end
+
 # utilities for working with Turing model parameter names using only the DynamicPPL API
 
 function Pathfinder.flattened_varnames_list(model::DynamicPPL.Model)
diff --git a/test/integration/Turing/Project.toml b/test/integration/Turing/Project.toml
index c9fb875e3..c40df5265 100644
--- a/test/integration/Turing/Project.toml
+++ b/test/integration/Turing/Project.toml
@@ -1,10 +1,12 @@
 [deps]
+LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
 Pathfinder = "b1d3bc72-d0e7-4279-b92f-7fa5d6d2d454"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
 
 [compat]
+LogDensityProblems = "2.1.0"
 Pathfinder = "0.9"
 Turing = "0.30.5, 0.31, 0.32, 0.33"
 julia = "1.6"
diff --git a/test/integration/Turing/runtests.jl b/test/integration/Turing/runtests.jl
index de5ee966f..58e1eae39 100644
--- a/test/integration/Turing/runtests.jl
+++ b/test/integration/Turing/runtests.jl
@@ -1,4 +1,11 @@
-using Pathfinder, Random, Test, Turing
+using LogDensityProblems, Pathfinder, Random, Test, Turing
+using Turing.Bijectors
+
+if isdefined(Base, :get_extension)
+    PathfinderTuringExt = Base.get_extension(Pathfinder, :PathfinderTuringExt)
+else
+    PathfinderTuringExt = Pathfinder.PathfinderTuringExt
+end
 
 Random.seed!(0)
 
@@ -12,6 +19,27 @@ Random.seed!(0)
 end
 
 @testset "Turing integration" begin
+    @testset "create_log_density_problem" begin
+        @testset for bijector in [elementwise(log), Bijectors.SimplexBijector()],
+            udist in [Normal(1, 2), Normal(3, 4)],
+            n in [1, 5]
+
+            binv = Bijectors.inverse(bijector)
+            dist = filldist(udist, n)
+            dist_trans = Bijectors.transformed(dist, binv)
+            @model function model()
+                return y ~ dist_trans
+            end
+            prob = PathfinderTuringExt.create_log_density_problem(model())
+            @test LogDensityProblems.capabilities(prob) isa
+                LogDensityProblems.LogDensityOrder{0}
+            x = rand(n, 10)
+            # after applying the Jacobian correction, the log-density of the model should
+            # be the same as the log-density of the distribution in unconstrained space
+            @test LogDensityProblems.logdensity.(Ref(prob), eachcol(x)) ≈ logpdf(dist, x)
+        end
+    end
+
     x = 0:0.01:1
     y = sin.(x) .+ randn.() .* 0.2 .+ x
     X = [x x .^ 2 x .^ 3]

From 564d3f3ba77a90bd156e2b6a780532ee7ea4c68a Mon Sep 17 00:00:00 2001
From: Seth Axen <seth@sethaxen.com>
Date: Thu, 4 Jul 2024 11:18:09 +0200
Subject: [PATCH 05/17] Add utility for converting draws to chains

---
 ext/PathfinderTuringExt.jl          | 21 +++++++++++++++++++++
 test/integration/Turing/runtests.jl | 21 +++++++++++++++++++++
 2 files changed, 42 insertions(+)

diff --git a/ext/PathfinderTuringExt.jl b/ext/PathfinderTuringExt.jl
index 51f2343e6..01c3c4fb3 100644
--- a/ext/PathfinderTuringExt.jl
+++ b/ext/PathfinderTuringExt.jl
@@ -39,6 +39,27 @@ function create_log_density_problem(model::DynamicPPL.Model)
     return prob
 end
 
+"""
+    draws_to_chains(model::DynamicPPL.Model, draws) -> MCMCChains.Chains
+
+Convert a `(nparams, ndraws)` matrix of unconstrained `draws` to an `MCMCChains.Chains`
+object with corresponding constrained draws and names according to `model`.
+"""
+function draws_to_chains(model::DynamicPPL.Model, draws::AbstractMatrix)
+    varinfo = DynamicPPL.link(DynamicPPL.VarInfo(model), model)
+    draw_con_varinfos = map(eachcol(draws)) do draw
+        # this re-evaluates the model but allows supporting dynamic bijectors
+        # https://github.com/TuringLang/Turing.jl/issues/2195
+        return Turing.Inference.getparams(model, DynamicPPL.unflatten(varinfo, draw))
+    end
+    param_con_names = map(first, first(draw_con_varinfos))
+    draws_con = reduce(
+        vcat, Iterators.map(transpose ∘ Base.Fix1(map, last), draw_con_varinfos)
+    )
+    chns = MCMCChains.Chains(draws_con, param_con_names)
+    return chns
+end
+
 # utilities for working with Turing model parameter names using only the DynamicPPL API
 
 function Pathfinder.flattened_varnames_list(model::DynamicPPL.Model)
diff --git a/test/integration/Turing/runtests.jl b/test/integration/Turing/runtests.jl
index 58e1eae39..c33bcb104 100644
--- a/test/integration/Turing/runtests.jl
+++ b/test/integration/Turing/runtests.jl
@@ -18,6 +18,15 @@ Random.seed!(0)
     return (; y)
 end
 
+# adapted from https://github.com/TuringLang/Turing.jl/issues/2195
+@model function dynamic_const_model()
+    lb ~ Uniform(0, 0.1)
+    ub ~ Uniform(0.11, 0.2)
+    return x ~ Bijectors.transformed(
+        Normal(0, 1), Bijectors.inverse(Bijectors.Logit(lb, ub))
+    )
+end
+
 @testset "Turing integration" begin
     @testset "create_log_density_problem" begin
         @testset for bijector in [elementwise(log), Bijectors.SimplexBijector()],
@@ -40,6 +49,18 @@ end
         end
     end
 
+    @testset "draws_to_chains" begin
+        draws = randn(3, 100)
+        model = dynamic_const_model()
+        chns = PathfinderTuringExt.draws_to_chains(model, draws)
+        @test chns isa MCMCChains.Chains
+        @test size(chns) == (100, 3, 1)
+        @test names(chns) == [:lb, :ub, :x]
+        @test all(0 .< chns[:, :lb, 1] .< 0.1)
+        @test all(0.11 .< chns[:, :ub, 1] .< 0.2)
+        @test all(chns[:, :lb, 1] .< chns[:, :x, 1] .< chns[:, :ub, 1])
+    end
+
     x = 0:0.01:1
     y = sin.(x) .+ randn.() .* 0.2 .+ x
     X = [x x .^ 2 x .^ 3]

From 61a99ff0ecc13df3a453bc2ca59601ceb9d5d418 Mon Sep 17 00:00:00 2001
From: Seth Axen <seth@sethaxen.com>
Date: Thu, 4 Jul 2024 11:21:39 +0200
Subject: [PATCH 06/17] Use draws_to_chains utility

---
 ext/PathfinderTuringExt.jl | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

diff --git a/ext/PathfinderTuringExt.jl b/ext/PathfinderTuringExt.jl
index 01c3c4fb3..2d8e1418b 100644
--- a/ext/PathfinderTuringExt.jl
+++ b/ext/PathfinderTuringExt.jl
@@ -204,15 +204,14 @@ function Pathfinder.pathfinder(
     adtype::ADTypes.AbstractADType=Pathfinder.default_ad(),
     kwargs...,
 )
-    var_names = flattened_varnames_list(model)
     # If no initial value is provided, sample from prior.
     init = init === nothing ? rand(Vector, model) : init
     f, init, transform_result = set_up_model_optimisation(model, init)
     prob = Optimization.OptimizationProblem(f, init)
     init_sampler(rng, init)
     result = Pathfinder.pathfinder(prob; rng, input=model, kwargs...)
-    draws = reduce(vcat, transpose.(transform_result.(eachcol(result.draws))))
-    chns = MCMCChains.Chains(draws, var_names; info=(; pathfinder_result=result))
+    chns_info = (; pathfinder_result=result)
+    chns = Accessors.@set draws_to_chains(model, result.draws).info = chns_info
     result_new = Accessors.@set result.draws_transformed = chns
     return result_new
 end
@@ -227,7 +226,6 @@ function Pathfinder.multipathfinder(
     adtype=Pathfinder.default_ad(),
     kwargs...,
 )
-    var_names = flattened_varnames_list(model)
     # Sample from prior.
     init1 = rand(Vector, model)
     fun, init1, transform_result = set_up_model_optimisation(model, init1)
@@ -236,8 +234,8 @@ function Pathfinder.multipathfinder(
         push!(init, init_sampler(rng, deepcopy(init1)))
     end
     result = Pathfinder.multipathfinder(fun, ndraws; rng, input=model, init, kwargs...)
-    draws = reduce(vcat, transpose.(transform_result.(eachcol(result.draws))))
-    chns = MCMCChains.Chains(draws, var_names; info=(; pathfinder_result=result))
+    chns_info = (; pathfinder_result=result)
+    chns = Accessors.@set draws_to_chains(model, result.draws).info = chns_info
     result_new = Accessors.@set result.draws_transformed = chns
     return result_new
 end

From 82463799f36845926d27aacea3b8c8d2d24d9a06 Mon Sep 17 00:00:00 2001
From: Seth Axen <seth@sethaxen.com>
Date: Thu, 4 Jul 2024 11:24:26 +0200
Subject: [PATCH 07/17] Remove now-unused `flattened_varnames_list` utility

---
 ext/PathfinderTuringExt.jl | 85 --------------------------------------
 src/Pathfinder.jl          |  2 -
 src/integration/turing.jl  | 41 ------------------
 3 files changed, 128 deletions(-)
 delete mode 100644 src/integration/turing.jl

diff --git a/ext/PathfinderTuringExt.jl b/ext/PathfinderTuringExt.jl
index 2d8e1418b..3f6b2fa55 100644
--- a/ext/PathfinderTuringExt.jl
+++ b/ext/PathfinderTuringExt.jl
@@ -9,7 +9,6 @@ if isdefined(Base, :get_extension)
     using Pathfinder: Pathfinder
     using Random: Random
     using Turing: Turing
-    import Pathfinder: flattened_varnames_list
 else  # using Requires
     using ..Accessors: Accessors
     using ..ADTypes: ADTypes
@@ -19,7 +18,6 @@ else  # using Requires
     using ..Pathfinder: Pathfinder
     using ..Random: Random
     using ..Turing: Turing
-    import ..Pathfinder: flattened_varnames_list
 end
 
 """
@@ -60,89 +58,6 @@ function draws_to_chains(model::DynamicPPL.Model, draws::AbstractMatrix)
     return chns
 end
 
-# utilities for working with Turing model parameter names using only the DynamicPPL API
-
-function Pathfinder.flattened_varnames_list(model::DynamicPPL.Model)
-    varnames_ranges = varnames_to_ranges(model)
-    nsyms = maximum(maximum, values(varnames_ranges))
-    syms = Vector{Symbol}(undef, nsyms)
-    for (var_name, range) in varnames_to_ranges(model)
-        sym = Symbol(var_name)
-        if length(range) == 1
-            syms[range[begin]] = sym
-            continue
-        end
-        for i in eachindex(range)
-            syms[range[i]] = Symbol("$sym[$i]")
-        end
-    end
-    return syms
-end
-
-# code snippet shared by @torfjelde
-"""
-    varnames_to_ranges(model::DynamicPPL.Model)
-    varnames_to_ranges(model::DynamicPPL.VarInfo)
-    varnames_to_ranges(model::DynamicPPL.Metadata)
-
-Get `Dict` mapping variable names in model to their ranges in a corresponding parameter vector.
-
-# Examples
-
-```julia
-julia> @model function demo()
-           s ~ Dirac(1)
-           x = Matrix{Float64}(undef, 2, 4)
-           x[1, 1] ~ Dirac(2)
-           x[2, 1] ~ Dirac(3)
-           x[3] ~ Dirac(4)
-           y ~ Dirac(5)
-           x[4] ~ Dirac(6)
-           x[:, 3] ~ arraydist([Dirac(7), Dirac(8)])
-           x[[2, 1], 4] ~ arraydist([Dirac(9), Dirac(10)])
-           return s, x, y
-       end
-demo (generic function with 2 methods)
-
-julia> demo()()
-(1, Any[2.0 4.0 7 10; 3.0 6.0 8 9], 5)
-
-julia> varnames_to_ranges(demo())
-Dict{AbstractPPL.VarName, UnitRange{Int64}} with 8 entries:
-  s           => 1:1
-  x[4]        => 5:5
-  x[:,3]      => 6:7
-  x[1,1]      => 2:2
-  x[2,1]      => 3:3
-  x[[2, 1],4] => 8:9
-  x[3]        => 4:4
-  y           => 10:10
-```
-"""
-function varnames_to_ranges end
-
-varnames_to_ranges(model::DynamicPPL.Model) = varnames_to_ranges(DynamicPPL.VarInfo(model))
-function varnames_to_ranges(varinfo::DynamicPPL.UntypedVarInfo)
-    return varnames_to_ranges(varinfo.metadata)
-end
-function varnames_to_ranges(varinfo::DynamicPPL.TypedVarInfo)
-    offset = 0
-    dicts = map(varinfo.metadata) do md
-        vns2ranges = varnames_to_ranges(md)
-        vals = collect(values(vns2ranges))
-        vals_offset = map(r -> offset .+ r, vals)
-        offset += reduce((curr, r) -> max(curr, r[end]), vals; init=0)
-        Dict(zip(keys(vns2ranges), vals_offset))
-    end
-
-    return reduce(merge, dicts)
-end
-function varnames_to_ranges(metadata::DynamicPPL.Metadata)
-    idcs = map(Base.Fix1(getindex, metadata.idcs), metadata.vns)
-    ranges = metadata.ranges[idcs]
-    return Dict(zip(metadata.vns, ranges))
-end
-
 """
     transform_to_constrained(
         p::AbstractArray, vi::DynamicPPL.VarInfo, model::DynamicPPL.Model
diff --git a/src/Pathfinder.jl b/src/Pathfinder.jl
index d7177ea42..5ba239272 100644
--- a/src/Pathfinder.jl
+++ b/src/Pathfinder.jl
@@ -50,8 +50,6 @@ include("resample.jl")
 include("singlepath.jl")
 include("multipath.jl")
 
-include("integration/turing.jl")
-
 function __init__()
     Requires.@require AdvancedHMC = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d" begin
         include("integration/advancedhmc.jl")
diff --git a/src/integration/turing.jl b/src/integration/turing.jl
deleted file mode 100644
index 1d8ddd392..000000000
--- a/src/integration/turing.jl
+++ /dev/null
@@ -1,41 +0,0 @@
-"""
-    flattened_varnames_list(model::DynamicPPL.Model) -> Vector{Symbol}
-
-Get a vector of varnames as `Symbol`s with one-to-one correspondence to the
-flattened parameter vector.
-
-!!! note
-    This function is only available when Turing has been loaded.
-
-# Examples
-
-```julia
-julia> @model function demo()
-           s ~ Dirac(1)
-           x = Matrix{Float64}(undef, 2, 4)
-           x[1, 1] ~ Dirac(2)
-           x[2, 1] ~ Dirac(3)
-           x[3] ~ Dirac(4)
-           y ~ Dirac(5)
-           x[4] ~ Dirac(6)
-           x[:, 3] ~ arraydist([Dirac(7), Dirac(8)])
-           x[[2, 1], 4] ~ arraydist([Dirac(9), Dirac(10)])
-           return s, x, y
-       end
-demo (generic function with 2 methods)
-
-julia> flattened_varnames_list(demo())
-10-element Vector{Symbol}:
- :s
- Symbol("x[1,1]")
- Symbol("x[2,1]")
- Symbol("x[3]")
- Symbol("x[4]")
- Symbol("x[:,3][1]")
- Symbol("x[:,3][2]")
- Symbol("x[[2, 1],4][1]")
- Symbol("x[[2, 1],4][2]")
- :y
-```
-"""
-function flattened_varnames_list end

From 71b8bc1092261e29804e4e8642d4d7891215075c Mon Sep 17 00:00:00 2001
From: Seth Axen <seth@sethaxen.com>
Date: Thu, 4 Jul 2024 11:27:17 +0200
Subject: [PATCH 08/17] Simplify Turing integration using LogDensityProblems

---
 ext/PathfinderTuringExt.jl | 90 +++++---------------------------------
 1 file changed, 12 insertions(+), 78 deletions(-)

diff --git a/ext/PathfinderTuringExt.jl b/ext/PathfinderTuringExt.jl
index 3f6b2fa55..bae0dfee0 100644
--- a/ext/PathfinderTuringExt.jl
+++ b/ext/PathfinderTuringExt.jl
@@ -2,21 +2,15 @@ module PathfinderTuringExt
 
 if isdefined(Base, :get_extension)
     using Accessors: Accessors
-    using ADTypes: ADTypes
     using DynamicPPL: DynamicPPL
     using MCMCChains: MCMCChains
-    using Optimization: Optimization
     using Pathfinder: Pathfinder
-    using Random: Random
     using Turing: Turing
 else  # using Requires
     using ..Accessors: Accessors
-    using ..ADTypes: ADTypes
     using ..DynamicPPL: DynamicPPL
     using ..MCMCChains: MCMCChains
-    using ..Optimization: Optimization
     using ..Pathfinder: Pathfinder
-    using ..Random: Random
     using ..Turing: Turing
 end
 
@@ -77,81 +71,21 @@ function transform_to_constrained(
     return p
 end
 
-"""
-    set_up_model_optimisation(model::DynamicPPL.Model, init)
-
-Create the necessary pieces for running optimisation on `model`.
-
-Returns
-* An `Optimization.OptimizationFunction` that evaluates the log density of the model and its
-gradient in the unconstrained space.
-* The initial value `init` transformed to unconstrained space.
-* A function `transform_result` that transforms the results back to constrained space. It
-takes a single vector argument.
-"""
-function set_up_model_optimisation(model::DynamicPPL.Model, init)
-    # The inner context deterimines whether we are solving MAP or MLE.
-    inner_context = DynamicPPL.DefaultContext()
-    ctx = Turing.Optimisation.OptimizationContext(inner_context)
-    log_density = Turing.Optimisation.OptimLogDensity(model, ctx)
-    # Initialise the varinfo with the initial value and then transform to unconstrained
-    # space.
-    Accessors.@set log_density.varinfo = DynamicPPL.unflatten(log_density.varinfo, init)
-    transformed_varinfo = DynamicPPL.link(log_density.varinfo, log_density.model)
-    log_density = Accessors.@set log_density.varinfo = transformed_varinfo
-    init = log_density.varinfo[:]
-    # Create a function that applies the appropriate inverse transformation to results, to
-    # bring them back to constrained space.
-    transform_result(p) = transform_to_constrained(p, log_density.varinfo, model)
-    f = Optimization.OptimizationFunction(
-        (x, _) -> log_density(x),;
-        grad = (G,x,p) -> log_density(nothing, G, x),
-    )
-    return f, init, transform_result
-end
-
-function Pathfinder.pathfinder(
-    model::DynamicPPL.Model;
-    rng=Random.GLOBAL_RNG,
-    init_scale=2,
-    init_sampler=Pathfinder.UniformSampler(init_scale),
-    init=nothing,
-    adtype::ADTypes.AbstractADType=Pathfinder.default_ad(),
-    kwargs...,
-)
-    # If no initial value is provided, sample from prior.
-    init = init === nothing ? rand(Vector, model) : init
-    f, init, transform_result = set_up_model_optimisation(model, init)
-    prob = Optimization.OptimizationProblem(f, init)
-    init_sampler(rng, init)
-    result = Pathfinder.pathfinder(prob; rng, input=model, kwargs...)
-    chns_info = (; pathfinder_result=result)
-    chns = Accessors.@set draws_to_chains(model, result.draws).info = chns_info
-    result_new = Accessors.@set result.draws_transformed = chns
+function Pathfinder.pathfinder(model::DynamicPPL.Model; kwargs...)
+    log_density_problem = create_log_density_problem(model)
+    result = Pathfinder.pathfinder(log_density_problem; input=model, kwargs...)
+    chains_info = (; pathfinder_result=result)
+    chains = Accessors.@set draws_to_chains(model, result.draws).info = chains_info
+    result_new = Accessors.@set result.draws_transformed = chains
     return result_new
 end
 
-function Pathfinder.multipathfinder(
-    model::DynamicPPL.Model,
-    ndraws::Int;
-    rng=Random.GLOBAL_RNG,
-    init_scale=2,
-    init_sampler=Pathfinder.UniformSampler(init_scale),
-    nruns::Int,
-    adtype=Pathfinder.default_ad(),
-    kwargs...,
-)
-    # Sample from prior.
-    init1 = rand(Vector, model)
-    fun, init1, transform_result = set_up_model_optimisation(model, init1)
-    init = [init_sampler(rng, init1)]
-    for _ in 2:nruns
-        push!(init, init_sampler(rng, deepcopy(init1)))
-    end
-    result = Pathfinder.multipathfinder(fun, ndraws; rng, input=model, init, kwargs...)
-    chns_info = (; pathfinder_result=result)
-    chns = Accessors.@set draws_to_chains(model, result.draws).info = chns_info
-    result_new = Accessors.@set result.draws_transformed = chns
+function Pathfinder.multipathfinder(model::DynamicPPL.Model, ndraws::Int; kwargs...)
+    log_density_problem = create_log_density_problem(model)
+    result = Pathfinder.multipathfinder(log_density_problem, ndraws; input=model, kwargs...)
+    chains_info = (; pathfinder_result=result)
+    chains = Accessors.@set draws_to_chains(model, result.draws).info = chains_info
+    result_new = Accessors.@set result.draws_transformed = chains
     return result_new
 end
 

From 82be18c1e135cc7df80a892fde2d94386e16f0d1 Mon Sep 17 00:00:00 2001
From: Seth Axen <seth@sethaxen.com>
Date: Thu, 4 Jul 2024 11:40:56 +0200
Subject: [PATCH 09/17] Bump compat lower bounds for mutual compatibility

---
 Project.toml                         | 8 ++++----
 docs/Project.toml                    | 2 +-
 test/integration/Turing/Project.toml | 2 +-
 3 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/Project.toml b/Project.toml
index b42d978d5..fbd88fe8a 100644
--- a/Project.toml
+++ b/Project.toml
@@ -41,8 +41,8 @@ ADTypes = "0.2, 1"
 Accessors = "0.1.12"
 Distributions = "0.25.87"
 DynamicHMC = "3.4.0"
-DynamicPPL = "0.24.7, 0.25, 0.27"
-Folds = "0.2.2"
+DynamicPPL = "0.25.1, 0.27"
+Folds = "0.2.9"
 ForwardDiff = "0.10.19"
 IrrationalConstants = "0.1.1, 0.2"
 LinearAlgebra = "1.6"
@@ -61,8 +61,8 @@ ReverseDiff = "1.4.5"
 SciMLBase = "1.95.0, 2"
 Statistics = "1.6"
 StatsBase = "0.33.7, 0.34"
-Transducers = "0.4.66"
-Turing = "0.30.5, 0.31, 0.32, 0.33"
+Transducers = "0.4.81"
+Turing = "0.31, 0.32, 0.33"
 UnPack = "1"
 julia = "1.6"
 
diff --git a/docs/Project.toml b/docs/Project.toml
index e0963dd3d..f4f316419 100644
--- a/docs/Project.toml
+++ b/docs/Project.toml
@@ -30,4 +30,4 @@ StatsFuns = "1"
 StatsPlots = "0.14.21, 0.15"
 TransformVariables = "0.6.2, 0.7, 0.8"
 TransformedLogDensities = "1.0.2"
-Turing = "0.30.5, 0.31, 0.32, 0.33"
+Turing = "0.31, 0.32, 0.33"
diff --git a/test/integration/Turing/Project.toml b/test/integration/Turing/Project.toml
index c40df5265..8411105e2 100644
--- a/test/integration/Turing/Project.toml
+++ b/test/integration/Turing/Project.toml
@@ -8,5 +8,5 @@ Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
 [compat]
 LogDensityProblems = "2.1.0"
 Pathfinder = "0.9"
-Turing = "0.30.5, 0.31, 0.32, 0.33"
+Turing = "0.31, 0.32, 0.33"
 julia = "1.6"

From 6df887cab10a5b9a10a7745a79c0bbc09ec164e4 Mon Sep 17 00:00:00 2001
From: Seth Axen <seth@sethaxen.com>
Date: Thu, 4 Jul 2024 13:07:18 +0200
Subject: [PATCH 10/17] Bump Turing lower bounds

Ensures `getparams` uses the changes introduced in TuringLang/Turing.jl#2202
---
 Project.toml                         | 2 +-
 docs/Project.toml                    | 2 +-
 test/integration/Turing/Project.toml | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/Project.toml b/Project.toml
index fbd88fe8a..706b94433 100644
--- a/Project.toml
+++ b/Project.toml
@@ -62,7 +62,7 @@ SciMLBase = "1.95.0, 2"
 Statistics = "1.6"
 StatsBase = "0.33.7, 0.34"
 Transducers = "0.4.81"
-Turing = "0.31, 0.32, 0.33"
+Turing = "0.31.4, 0.32, 0.33"
 UnPack = "1"
 julia = "1.6"
 
diff --git a/docs/Project.toml b/docs/Project.toml
index f4f316419..1bd3cf4d5 100644
--- a/docs/Project.toml
+++ b/docs/Project.toml
@@ -30,4 +30,4 @@ StatsFuns = "1"
 StatsPlots = "0.14.21, 0.15"
 TransformVariables = "0.6.2, 0.7, 0.8"
 TransformedLogDensities = "1.0.2"
-Turing = "0.31, 0.32, 0.33"
+Turing = "0.31.4, 0.32, 0.33"
diff --git a/test/integration/Turing/Project.toml b/test/integration/Turing/Project.toml
index 8411105e2..36e1c1b69 100644
--- a/test/integration/Turing/Project.toml
+++ b/test/integration/Turing/Project.toml
@@ -8,5 +8,5 @@ Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
 [compat]
 LogDensityProblems = "2.1.0"
 Pathfinder = "0.9"
-Turing = "0.31, 0.32, 0.33"
+Turing = "0.31.4, 0.32, 0.33"
 julia = "1.6"

From f74e1f37bc613c7796cefeb6f7491a1bdc7f3a59 Mon Sep 17 00:00:00 2001
From: Seth Axen <seth@sethaxen.com>
Date: Thu, 4 Jul 2024 13:18:49 +0200
Subject: [PATCH 11/17] Reorganize Turing tests

---
 test/integration/Turing/runtests.jl | 74 ++++++++++++++++-------------
 1 file changed, 41 insertions(+), 33 deletions(-)

diff --git a/test/integration/Turing/runtests.jl b/test/integration/Turing/runtests.jl
index c33bcb104..a6f5120e3 100644
--- a/test/integration/Turing/runtests.jl
+++ b/test/integration/Turing/runtests.jl
@@ -9,6 +9,7 @@ end
 
 Random.seed!(0)
 
+#! format: off
 @model function regression_model(x, y)
     σ ~ truncated(Normal(); lower=0)
     α ~ Normal()
@@ -22,11 +23,16 @@ end
 @model function dynamic_const_model()
     lb ~ Uniform(0, 0.1)
     ub ~ Uniform(0.11, 0.2)
-    return x ~ Bijectors.transformed(
+    x ~ Bijectors.transformed(
         Normal(0, 1), Bijectors.inverse(Bijectors.Logit(lb, ub))
     )
 end
 
+@model function transformed_model(dist, bijector)
+    y ~ Bijectors.transformed(dist, bijector)
+end
+#! format: on
+
 @testset "Turing integration" begin
     @testset "create_log_density_problem" begin
         @testset for bijector in [elementwise(log), Bijectors.SimplexBijector()],
@@ -36,10 +42,8 @@ end
             binv = Bijectors.inverse(bijector)
             dist = filldist(udist, n)
             dist_trans = Bijectors.transformed(dist, binv)
-            @model function model()
-                return y ~ dist_trans
-            end
-            prob = PathfinderTuringExt.create_log_density_problem(model())
+            model = transformed_model(dist, binv)
+            prob = PathfinderTuringExt.create_log_density_problem(model)
             @test LogDensityProblems.capabilities(prob) isa
                 LogDensityProblems.LogDensityOrder{0}
             x = rand(n, 10)
@@ -61,34 +65,38 @@ end
         @test all(chns[:, :lb, 1] .< chns[:, :x, 1] .< chns[:, :ub, 1])
     end
 
-    x = 0:0.01:1
-    y = sin.(x) .+ randn.() .* 0.2 .+ x
-    X = [x x .^ 2 x .^ 3]
-    model = regression_model(X, y)
-    expected_param_names = Symbol.(["α", "β[1]", "β[2]", "β[3]", "σ"])
+    @testset "integration tests" begin
+        @testset "regression model" begin
+            x = 0:0.01:1
+            y = sin.(x) .+ randn.() .* 0.2 .+ x
+            X = [x x .^ 2 x .^ 3]
+            model = regression_model(X, y)
+            expected_param_names = Symbol.(["α", "β[1]", "β[2]", "β[3]", "σ"])
 
-    result = pathfinder(model; ndraws=10_000)
-    @test result isa PathfinderResult
-    @test result.input === model
-    @test size(result.draws) == (5, 10_000)
-    @test result.draws_transformed isa MCMCChains.Chains
-    @test result.draws_transformed.info.pathfinder_result isa PathfinderResult
-    @test sort(names(result.draws_transformed)) == expected_param_names
-    @test all(>(0), result.draws_transformed[:σ])
-    init_params = Vector(result.draws_transformed.value[1, :, 1])
-    chns = sample(model, NUTS(), 10_000; init_params, progress=false)
-    @test mean(chns).nt.mean ≈ mean(result.draws_transformed).nt.mean rtol = 0.1
+            result = pathfinder(model; ndraws=10_000)
+            @test result isa PathfinderResult
+            @test result.input === model
+            @test size(result.draws) == (5, 10_000)
+            @test result.draws_transformed isa MCMCChains.Chains
+            @test result.draws_transformed.info.pathfinder_result isa PathfinderResult
+            @test sort(names(result.draws_transformed)) == expected_param_names
+            @test all(>(0), result.draws_transformed[:σ])
+            init_params = Vector(result.draws_transformed.value[1, :, 1])
+            chns = sample(model, NUTS(), 10_000; init_params, progress=false)
+            @test mean(chns).nt.mean ≈ mean(result.draws_transformed).nt.mean rtol = 0.1
 
-    result = multipathfinder(model, 10_000; nruns=4)
-    @test result isa MultiPathfinderResult
-    @test result.input === model
-    @test size(result.draws) == (5, 10_000)
-    @test length(result.pathfinder_results) == 4
-    @test result.draws_transformed isa MCMCChains.Chains
-    @test result.draws_transformed.info.pathfinder_result isa MultiPathfinderResult
-    @test sort(names(result.draws_transformed)) == expected_param_names
-    @test all(>(0), result.draws_transformed[:σ])
-    init_params = Vector(result.draws_transformed.value[1, :, 1])
-    chns = sample(model, NUTS(), 10_000; init_params, progress=false)
-    @test mean(chns).nt.mean ≈ mean(result.draws_transformed).nt.mean rtol = 0.1
+            result = multipathfinder(model, 10_000; nruns=4)
+            @test result isa MultiPathfinderResult
+            @test result.input === model
+            @test size(result.draws) == (5, 10_000)
+            @test length(result.pathfinder_results) == 4
+            @test result.draws_transformed isa MCMCChains.Chains
+            @test result.draws_transformed.info.pathfinder_result isa MultiPathfinderResult
+            @test sort(names(result.draws_transformed)) == expected_param_names
+            @test all(>(0), result.draws_transformed[:σ])
+            init_params = Vector(result.draws_transformed.value[1, :, 1])
+            chns = sample(model, NUTS(), 10_000; init_params, progress=false)
+            @test mean(chns).nt.mean ≈ mean(result.draws_transformed).nt.mean rtol = 0.1
+        end
+    end
 end

From 109b89c36fbd6a4f47318b15ad792fad7b143c5f Mon Sep 17 00:00:00 2001
From: Seth Axen <seth@sethaxen.com>
Date: Thu, 4 Jul 2024 13:27:38 +0200
Subject: [PATCH 12/17] Bump DynamicPPL lower bound

---
 Project.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Project.toml b/Project.toml
index 706b94433..3cbbe11d0 100644
--- a/Project.toml
+++ b/Project.toml
@@ -41,7 +41,7 @@ ADTypes = "0.2, 1"
 Accessors = "0.1.12"
 Distributions = "0.25.87"
 DynamicHMC = "3.4.0"
-DynamicPPL = "0.25.1, 0.27"
+DynamicPPL = "0.25.2, 0.27"
 Folds = "0.2.9"
 ForwardDiff = "0.10.19"
 IrrationalConstants = "0.1.1, 0.2"

From 101caa077ca0686b50aee6dc00e6d7f61c285b8e Mon Sep 17 00:00:00 2001
From: Seth Axen <seth@sethaxen.com>
Date: Thu, 4 Jul 2024 13:53:21 +0200
Subject: [PATCH 13/17] Also add Chains to individual paths in multipathfinder
 result

---
 ext/PathfinderTuringExt.jl          | 17 ++++++++++++++++-
 test/integration/Turing/runtests.jl |  4 ++++
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/ext/PathfinderTuringExt.jl b/ext/PathfinderTuringExt.jl
index bae0dfee0..aa3906c85 100644
--- a/ext/PathfinderTuringExt.jl
+++ b/ext/PathfinderTuringExt.jl
@@ -74,6 +74,8 @@ end
 function Pathfinder.pathfinder(model::DynamicPPL.Model; kwargs...)
     log_density_problem = create_log_density_problem(model)
     result = Pathfinder.pathfinder(log_density_problem; input=model, kwargs...)
+
+    # add transformed draws as Chains
     chains_info = (; pathfinder_result=result)
     chains = Accessors.@set draws_to_chains(model, result.draws).info = chains_info
     result_new = Accessors.@set result.draws_transformed = chains
@@ -83,9 +85,22 @@ end
 function Pathfinder.multipathfinder(model::DynamicPPL.Model, ndraws::Int; kwargs...)
     log_density_problem = create_log_density_problem(model)
     result = Pathfinder.multipathfinder(log_density_problem, ndraws; input=model, kwargs...)
+
+    # add transformed draws as Chains
     chains_info = (; pathfinder_result=result)
     chains = Accessors.@set draws_to_chains(model, result.draws).info = chains_info
-    result_new = Accessors.@set result.draws_transformed = chains
+
+    # add transformed draws as Chains for each individual path
+    single_path_results_new = map(result.pathfinder_results) do r
+        single_chains_info = (; pathfinder_result=r)
+        single_chains = Accessors.@set draws_to_chains(model, r.draws).info =
+            single_chains_info
+        r_new = Accessors.@set r.draws_transformed = single_chains
+        return r_new
+    end
+
+    result_new = Accessors.@set (Accessors.@set result.draws_transformed =
+        chains).pathfinder_results = single_path_results_new
     return result_new
 end
 
diff --git a/test/integration/Turing/runtests.jl b/test/integration/Turing/runtests.jl
index a6f5120e3..01fd92c3e 100644
--- a/test/integration/Turing/runtests.jl
+++ b/test/integration/Turing/runtests.jl
@@ -97,6 +97,10 @@ end
             init_params = Vector(result.draws_transformed.value[1, :, 1])
             chns = sample(model, NUTS(), 10_000; init_params, progress=false)
             @test mean(chns).nt.mean ≈ mean(result.draws_transformed).nt.mean rtol = 0.1
+
+            for r in result.pathfinder_results
+                @test r.draws_transformed isa MCMCChains.Chains
+            end
         end
     end
 end

From dd7e402ff3a91009e209705f307cdaf4cbdc9fab Mon Sep 17 00:00:00 2001
From: Seth Axen <seth@sethaxen.com>
Date: Thu, 4 Jul 2024 13:53:41 +0200
Subject: [PATCH 14/17] Test additional models

---
 test/integration/Turing/Project.toml |  1 +
 test/integration/Turing/runtests.jl  | 39 +++++++++++++++++++++++++++-
 2 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/test/integration/Turing/Project.toml b/test/integration/Turing/Project.toml
index 36e1c1b69..8c0b5df09 100644
--- a/test/integration/Turing/Project.toml
+++ b/test/integration/Turing/Project.toml
@@ -1,4 +1,5 @@
 [deps]
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
 Pathfinder = "b1d3bc72-d0e7-4279-b92f-7fa5d6d2d454"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
diff --git a/test/integration/Turing/runtests.jl b/test/integration/Turing/runtests.jl
index 01fd92c3e..58799c8da 100644
--- a/test/integration/Turing/runtests.jl
+++ b/test/integration/Turing/runtests.jl
@@ -1,4 +1,4 @@
-using LogDensityProblems, Pathfinder, Random, Test, Turing
+using LogDensityProblems, LinearAlgebra, Pathfinder, Random, Test, Turing
 using Turing.Bijectors
 
 if isdefined(Base, :get_extension)
@@ -103,4 +103,41 @@ end
             end
         end
     end
+
+    @testset "transformed IID normal solved exactly" begin
+        @testset for bijector in [elementwise(log), Bijectors.SimplexBijector()],
+            udist in [Normal(1, 2), Normal(3, 4)],
+            n in [1, 5]
+
+            binv = Bijectors.inverse(bijector)
+            dist = filldist(udist, n)
+            model = transformed_model(dist, binv)
+            result = pathfinder(model)
+            @test mean(result.fit_distribution) ≈ fill(mean(udist), n)
+            @test cov(result.fit_distribution) ≈ Diagonal(fill(var(udist), n))
+
+            result = multipathfinder(model, 100; nruns=4, ndraws_per_run=100)
+            @test result isa MultiPathfinderResult
+            for r in result.pathfinder_results
+                @test mean(r.fit_distribution) ≈ fill(mean(udist), n)
+                @test cov(r.fit_distribution) ≈ Diagonal(fill(var(udist), n))
+            end
+        end
+    end
+
+    @testset "models with dynamic constraints successfully fitted" begin
+        result = pathfinder(dynamic_const_model(); ndraws=10_000)
+        chns = result.draws_transformed
+        @test all(0 .< chns[:, :lb, 1] .< 0.1)
+        @test all(0.11 .< chns[:, :ub, 1] .< 0.2)
+        @test all(chns[:, :lb, 1] .< chns[:, :x, 1] .< chns[:, :ub, 1])
+
+        result = multipathfinder(dynamic_const_model(), 10_000; nruns=4)
+        for r in result.pathfinder_results
+            chns = r.draws_transformed
+            @test all(0 .< chns[:, :lb, 1] .< 0.1)
+            @test all(0.11 .< chns[:, :ub, 1] .< 0.2)
+            @test all(chns[:, :lb, 1] .< chns[:, :x, 1] .< chns[:, :ub, 1])
+        end
+    end
 end

From d50253ddc259e81fbc4158540c27f0a517784fc6 Mon Sep 17 00:00:00 2001
From: Seth Axen <seth@sethaxen.com>
Date: Fri, 5 Jul 2024 14:15:02 +0200
Subject: [PATCH 15/17] Remove unused method

---
 ext/PathfinderTuringExt.jl | 19 -------------------
 1 file changed, 19 deletions(-)

diff --git a/ext/PathfinderTuringExt.jl b/ext/PathfinderTuringExt.jl
index aa3906c85..d959f5ec2 100644
--- a/ext/PathfinderTuringExt.jl
+++ b/ext/PathfinderTuringExt.jl
@@ -52,25 +52,6 @@ function draws_to_chains(model::DynamicPPL.Model, draws::AbstractMatrix)
     return chns
 end
 
-"""
-    transform_to_constrained(
-        p::AbstractArray, vi::DynamicPPL.VarInfo, model::DynamicPPL.Model
-    )
-
-Transform a vector of parameters `p` from unconstrained to constrained space.
-"""
-function transform_to_constrained(
-    p::AbstractArray, vi::DynamicPPL.VarInfo, model::DynamicPPL.Model
-)
-    p = copy(p)
-    @assert DynamicPPL.istrans(vi)
-    vi = DynamicPPL.unflatten(vi, p)
-    p .= DynamicPPL.invlink!!(vi, model)[:]
-    # Restore the linking, since we mutated vi.
-    DynamicPPL.link!!(vi, model)
-    return p
-end
-
 function Pathfinder.pathfinder(model::DynamicPPL.Model; kwargs...)
     log_density_problem = create_log_density_problem(model)
     result = Pathfinder.pathfinder(log_density_problem; input=model, kwargs...)

From 625d5f90797cc92861a0b42f81b23ee5b41c6007 Mon Sep 17 00:00:00 2001
From: Seth Axen <seth@sethaxen.com>
Date: Fri, 16 Aug 2024 16:39:18 +0200
Subject: [PATCH 16/17] Bump AdvancedHMC integration test compat  to AHMC v0.6

---
 test/integration/AdvancedHMC/Project.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/integration/AdvancedHMC/Project.toml b/test/integration/AdvancedHMC/Project.toml
index 02e9ff1d0..d90ca4f6f 100644
--- a/test/integration/AdvancedHMC/Project.toml
+++ b/test/integration/AdvancedHMC/Project.toml
@@ -16,7 +16,7 @@ TransformVariables = "84d833dd-6860-57f9-a1a7-6da5db126cff"
 TransformedLogDensities = "f9bc47f6-f3f8-4f3b-ab21-f8bc73906f26"
 
 [compat]
-AdvancedHMC = "0.4, 0.5.2, 0.6"
+AdvancedHMC = "0.6"
 Distributions = "0.25.87"
 ForwardDiff = "0.10.19"
 LogDensityProblems = "2.1.0"

From 823a515dc895d79fa99ecf5c5cd5aa6dc914f27a Mon Sep 17 00:00:00 2001
From: Seth Axen <seth@sethaxen.com>
Date: Fri, 16 Aug 2024 16:40:22 +0200
Subject: [PATCH 17/17] Increment patch number

---
 Project.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Project.toml b/Project.toml
index 3cbbe11d0..642dfd8c2 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "Pathfinder"
 uuid = "b1d3bc72-d0e7-4279-b92f-7fa5d6d2d454"
 authors = ["Seth Axen <seth.axen@gmail.com> and contributors"]
-version = "0.9.0-DEV"
+version = "0.9.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"