Make pobserve work with VarInfos that have no likelihood accumulator

Author: penelopeysm

src/pobserve_macro.jl

@@ -25,10 +25,11 @@ function _pobserve(expr::Expr)
         end
         retvals_and_likelihoods = fetch.(likelihood_tasks)
         total_likelihoods = sum(last, retvals_and_likelihoods)
-        # println("Total likelihoods: ", total_likelihoods)
-        $(esc(:(__varinfo__))) = $(DynamicPPL.accloglikelihood!!)(
-            $(esc(:(__varinfo__))), total_likelihoods
-        )
+        if $(DynamicPPL.hasacc)($(esc(:(__varinfo__))), Val(:LogLikelihood))
+            $(esc(:(__varinfo__))) = $(DynamicPPL.accloglikelihood!!)(
+                $(esc(:(__varinfo__))), total_likelihoods
+            )
+        end
         map(first, retvals_and_likelihoods)
     end
     return return_expr
@@ -49,8 +50,13 @@ function process_tilde_statements(expr::Expr)
         end
     ) || error("expected block")
     @gensym loglike
-    beginning_statement =
-        :($loglike = zero($(DynamicPPL.getloglikelihood)($(esc(:(__varinfo__))))))
+    beginning_expr = quote
+        $loglike = if $(DynamicPPL.hasacc)($(esc(:(__varinfo__))), Val(:LogLikelihood))
+            zero($(DynamicPPL.getloglikelihood)($(esc(:(__varinfo__)))))
+        else
+            zero($(DynamicPPL.LogProbType))
+        end
+    end
     n_statements = length(statements)
     transformed_statements::Vector{Vector{Expr}} = map(enumerate(statements)) do (i, stmt)
         is_last = i == n_statements
@@ -79,6 +85,6 @@ function process_tilde_statements(expr::Expr)
             e
         end
     end
-    new_statements = [beginning_statement, reduce(vcat, transformed_statements)...]
+    new_statements = [beginning_expr.args..., reduce(vcat, transformed_statements)...]
     return Expr(:block, new_statements...)
 end

test/pobserve_macro.jl

@@ -15,6 +15,18 @@ using DynamicPPL, Distributions, Test
         @test isapprox(DynamicPPL.getloglikelihood(vi), expected_loglike)
     end
 
+    @testset "doesn't error when varinfo has no likelihood acc" begin
+        @model function f(x)
+            @pobserve for i in eachindex(x)
+                x[i] ~ Normal()
+            end
+        end
+        x = randn(3)
+        vi = VarInfo()
+        vi = DynamicPPL.setaccs!!(vi, (DynamicPPL.LogPriorAccumulator(),))
+        @test DynamicPPL.evaluate!!(f(x), vi) isa Any
+    end
+
     @testset "return values are correct" begin
         @testset "single expression at the end" begin
             @model function f(x)