Fix type-inference failure for nested submodels (#1427)

Evaluating a submodel used to recurse through the shared
`_evaluate!!(::Model, ::AbstractVarInfo)` method. Each submodel level
constructs a contextualized `Model{..., Ctx}` whose `Ctx` parameter
gains another context layer. In nested submodels, Julia's recursion
limiter (`limit_type_size`) sees the same `_evaluate!!(::Model, ...)`
method recurring with a more complex `Model{..., Ctx}` call signature,
widens the `Model` argument to abstract `Model`, and inference can no
longer resolve `model.f` precisely. The model return type then collapses
to `Any`, evaluation falls back to runtime dispatch, and primal/gradient
evaluation slows down.

Submodel evaluation now calls the generated model function directly
instead of going through `_evaluate!!(::Model, ...)`. This removes the
recursive `_evaluate!!(::Model, ...)` call edge while leaving the
recursive `_evaluate!!(::Submodel, ...)` edge intact; the remaining
submodel edge is not widened by the limiter. As a result, inference
stays type stable for nested submodels.

Adds type-stability (`@inferred`) and allocation (`iszero(allocs)`)
regression tests at both the evaluation and `LogDensityFunction` levels,
covering several nesting depths under `UnlinkAll`/`LinkAll`. The tests
were verified to fail on the unfixed code. Also bumps the version to
0.42.1 and adds a HISTORY.md entry.

See TuringLang/Turing.jl#2844

Above is written by Claude/Codex, lightly touched by me.

Author: sunxd3

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.

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"

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!!(

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)

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