Skip to content

Commit 4a11560

Browse files
penelopeysmpenelopeysm
authored
Allow generation of ParamsWithStats from FastLDF plus parameters, and also bundle_samples (#1129)
* Implement `ParamsWithStats` for `FastLDF` * Add comments * Implement `bundle_samples` for ParamsWithStats -> MCMCChains * Remove redundant comment * don't need Statistics?
1 parent 3cd8d34 commit 4a11560

5 files changed

Lines changed: 180 additions & 25 deletions

File tree

‎ext/DynamicPPLMCMCChainsExt.jl‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -140,6 +140,43 @@ function AbstractMCMC.to_samples(
140140
end
141141
end
142142

143+
function AbstractMCMC.bundle_samples(
144+
ts::Vector{<:DynamicPPL.ParamsWithStats},
145+
model::DynamicPPL.Model,
146+
spl::AbstractMCMC.AbstractSampler,
147+
state,
148+
chain_type::Type{MCMCChains.Chains};
149+
save_state=false,
150+
stats=missing,
151+
sort_chain=false,
152+
discard_initial=0,
153+
thinning=1,
154+
kwargs...,
155+
)
156+
bare_chain = AbstractMCMC.from_samples(MCMCChains.Chains, reshape(ts, :, 1))
157+
158+
# Add additional MCMC-specific info
159+
info = bare_chain.info
160+
if save_state
161+
info = merge(info, (model=model, sampler=spl, samplerstate=state))
162+
end
163+
if !ismissing(stats)
164+
info = merge(info, (start_time=stats.start, stop_time=stats.stop))
165+
end
166+
167+
# Reconstruct the chain with the extra information
168+
# Yeah, this is quite ugly. Blame MCMCChains.
169+
chain = MCMCChains.Chains(
170+
bare_chain.value.data,
171+
names(bare_chain),
172+
bare_chain.name_map;
173+
info=info,
174+
start=discard_initial + 1,
175+
thin=thinning,
176+
)
177+
return sort_chain ? sort(chain) : chain
178+
end
179+
143180
"""
144181
predict([rng::AbstractRNG,] model::Model, chain::MCMCChains.Chains; include_all=false)
145182

‎src/DynamicPPL.jl‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -202,14 +202,14 @@ include("logdensityfunction.jl")
202202
include("model_utils.jl")
203203
include("extract_priors.jl")
204204
include("values_as_in_model.jl")
205+
include("experimental.jl")
205206
include("chains.jl")
206207
include("bijector.jl")
207208

208209
include("debug_utils.jl")
209210
using .DebugUtils
210211
include("test_utils.jl")
211212

212-
include("experimental.jl")
213213
include("deprecated.jl")
214214

215215
if isdefined(Base.Experimental, :register_error_hint)

‎src/chains.jl‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -133,3 +133,60 @@ function ParamsWithStats(
133133
end
134134
return ParamsWithStats(params, stats)
135135
end
136+
137+
"""
138+
ParamsWithStats(
139+
param_vector::AbstractVector,
140+
ldf::DynamicPPL.Experimental.FastLDF,
141+
stats::NamedTuple=NamedTuple();
142+
include_colon_eq::Bool=true,
143+
include_log_probs::Bool=true,
144+
)
145+
146+
Generate a `ParamsWithStats` by re-evaluating the given `ldf` with the provided
147+
`param_vector`.
148+
149+
This method is intended to replace the old method of obtaining parameters and statistics
150+
via `unflatten` plus re-evaluation. It is faster for two reasons:
151+
152+
1. It does not rely on `deepcopy`-ing the VarInfo object (this used to be mandatory as
153+
otherwise re-evaluation would mutate the VarInfo, rendering it unusable for subsequent
154+
MCMC iterations).
155+
2. The re-evaluation is faster as it uses `OnlyAccsVarInfo`.
156+
"""
157+
function ParamsWithStats(
158+
param_vector::AbstractVector,
159+
ldf::DynamicPPL.Experimental.FastLDF,
160+
stats::NamedTuple=NamedTuple();
161+
include_colon_eq::Bool=true,
162+
include_log_probs::Bool=true,
163+
)
164+
strategy = InitFromParams(
165+
VectorWithRanges(ldf._iden_varname_ranges, ldf._varname_ranges, param_vector),
166+
nothing,
167+
)
168+
accs = if include_log_probs
169+
(
170+
DynamicPPL.LogPriorAccumulator(),
171+
DynamicPPL.LogLikelihoodAccumulator(),
172+
DynamicPPL.ValuesAsInModelAccumulator(include_colon_eq),
173+
)
174+
else
175+
(DynamicPPL.ValuesAsInModelAccumulator(include_colon_eq),)
176+
end
177+
_, vi = DynamicPPL.Experimental.fast_evaluate!!(
178+
ldf.model, strategy, AccumulatorTuple(accs)
179+
)
180+
params = DynamicPPL.getacc(vi, Val(:ValuesAsInModel)).values
181+
if include_log_probs
182+
stats = merge(
183+
stats,
184+
(
185+
logprior=DynamicPPL.getlogprior(vi),
186+
loglikelihood=DynamicPPL.getloglikelihood(vi),
187+
lp=DynamicPPL.getlogjoint(vi),
188+
),
189+
)
190+
end
191+
return ParamsWithStats(params, stats)
192+
end

‎src/fasteval.jl‎

Lines changed: 58 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ using DynamicPPL:
33
AccumulatorTuple,
44
InitContext,
55
InitFromParams,
6+
AbstractInitStrategy,
67
LogJacobianAccumulator,
78
LogLikelihoodAccumulator,
89
LogPriorAccumulator,
@@ -28,6 +29,60 @@ using LogDensityProblems: LogDensityProblems
2829
import DifferentiationInterface as DI
2930
using Random: Random
3031

32+
"""
33+
DynamicPPL.Experimental.fast_evaluate!!(
34+
[rng::Random.AbstractRNG,]
35+
model::Model,
36+
strategy::AbstractInitStrategy,
37+
accs::AccumulatorTuple, params::AbstractVector{<:Real}
38+
)
39+
40+
Evaluate a model using parameters obtained via `strategy`, and only computing the results in
41+
the provided accumulators.
42+
43+
It is assumed that the accumulators passed in have been initialised to appropriate values,
44+
as this function will not reset them. The default constructors for each accumulator will do
45+
this for you correctly.
46+
47+
Returns a tuple of the model's return value, plus an `OnlyAccsVarInfo`. Note that the `accs`
48+
argument may be mutated (depending on how the accumulators are implemented); hence the `!!`
49+
in the function name.
50+
"""
51+
@inline function fast_evaluate!!(
52+
# Note that this `@inline` is mandatory for performance. If it's not inlined, it leads
53+
# to extra allocations (even for trivial models) and much slower runtime.
54+
rng::Random.AbstractRNG,
55+
model::Model,
56+
strategy::AbstractInitStrategy,
57+
accs::AccumulatorTuple,
58+
)
59+
ctx = InitContext(rng, strategy)
60+
model = DynamicPPL.setleafcontext(model, ctx)
61+
# Calling `evaluate!!` would be fine, but would lead to an extra call to resetaccs!!,
62+
# which is unnecessary. So we shortcircuit this by simply calling `_evaluate!!`
63+
# directly. To preserve thread-safety we need to reproduce the ThreadSafeVarInfo logic
64+
# here.
65+
# TODO(penelopeysm): This should _not_ check Threads.nthreads(). I still don't know what
66+
# it _should_ do, but this is wrong regardless.
67+
# https://github.com/TuringLang/DynamicPPL.jl/issues/1086
68+
vi = if Threads.nthreads() > 1
69+
param_eltype = DynamicPPL.get_param_eltype(strategy)
70+
accs = map(accs) do acc
71+
DynamicPPL.convert_eltype(float_type_with_fallback(param_eltype), acc)
72+
end
73+
ThreadSafeVarInfo(OnlyAccsVarInfo(accs))
74+
else
75+
OnlyAccsVarInfo(accs)
76+
end
77+
return DynamicPPL._evaluate!!(model, vi)
78+
end
79+
@inline function fast_evaluate!!(
80+
model::Model, strategy::AbstractInitStrategy, accs::AccumulatorTuple
81+
)
82+
# This `@inline` is also mandatory for performance
83+
return fast_evaluate!!(Random.default_rng(), model, strategy, accs)
84+
end
85+
3186
"""
3287
FastLDF(
3388
model::Model,
@@ -213,31 +268,11 @@ struct FastLogDensityAt{M<:Model,F<:Function,N<:NamedTuple}
213268
varname_ranges::Dict{VarName,RangeAndLinked}
214269
end
215270
function (f::FastLogDensityAt)(params::AbstractVector{<:Real})
216-
ctx = InitContext(
217-
Random.default_rng(),
218-
InitFromParams(
219-
VectorWithRanges(f.iden_varname_ranges, f.varname_ranges, params), nothing
220-
),
271+
strategy = InitFromParams(
272+
VectorWithRanges(f.iden_varname_ranges, f.varname_ranges, params), nothing
221273
)
222-
model = DynamicPPL.setleafcontext(f.model, ctx)
223274
accs = fast_ldf_accs(f.getlogdensity)
224-
# Calling `evaluate!!` would be fine, but would lead to an extra call to resetaccs!!,
225-
# which is unnecessary. So we shortcircuit this by simply calling `_evaluate!!`
226-
# directly. To preserve thread-safety we need to reproduce the ThreadSafeVarInfo logic
227-
# here.
228-
# TODO(penelopeysm): This should _not_ check Threads.nthreads(). I still don't know what
229-
# it _should_ do, but this is wrong regardless.
230-
# https://github.com/TuringLang/DynamicPPL.jl/issues/1086
231-
vi = if Threads.nthreads() > 1
232-
accs = map(
233-
acc -> DynamicPPL.convert_eltype(float_type_with_fallback(eltype(params)), acc),
234-
accs,
235-
)
236-
ThreadSafeVarInfo(OnlyAccsVarInfo(accs))
237-
else
238-
OnlyAccsVarInfo(accs)
239-
end
240-
_, vi = DynamicPPL._evaluate!!(model, vi)
275+
_, vi = fast_evaluate!!(f.model, strategy, accs)
241276
return f.getlogdensity(vi)
242277
end
243278

‎test/chains.jl‎

Lines changed: 27 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ using DynamicPPL
44
using Distributions
55
using Test
66

7-
@testset "ParamsWithStats" begin
7+
@testset "ParamsWithStats from VarInfo" begin
88
@model function f(z)
99
x ~ Normal()
1010
y := x + 1
@@ -66,4 +66,30 @@ using Test
6666
end
6767
end
6868

69+
@testset "ParamsWithStats from FastLDF" begin
70+
@testset "$(m.f)" for m in DynamicPPL.TestUtils.DEMO_MODELS
71+
unlinked_vi = VarInfo(m)
72+
@testset "$islinked" for islinked in (false, true)
73+
vi = if islinked
74+
DynamicPPL.link!!(unlinked_vi, m)
75+
else
76+
unlinked_vi
77+
end
78+
params = [x for x in vi[:]]
79+
80+
# Get the ParamsWithStats using FastLDF
81+
fldf = DynamicPPL.Experimental.FastLDF(m, getlogjoint, vi)
82+
ps = ParamsWithStats(params, fldf)
83+
84+
# Check that length of parameters is as expected
85+
@test length(ps.params) == length(keys(vi))
86+
87+
# Iterate over all variables to check that their values match
88+
for vn in keys(vi)
89+
@test ps.params[vn] == vi[vn]
90+
end
91+
end
92+
end
93+
end
94+
6995
end # module

0 commit comments

Comments
 (0)