Fix correctness of submodels that aren't inlined under SMC/PG

HISTORY.md

@@ -1,3 +1,11 @@
+# 0.42.9
+
+Mark extra methods with `Libtask.might_produce` in order to prevent inlining from affecting the correctness of particle methods on models with submodels.
+(Prior to this patch, if your submodel contained observations and was not inlined by the compiler, this could potentially lead to silent wrong results when using SMC or PG.)
+
+Note that if your submodel takes keyword arguments, you must still use `Libtask.@might_produce` yourself, just like if the top-level model has keyword arguments.
+See the patch notes for v0.42.5 for more details.
+
 # 0.42.8
 
 Add support for `TensorBoardLogger.jl` via `AbstractMCMC.mcmc_callback`.

Project.toml

@@ -1,6 +1,6 @@
 name = "Turing"
 uuid = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
-version = "0.42.8"
+version = "0.42.9"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

src/mcmc/particle_mcmc.jl

@@ -537,3 +537,23 @@ Libtask.@might_produce(DynamicPPL.tilde_assume!!)
 Libtask.@might_produce(DynamicPPL.evaluate!!)
 Libtask.@might_produce(DynamicPPL.init!!)
 Libtask.might_produce(::Type{<:Tuple{<:DynamicPPL.Model,Vararg}}) = true
+
+# We need these to automatically catch any submodels whose bodies aren't inlined. See
+# https://github.com/TuringLang/Turing.jl/issues/2772. The line that marks ANY function with
+# the signature f(::Model, ::AbstractVarInfo, ...) might seem completely overkill. However,
+# in practice, the only function that will have this signature is the model evaluation
+# function itself (which currently has the signature model.f(model, varinfo,
+# model.args...)). Of course, there are other functions like logjoint(model, varinfo) that
+# also have this signature, but those are unlikely to be found inside the body of a model,
+# so marking those as might_produce should not cause any real performance issues in particle
+# MCMC.
+# Furthermore, note that these are not sufficient for submodels that have keyword arguments.
+# Those need to be handled by the user in the same way that top-level models with keyword
+# arguments do (i.e. by marking the submodel function itself with @might_produce). We
+# currently don't have a check for this, but it would be good to do so.
+Libtask.@might_produce(DynamicPPL._evaluate!!)
+function Libtask.might_produce(
+    ::Type{<:Tuple{<:Any,<:DynamicPPL.Model,<:DynamicPPL.AbstractVarInfo,Vararg}}
+)
+    return true
+end

test/Aqua.jl

@@ -1,11 +1,29 @@
 module AquaTests
 
 using Aqua: Aqua
+using Libtask: Libtask
 using Turing
 
-# We test ambiguities separately because it catches a lot of problems
-# in dependencies but we test it for Turing.
-Aqua.test_ambiguities([Turing])
+# We test ambiguities specifically only for Turing, because testing ambiguities for all
+# packages in the environment leads to a lot of ambiguities from dependencies that we cannot
+# control.
+#
+# `Libtask.might_produce` is excluded because the `@might_produce` macro generates a lot of
+# ambiguities that will never happen in practice.
+#
+# Specifically, when you write `@might_produce f` for a function `f` that has methods that
+# take keyword arguments, we have to generate a `might_produce` method for
+# `Type{<:Tuple{<:Function,...,typeof(f)}}`. There is no way to circumvent this: see
+# https://github.com/TuringLang/Libtask.jl/issues/197. This in turn will cause method
+# ambiguities with any other function, say `g`, for which
+# `::Type{<:Tuple{typeof(g),Vararg}}` is marked as produceable.
+#
+# To avoid the method ambiguities, we *could* manually spell out `might_produce` methods for
+# each method of `g` manually instead of using Vararg, but that would be both very verbose
+# and fragile. It would also not provide any real benefit since those ambiguities are not
+# meaningful in practice (in particular, to trigger this we would need to call `g(..., f)`,
+# which is incredibly unlikely).
+Aqua.test_ambiguities([Turing]; exclude=[Libtask.might_produce])
 Aqua.test_all(Turing; ambiguities=false)
 
 end

test/Project.toml

@@ -17,6 +17,7 @@ DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8"
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 HypothesisTests = "09f84164-cd44-5f33-b23f-e6b0d136a0d5"
+Libtask = "6f1fad26-d15e-5dc8-ae53-837a1d7b8c9f"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
 LogDensityProblemsAD = "996a588d-648d-4e1f-a8f0-a84b347e47b1"
@@ -57,6 +58,7 @@ DynamicPPL = "0.39.6"
 FiniteDifferences = "0.10.8, 0.11, 0.12"
 ForwardDiff = "0.10.12 - 0.10.32, 0.10, 1"
 HypothesisTests = "0.11"
+Libtask = "0.9"
 LinearAlgebra = "1"
 LogDensityProblems = "2"
 LogDensityProblemsAD = "1.4"

test/mcmc/particle_mcmc.jl

@@ -181,6 +181,21 @@ end
         @test mean(chain2[:x]) ≈ 7.5 atol = 0.2
     end
 
+    @testset "correctness with submodels that aren't inlined" begin
+        # cf. https://github.com/TuringLang/Turing.jl/issues/2772
+        @model function inner(y, x)
+            @noinline
+            return y ~ Normal(x)
+        end
+        @model function nested(y)
+            x ~ Normal()
+            return a ~ to_submodel(inner(y, x))
+        end
+        m1 = nested(1.0)
+        chn = sample(StableRNG(468), m1, PG(10), 500)
+        @test mean(chn[:x]) ≈ 0.5 atol = 0.05
+    end
+
     @testset "refuses to run threadsafe eval" begin
         # PG can't run models that have nondeterministic evaluation order,
         # so it should refuse to run models marked as threadsafe.