diff --git a/HISTORY.md b/HISTORY.md index f1245ed48..7f90d3af7 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -1,3 +1,9 @@ +# 0.42.1 + +Fixed a type-inference failure that made nested submodels (a `~ to_submodel(...)` statement inside a model that is itself evaluated as a submodel) very slow. + +Previously, evaluating a submodel recursed through the shared `_evaluate!!(::Model, ::AbstractVarInfo)` method. Each level of nesting adds another context layer to the `Model` type, which tripped Julia's type-inference recursion limit: from the first level of nesting onwards the return type was inferred as `Any` and evaluation fell back to runtime dispatch, slowing down both primal and gradient evaluation (the primal slowdown was roughly 15x, growing with nesting depth). Submodel evaluation now calls the model function directly, keeping nested submodels type-stable. See [Turing.jl#2844](https://github.com/TuringLang/Turing.jl/issues/2844). + # 0.42.0 `LogDensityFunction` now performs AD preparation through AbstractPPL's `prepare` / `value_and_gradient!!` interface instead of calling DifferentiationInterface directly. Internally this removes the `_use_closure` heuristic and the explicit `DI.Constant` plumbing; the choice between closure and constants now lives in AbstractPPL. diff --git a/Project.toml b/Project.toml index 07fdb81ea..cbabc7812 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "DynamicPPL" uuid = "366bfd00-2699-11ea-058f-f148b4cae6d8" -version = "0.42" +version = "0.42.1" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" diff --git a/src/submodel.jl b/src/submodel.jl index 97618db66..f8acf83c0 100644 --- a/src/submodel.jl +++ b/src/submodel.jl @@ -181,9 +181,17 @@ function _evaluate!!( # (4) Finally, we need to store that context inside the submodel. model = contextualize(submodel.model, eval_context) - # Once that's all set up nicely, we can just _evaluate!! the wrapped model. This - # returns a tuple of submodel.model's return value and the new varinfo. - return _evaluate!!(model, vi) + # Evaluate the wrapped model. These two lines are a verbatim copy of the body of + # `_evaluate!!(model::Model, ::AbstractVarInfo)` (in `model.jl`), and the duplication is + # deliberate: DO NOT replace them with `return _evaluate!!(model, vi)`. Each level of + # submodel nesting grows the contextualised `Model`'s context type, and routing the + # recursion through the shared `_evaluate!!(::Model, ...)` method trips Julia's recursion + # limiter, which widens the `Model` argument to abstract and collapses the return type to + # `Any`. Calling `model.f` directly avoids that. See + # https://github.com/TuringLang/DynamicPPL.jl/pull/1427 and + # https://github.com/TuringLang/Turing.jl/issues/2844 for the full explanation. + args, kwargs = make_evaluate_args_and_kwargs(model, vi) + return model.f(args...; kwargs...) end function tilde_observe!!( diff --git a/test/logdensityfunction.jl b/test/logdensityfunction.jl index 8a5d37ed3..66fe01ed1 100644 --- a/test/logdensityfunction.jl +++ b/test/logdensityfunction.jl @@ -17,6 +17,30 @@ using DifferentiationInterface: DifferentiationInterface using ForwardDiff: ForwardDiff using Mooncake: Mooncake +@model function issue_2844_nested_inner() + p ~ Normal() + return (; p) +end +@model function issue_2844_nested_middle() + inner ~ to_submodel(issue_2844_nested_inner()) + return (; p=inner.p) +end +@model function issue_2844_nested_outer() + middle ~ to_submodel(issue_2844_nested_middle()) + return 0.0 ~ Normal(middle.p) +end +# One submodel boundary deeper than `issue_2844_nested_outer`, so the inference test +# guards against a regression that merely pushes the recursion limit out by one level +# rather than removing it. +@model function issue_2844_nested_middle2() + middle ~ to_submodel(issue_2844_nested_middle()) + return (; p=middle.p) +end +@model function issue_2844_nested_outer_deep() + middle2 ~ to_submodel(issue_2844_nested_middle2()) + return 0.0 ~ Normal(middle2.p) +end + @testset "LogDensityFunction: constructors" begin dist = Beta(2, 2) @model f() = x ~ dist @@ -420,6 +444,21 @@ end end skip = skip end end + + # See https://github.com/TuringLang/DynamicPPL.jl/pull/1427. + @testset "nested to_submodel return values" begin + @testset "$(nameof(model.f))" for model in ( + issue_2844_nested_outer(), issue_2844_nested_outer_deep() + ) + @testset "$tfm_strategy" for tfm_strategy in (UnlinkAll(), LinkAll()) + ldf = DynamicPPL.LogDensityFunction( + model, DynamicPPL.getlogjoint_internal, tfm_strategy + ) + x = rand(ldf) + @test @inferred(LogDensityProblems.logdensity(ldf, x)) isa Float64 + end + end + end end @testset "LogDensityFunction: performance" begin @@ -442,8 +481,18 @@ end y ~ Normal(params.m, params.s) return 1.0 ~ Normal(y) end - @testset for model in - (f(), submodel_inner() | (; s=0.0), submodel_outer(submodel_inner())) + # A submodel nested inside another submodel; see + # https://github.com/TuringLang/DynamicPPL.jl/pull/1427. + @model function submodel_middle(inner) + mid ~ to_submodel(inner) + return (m=mid.m, s=mid.s) + end + @testset for model in ( + f(), + submodel_inner() | (; s=0.0), + submodel_outer(submodel_inner()), + submodel_outer(submodel_middle(submodel_inner())), + ) @testset for tfm_strategy in (UnlinkAll(), LinkAll()) ldf = DynamicPPL.LogDensityFunction(model, getlogjoint_internal, tfm_strategy) x = rand(ldf) diff --git a/test/submodels.jl b/test/submodels.jl index b0fe8184c..d75f4f99d 100644 --- a/test/submodels.jl +++ b/test/submodels.jl @@ -17,6 +17,15 @@ function get_logp_and_rawval_accs(model::Model) return accs end +# Models for the nested-submodel type-stability tests; see +# https://github.com/TuringLang/DynamicPPL.jl/pull/1427. Each level wraps the previous one in +# a `to_submodel`. They must be defined at module scope: a model defined in local (testset) +# scope is not type-inferrable, which would mask the property under test. +@model t2844_inner() = (x ~ Normal(); return (; x)) +@model t2844_middle() = (a ~ to_submodel(t2844_inner()); return (; x=a.x)) +@model t2844_outer() = (b ~ to_submodel(t2844_middle()); return (; x=b.x)) +@model t2844_deeper() = (c ~ to_submodel(t2844_outer()); return (; x=c.x)) + @testset "submodels.jl" begin @testset "$op with AbstractPPL API" for op in [condition, fix] x_val = 1.0 @@ -338,6 +347,22 @@ end @test vnt.data.a.data.b.data.x.data isa Matrix{Float64} @test size(vnt.data.a.data.b.data.x.data) == (2, 2) end + + @testset "type stability of nested submodels (issue #2844)" begin + # See https://github.com/TuringLang/DynamicPPL.jl/pull/1427. + @testset "$(nameof(model.f))" for model in ( + t2844_inner(), t2844_middle(), t2844_outer(), t2844_deeper() + ) + # The fast evaluation path: `init!!` into an `OnlyAccsVarInfo`, under both + # transform strategies. + @testset "$tfm" for tfm in (UnlinkAll(), LinkAll()) + accs = setacc!!(OnlyAccsVarInfo(), LogPriorAccumulator()) + @test @inferred(init!!(model, accs, InitFromPrior(), tfm)) isa Tuple + end + # Evaluating a pre-populated `VarInfo` must also stay type stable. + @test @inferred(DynamicPPL.evaluate_nowarn!!(model, VarInfo(model))) isa Tuple + end + end end end