Add DynamicPPL extension, use NoCache for performance
#644
Conversation
|
with
|
Codecov ReportAttention: Patch coverage is
📢 Thoughts on this report? Let us know! |
|
Mooncake.jl documentation for PR #644 is available at: |
|
Performance Ratio: |
|
As Penny pointed out in TuringLang/DynamicPPL.jl#856, when a model is declared in a function, it becomes a closure. julia> function create_closure()
@model function closure_model(x)
s ~ InverseGamma(2, 3)
return x ~ Normal(0, s)
end
return closure_model
end
create_closure (generic function with 1 method)
julia> closure_fn = create_closure()
(::var"#closure_model#17") (generic function with 2 methods)
julia> model = closure_fn([1.0, 2.0])
Model{var"#closure_model#17", (:x,), (), (), Tuple{Vector{Float64}}, Tuple{}, DefaultContext}(var"#closure_model#17"(Core.Box(var"#closure_model#17"(#= circular reference @-2 =#))), (x = [1.0, 2.0],), NamedTuple(), DefaultContext())
julia> vi = DynamicPPL.VarInfo(Random.default_rng(), model); ldf = DynamicPPL.LogDensityFunction(model, vi, DynamicPPL.DefaultContext()); tangent = zero_tangent(ldf);
julia> tangent.fields.model.fields.f.fields.closure_model.fields.contents.tangent === model_f_tangent
trueThis is a special case where tangent might refer to itself, and reason for Mooncake.jl/ext/MooncakeDynamicPPLExt.jl Lines 88 to 96 in 2f9a1c1 |
There was a problem hiding this comment.
The code looks sensible. Quick comment on the interface, I suggest we introduce a function
Mooncake._build_iddict(::Type{Tangent}) = IdDict{Any,Bool}()Then, DynamicPPL can overload this function for various DynamicPPL related tangent types, e.g. LDF, Model, VarInfo:
Mooncake._build_iddict(::Type{Tangent{VarInfo}}) = NoCache()EDIT: We could also have a two-argument version to pass runtime information
Mooncake._build_iddict(::Type{Tangent{VarInfo}}, extra) = ...
|
The interface is fine, but achieving feature parity with the current code is difficult to do in a clean way. The main reason is that closure detection requires runtime information. We can go for a hybrid approach, which somewhat defeats the purpose of introducing an interface function. Or we can opt for cleanliness and leave some performance gains on the table. |
Yes, that is okay if we do it robustly, excluding invalid cases like closures. |
|
I haven't yet looked at the code, but I wanted to point out that if Mooncake has a DynamicPPLExt (and thus a compat entry for DynamicPPL), then breaking changes of DynamicPPL cannot be tested with Mooncake. That isn't a massive problem in and of itself; we could separate the Mooncake part out from the test suite (in exactly the same way Enzyme is). However, because breaking changes in DPPL are much more common than breaking changes in Mooncake, I think it may be easier to implement this code inside DynamicPPLMooncakeExt rather than the other way around. (see: TuringLang/DynamicPPL.jl#740 where this issue was discussed) |
Co-authored-by: Penelope Yong <penelopeysm@gmail.com> Signed-off-by: Xianda Sun <5433119+sunxd3@users.noreply.github.com>
Co-authored-by: Penelope Yong <penelopeysm@gmail.com> Signed-off-by: Xianda Sun <5433119+sunxd3@users.noreply.github.com>
|
Closed in favour of TuringLang/DynamicPPL.jl#975 |
ref #552