Add a test, plus a bunch of == methods

Author: penelopeysm

src/debug_utils.jl

@@ -156,14 +156,22 @@ end
 const _DEBUG_ACC_NAME = :Debug
 DynamicPPL.accumulator_name(::Type{<:DebugAccumulator}) = _DEBUG_ACC_NAME
 
+function Base.:(==)(acc1::DebugAccumulator, acc2::DebugAccumulator)
+    return (
+        acc1.varnames_seen == acc2.varnames_seen &&
+        acc1.statements == acc2.statements &&
+        acc1.error_on_failure == acc2.error_on_failure
+    )
+end
+
 function _zero(acc::DebugAccumulator)
     return DebugAccumulator(
         OrderedDict{VarName,Int}(), Vector{Stmt}(), acc.error_on_failure
     )
 end
-reset(acc::DebugAccumulator) = _zero(acc)
-split(acc::DebugAccumulator) = _zero(acc)
-function combine(acc1::DebugAccumulator, acc2::DebugAccumulator)
+DynamicPPL.reset(acc::DebugAccumulator) = _zero(acc)
+DynamicPPL.split(acc::DebugAccumulator) = _zero(acc)
+function DynamicPPL.combine(acc1::DebugAccumulator, acc2::DebugAccumulator)
     return DebugAccumulator(
         merge(acc1.varnames_seen, acc2.varnames_seen),
         vcat(acc1.statements, acc2.statements),

src/extract_priors.jl

@@ -10,6 +10,10 @@ end
 
 accumulator_name(::PriorDistributionAccumulator) = :PriorDistributionAccumulator
 
+function Base.:(==)(acc1::PriorDistributionAccumulator, acc2::PriorDistributionAccumulator)
+    return acc1.priors == acc2.priors
+end
+
 _zero(acc::PriorDistributionAccumulator) = PriorDistributionAccumulator(empty(acc.priors))
 reset(acc::PriorDistributionAccumulator) = _zero(acc)
 split(acc::PriorDistributionAccumulator) = _zero(acc)

src/pointwise_logdensities.jl

@@ -32,6 +32,12 @@ function PointwiseLogProbAccumulator{whichlogprob,KeyType}() where {whichlogprob
     return PointwiseLogProbAccumulator{whichlogprob,KeyType,typeof(logps)}(logps)
 end
 
+function Base.:(==)(
+    acc1::PointwiseLogProbAccumulator{wlp1}, acc2::PointwiseLogProbAccumulator{wlp2}
+) where {wlp1,wlp2}
+    return (wlp1 == wlp2 && acc1.logps == acc2.logps)
+end
+
 function Base.copy(acc::PointwiseLogProbAccumulator{whichlogprob}) where {whichlogprob}
     return PointwiseLogProbAccumulator{whichlogprob}(copy(acc.logps))
 end

src/values_as_in_model.jl

@@ -20,6 +20,10 @@ function ValuesAsInModelAccumulator(include_colon_eq)
     return ValuesAsInModelAccumulator(OrderedDict{VarName,Any}(), include_colon_eq)
 end
 
+function Base.:(==)(acc1::ValuesAsInModelAccumulator, acc2::ValuesAsInModelAccumulator)
+    return (acc1.include_colon_eq == acc2.include_colon_eq && acc1.values == acc2.values)
+end
+
 function Base.copy(acc::ValuesAsInModelAccumulator)
     return ValuesAsInModelAccumulator(copy(acc.values), acc.include_colon_eq)
 end

test/accumulators.jl

@@ -13,6 +13,7 @@ using DynamicPPL:
     convert_eltype,
     getacc,
     map_accumulator,
+    reset,
     setacc!!,
     split
 

test/simple_varinfo.jl

@@ -204,7 +204,7 @@
             end
 
             # Reset the logp accumulators.
-            svi_eval = DynamicPPL.resetlogp!!(svi_eval)
+            svi_eval = DynamicPPL.resetaccs!!(svi_eval)
 
             # Compute `logjoint` using the varinfo.
             logπ = logjoint(model, svi_eval)

test/threadsafe.jl

@@ -25,7 +25,7 @@
         @test getlogjoint(vi) == lp
         @test getlogjoint(threadsafe_vi) == lp + 42
 
-        threadsafe_vi = resetlogp!!(threadsafe_vi)
+        threadsafe_vi = DynamicPPL.resetaccs!!(threadsafe_vi)
         @test iszero(getlogjoint(threadsafe_vi))
         expected_accs = DynamicPPL.AccumulatorTuple(
             (DynamicPPL.split(acc) for acc in DynamicPPL.getaccs(threadsafe_vi.varinfo))...

test/varinfo.jl

@@ -112,7 +112,7 @@ end
         test_base(SimpleVarInfo(DynamicPPL.VarNamedVector()))
     end
 
-    @testset "get/set/acc/resetlogp" begin
+    @testset "get/set/acclogp" begin
         function test_varinfo_logp!(vi)
             @test DynamicPPL.getlogjoint(vi) === 0.0
             vi = DynamicPPL.setlogprior!!(vi, 1.0)
@@ -131,8 +131,6 @@ end
             @test DynamicPPL.getlogprior(vi) === 2.0
             @test DynamicPPL.getloglikelihood(vi) === 2.0
             @test DynamicPPL.getlogjoint(vi) === 4.0
-            vi = DynamicPPL.resetlogp!!(vi)
-            @test DynamicPPL.getlogjoint(vi) === 0.0
         end
 
         vi = VarInfo()
@@ -143,7 +141,7 @@ end
         test_varinfo_logp!(SimpleVarInfo(DynamicPPL.VarNamedVector()))
     end
 
-    @testset "accumulators" begin
+    @testset "logp accumulators" begin
         @model function demo()
             a ~ Normal()
             b ~ Normal()
@@ -227,6 +225,71 @@ end
         @test_throws r"has no field `?LogPrior" getlogjoint(vi)
     end
 
+    @testset "resetaccs" begin
+        # Put in a bunch of accumulators, check that they're all reset either
+        # when we call resetaccs!!, empty!!, or evaluate!!.
+        @model function demo()
+            a ~ Normal()
+            return x ~ Normal(a)
+        end
+        model = demo()
+        vi_orig = VarInfo(model)
+        # It already has the logp accumulators, so let's add in some more.
+        vi_orig = DynamicPPL.setacc!!(vi_orig, DynamicPPL.DebugUtils.DebugAccumulator(true))
+        vi_orig = DynamicPPL.setacc!!(vi_orig, DynamicPPL.ValuesAsInModelAccumulator(true))
+        vi_orig = DynamicPPL.setacc!!(vi_orig, DynamicPPL.PriorDistributionAccumulator())
+        vi_orig = DynamicPPL.setacc!!(
+            vi_orig, DynamicPPL.PointwiseLogProbAccumulator{:both}()
+        )
+        # And evaluate the model once so that they are populated.
+        _, vi_orig = DynamicPPL.evaluate!!(model, vi_orig)
+
+        function all_accs_empty(vi::AbstractVarInfo)
+            for acc_key in keys(DynamicPPL.getaccs(vi))
+                acc = DynamicPPL.getacc(vi, Val(acc_key))
+                acc == DynamicPPL.reset(acc) || return false
+            end
+            return true
+        end
+
+        @test !all_accs_empty(vi_orig)
+
+        vi = DynamicPPL.resetaccs!!(deepcopy(vi_orig))
+        @test all_accs_empty(vi)
+        @test getlogjoint(vi) == 0.0 # for good measure
+        @test getlogprior(vi) == 0.0
+        @test getloglikelihood(vi) == 0.0
+
+        vi = DynamicPPL.empty!!(deepcopy(vi_orig))
+        @test all_accs_empty(vi)
+        @test getlogjoint(vi) == 0.0
+        @test getlogprior(vi) == 0.0
+        @test getloglikelihood(vi) == 0.0
+
+        function all_accs_same(vi1::AbstractVarInfo, vi2::AbstractVarInfo)
+            # Check that they have the same accs
+            keys1 = Set(keys(DynamicPPL.getaccs(vi1)))
+            keys2 = Set(keys(DynamicPPL.getaccs(vi2)))
+            keys1 == keys2 || return false
+            # Check that they have the same values
+            for acc_key in keys1
+                acc1 = DynamicPPL.getacc(vi1, Val(acc_key))
+                acc2 = DynamicPPL.getacc(vi2, Val(acc_key))
+                if acc1 != acc2
+                    @show acc1, acc2
+                end
+                acc1 == acc2 || return false
+            end
+            return true
+        end
+        # Hopefully this doesn't matter
+        @test all_accs_same(vi_orig, deepcopy(vi_orig))
+        # If we re-evaluate, then we expect the accs to be reset prior to evaluation.
+        # Thus after re-evaluation, the accs should be exactly the same as before.
+        _, vi = DynamicPPL.evaluate!!(model, deepcopy(vi_orig))
+        @test all_accs_same(vi, vi_orig)
+    end
+
     @testset "flags" begin
         # Test flag setting:
         #    is_flagged, set_flag!, unset_flag!