Fix tests

Author: penelopeysm

src/debug_utils.jl

@@ -426,7 +426,7 @@ function check_model_and_trace(
     issuccess = check_model_pre_evaluation(model)
 
     # Force single-threaded execution.
-    DynamicPPL.evaluate_threadunsafe!!(model, varinfo)
+    _, varinfo = DynamicPPL.evaluate_threadunsafe!!(model, varinfo)
 
     # Perform checks after evaluating the model.
     debug_acc = DynamicPPL.getacc(varinfo, Val(_DEBUG_ACC_NAME))

src/pointwise_logdensities.jl

@@ -230,23 +230,36 @@ function pointwise_logdensities(
     # Get the data by executing the model once
     vi = VarInfo(model)
 
+    # This accumulator tracks the pointwise log-probabilities in a single iteration.
     AccType = PointwiseLogProbAccumulator{whichlogprob,KeyType}
     vi = setaccs!!(vi, (AccType(),))
 
     iters = Iterators.product(1:size(chain, 1), 1:size(chain, 3))
+
+    # Maintain a separate accumulator that isn't tied to a VarInfo but rather
+    # tracks _all_ iterations.
+    all_logps = AccType()
     for (sample_idx, chain_idx) in iters
         # Update the values
         setval!(vi, chain, sample_idx, chain_idx)
 
         # Execute model
-        vi = last(evaluate!!(model, vi))
+        vi = setaccs!!(vi, (AccType(),))
+        vi = last(_evaluate!!(model, vi))
+
+        # Get the log-probabilities
+        this_iter_logps = getacc(vi, Val(accumulator_name(AccType))).logps
+
+        # Merge into main acc
+        for (varname, this_lp) in this_iter_logps
+            push!(all_logps, varname, only(this_lp))
+        end
     end
 
-    logps = getacc(vi, Val(accumulator_name(AccType))).logps
     niters = size(chain, 1)
     nchains = size(chain, 3)
     logdensities = OrderedDict(
-        varname => reshape(vals, niters, nchains) for (varname, vals) in logps
+        varname => reshape(vals, niters, nchains) for (varname, vals) in all_logps.logps
     )
     return logdensities
 end

test/accumulators.jl

@@ -24,12 +24,12 @@ using DynamicPPL:
                 LogPriorAccumulator() ==
                 LogPriorAccumulator{Float64}() ==
                 LogPriorAccumulator{Float64}(0.0) ==
-                zero(LogPriorAccumulator(1.0))
+                DynamicPPL.reset(LogPriorAccumulator(1.0))
             @test LogLikelihoodAccumulator(0.0) ==
                 LogLikelihoodAccumulator() ==
                 LogLikelihoodAccumulator{Float64}() ==
                 LogLikelihoodAccumulator{Float64}(0.0) ==
-                zero(LogLikelihoodAccumulator(1.0))
+                DynamicPPL.reset(LogLikelihoodAccumulator(1.0))
         end
 
         @testset "addition and incrementation" begin
@@ -136,7 +136,7 @@ using DynamicPPL:
         @testset "map_accumulator(s)!!" begin
             # map over all accumulators
             accs = AccumulatorTuple(lp_f32, ll_f32)
-            @test map(zero, accs) == AccumulatorTuple(
+            @test map(DynamicPPL.reset, accs) == AccumulatorTuple(
                 LogPriorAccumulator(0.0f0), LogLikelihoodAccumulator(0.0f0)
             )
             # Test that the original wasn't modified.
@@ -147,7 +147,7 @@ using DynamicPPL:
                 AccumulatorTuple(LogPriorAccumulator(1.0), LogLikelihoodAccumulator(1.0))
 
             # only apply to a particular accumulator
-            @test map_accumulator(zero, accs, Val(:LogLikelihood)) ==
+            @test map_accumulator(DynamicPPL.reset, accs, Val(:LogLikelihood)) ==
                 AccumulatorTuple(lp_f32, LogLikelihoodAccumulator(0.0f0))
             @test map_accumulator(
                 acc -> convert_eltype(Float64, acc), accs, Val(:LogLikelihood)