made filtering for errors only in the tilde pipeline optional

Author: torfjelde

ext/DynamicPPLJETExt.jl

@@ -3,10 +3,37 @@ module DynamicPPLJETExt
 using DynamicPPL: DynamicPPL
 using JET: JET
 
+"""
+    is_tilde_instance(x)
+
+Return `true` if `x` is a method instance of a tilde function, otherwise `false`.
+"""
+is_tilde_instance(x) = false
+is_tilde_instance(frame::JET.VirtualFrame) = is_tilde_instance(frame.linfo)
+is_tilde_instance(mi::Core.MethodInstance) = is_tilde_instance(mi.specTypes.parameters[1])
+is_tilde_instance(::Type{typeof(DynamicPPL.tilde_assume!!)}) = true
+is_tilde_instance(::Type{typeof(DynamicPPL.tilde_observe!!)}) = true
+is_tilde_instance(::Type{typeof(DynamicPPL.dot_tilde_assume!!)}) = true
+is_tilde_instance(::Type{typeof(DynamicPPL.dot_tilde_observe!!)}) = true
+
+"""
+    report_has_error_in_tilde(report)
+
+Return `true` if the given error `report` contains a tilde function in its frames, otherwise `false`.
+
+This is used to filter out reports that occur outside of the tilde pipeline, in an attempt to avoid
+warning the user about DynamicPPL doing something wrong when it is in fact an issue with the user's code.
+"""
+function report_has_error_in_tilde(report)
+    frames = report.vst
+    return any(is_tilde_instance, frames)
+end
+
 function DynamicPPL.determine_varinfo(
     model::DynamicPPL.Model,
     context::DynamicPPL.AbstractContext=DynamicPPL.DefaultContext();
     verbose::Bool=false,
+    only_tilde::Bool=true
 )
     # First we try with the typed varinfo.
     varinfo = DynamicPPL.typed_varinfo(model)
@@ -18,6 +45,9 @@ function DynamicPPL.determine_varinfo(
     )
     result_eval = JET.report_call(f_eval, argtypes_eval)
     reports_eval = JET.get_reports(result_eval)
+    if only_tilde
+        reports_eval = filter(report_has_error_in_tilde, reports_eval)
+    end
     # If we get reports => we had issues so we use the untyped varinfo.
     issuccess &= length(reports_eval) == 0
     if issuccess
@@ -27,6 +57,9 @@ function DynamicPPL.determine_varinfo(
         )
         result_sample = JET.report_call(f_sample, argtypes_sample)
         reports_sample = JET.get_reports(result_sample)
+        if only_tilde
+            reports_sample = filter(report_has_error_in_tilde, reports_sample)
+        end
         # If we get reports => we had issues so we use the untyped varinfo.
         issuccess &= length(reports_sample) == 0
         if !issuccess && verbose
@@ -55,3 +88,4 @@ function DynamicPPL.determine_varinfo(
 end
 
 end
+ 

test/ext/DynamicPPLJETExt.jl

@@ -1,5 +1,3 @@
-using JET: JET
-
 @testset "DynamicPPLJETExt.jl" begin
     @testset "determine_varinfo" begin
         @model function demo1()
@@ -11,7 +9,7 @@ using JET: JET
             end
         end
         model = demo1()
-        @test DynamicPPL.determine_varinfo(model) isa DynamicPPL.UntypedVarInfo
+        @test DynamicPPL.determine_varinfo(model; verbose=true) isa DynamicPPL.UntypedVarInfo
 
         @model demo2() = x ~ Normal()
         @test DynamicPPL.determine_varinfo(demo2()) isa DynamicPPL.TypedVarInfo
@@ -26,7 +24,7 @@ using JET: JET
                 z ~ Normal()
             end
         end
-        @test DynamicPPL.determine_varinfo(demo3()) isa DynamicPPL.UntypedVarInfo
+        @test DynamicPPL.determine_varinfo(demo3(); verbose=true) isa DynamicPPL.UntypedVarInfo
 
         # Evaluation works (and it would even do so in practice), but sampling
         # fill fail due to storing `Cauchy{Float64}` in `Vector{Normal{Float64}}`.
@@ -38,24 +36,38 @@ using JET: JET
                 y ~ Cauchy() # different distibution, but same transformation => should work
             end
         end
-        @test DynamicPPL.determine_varinfo(demo4()) isa DynamicPPL.UntypedVarInfo
+        @test DynamicPPL.determine_varinfo(demo4(); verbose=true) isa DynamicPPL.UntypedVarInfo
+
+        # In this model, the type error occurs in the user code rather than in DynamicPPL.
+        @model function demo5()
+            x ~ Normal()
+            xs = Any[]
+            push!(xs, x)
+            # `sum(::Vector{Any})` can potentially error unless the dynamic manages to resolve the
+            # correct `zero` method. As a result, this code will run, but JET will raise this is an issue.
+            return sum(xs)
+        end
+        # Should pass if we're only checking the tilde statements.
+        @test DynamicPPL.determine_varinfo(demo5(); verbose=true) isa DynamicPPL.TypedVarInfo
+        # Should fail if we're including errors in the model body.
+        @test DynamicPPL.determine_varinfo(demo5(); verbose=true, only_tilde=false) isa DynamicPPL.UntypedVarInfo
     end
 
-    # @testset "demo models" begin
-    #     @testset "$(model.f)" for model in DynamicPPL.TestUtils.DEMO_MODELS
-    #         varinfo = DynamicPPL.DynamicPPL.determine_varinfo(model)
-    #         # They should all result in typed.
-    #         @test varinfo isa DynamicPPL.TypedVarInfo
-    #         # But let's also make sure that they're not lying.
-    #         f_eval, argtypes_eval = DynamicPPL.DebugUtils.gen_evaluator_call_with_types(
-    #             model, varinfo
-    #         )
-    #         JET.test_call(f_eval, argtypes_eval)
+    @testset "demo models" begin
+        @testset "$(model.f)" for model in DynamicPPL.TestUtils.DEMO_MODELS
+            varinfo = DynamicPPL.DynamicPPL.determine_varinfo(model)
+            # They should all result in typed.
+            @test varinfo isa DynamicPPL.TypedVarInfo
+            # But let's also make sure that they're not lying.
+            f_eval, argtypes_eval = DynamicPPL.DebugUtils.gen_evaluator_call_with_types(
+                model, varinfo
+            )
+            JET.test_call(f_eval, argtypes_eval)
 
-    #         f_sample, argtypes_sample = DynamicPPL.DebugUtils.gen_evaluator_call_with_types(
-    #             model, varinfo, DynamicPPL.SamplingContext()
-    #         )
-    #         JET.test_call(f_sample, argtypes_sample)
-    #     end
-    # end
+            f_sample, argtypes_sample = DynamicPPL.DebugUtils.gen_evaluator_call_with_types(
+                model, varinfo, DynamicPPL.SamplingContext()
+            )
+            JET.test_call(f_sample, argtypes_sample)
+        end
+    end
 end

test/runtests.jl

@@ -25,6 +25,8 @@ using Test
 using Distributions
 using LinearAlgebra # Diagonal
 
+using JET: JET
+
 using Combinatorics: combinations
 
 using DynamicPPL: getargs_dottilde, getargs_tilde, Selector
@@ -71,6 +73,7 @@ include("test_util.jl")
 
         @testset "extensions" begin
             include("ext/DynamicPPLMCMCChainsExt.jl")
+            include("ext/DynamicPPLJETExt.jl")
         end
 
         @testset "ad" begin