Traces of transformed variables within Turing.jl models

Nice and thanks for the quick response @dlakelan! For future reference, here is the example from the docs using variational inference:

using Turing
using Optim
using StatsPlots
using Turing: Variational

@model function demo(xs)
    s ~ InverseGamma(2, 3)
    m_shifted ~ Normal(10, √s)
    m = m_shifted - 10

    for i in eachindex(xs) 
        xs[i] ~ Normal(m, √s)
    end

    return (m, )
end


q = vi(model, ADVI(10, 2000))
samples = rand(q, 1000)

#NOTE: constructing the chain becomes slightly more involved with sampled vectors and matrices
_, sym2range = bijector(model, Val(true))
val = vcat([samples[range...,:] for (sym, range) in pairs(sym2range)]...)'
parameter_names = [keys(sym2range)...]
chain = Chains(val, parameter_names)

gq = generated_quantities(model, chain_vi)
m = vec(getindex.(gq, 1))
density!(m, lab="generated quantity (VI)")
vline!([0], lab="true value")

example