Introduce merge and subset for accs

Author: mhauru

src/accumulators.jl

@@ -30,6 +30,13 @@ To be able to work with multi-threading, it should also implement:
 - `split(acc::T)`
 - `combine(acc::T, acc2::T)`
 
+If two accumulators of the same type should be merged in some non-trivial way, other than
+always keeping the second one over the first, `merge(acc1::T, acc2::T)` should be defined.
+
+If limiting the accumulator to a subset of `VarName`s is a meaningful operation and should
+do something other than copy the original accumulator, then
+`subset(acc::T, vns::AbstractVector{<:VarnName})` should be defined.`
+
 See the documentation for each of these functions for more details.
 """
 abstract type AbstractAccumulator end
@@ -113,6 +120,24 @@ used by various AD backends, should implement a method for this function.
 """
 convert_eltype(::Type, acc::AbstractAccumulator) = acc
 
+"""
+    subset(acc::AbstractAccumulator, vns::AbstractVector{<:VarName})
+
+Return a new accumulator that only contains the information for the `VarName`s in `vns`.
+
+By default returns a copy of `acc`. Subtypes should override this behaviour as needed.
+"""
+subset(acc::AbstractAccumulator, ::AbstractVector{<:VarName}) = copy(acc)
+
+"""
+    merge(acc1::AbstractAccumulator, acc2::AbstractAccumulator)
+
+Merge two accumulators of the same type. Returns a new accumulator of the same type.
+
+By default returns a copy of `acc2`. Subtypes should override this behaviour as needed.
+"""
+Base.merge(acc1::AbstractAccumulator, acc2::AbstractAccumulator) = copy(acc2)
+
 """
     AccumulatorTuple{N,T<:NamedTuple}
 
@@ -158,6 +183,52 @@ function Base.convert(::Type{AccumulatorTuple{N,T}}, accs::AccumulatorTuple{N})
     return AccumulatorTuple(convert(T, accs.nt))
 end
 
+"""
+    subset(at::AccumulatorTuple, vns::AbstractVector{<:VarName})
+
+Replace each accumulator `acc` in `at` with `subset(acc, vns)`.
+"""
+function subset(at::AccumulatorTuple, vns::AbstractVector{<:VarName})
+    return AccumulatorTuple(map(Base.Fix2(subset, vns), at.nt))
+end
+
+"""
+    _joint_keys(nt1::NamedTuple, nt2::NamedTuple)
+
+A helper function that returns three tuples of keys given two `NamedTuple`s:
+The keys only in `nt1`, only in `nt2`, and in both, and in that order.
+
+Implemented as a generated function to enabled constant propagation of the result in `merge`.
+"""
+@generated function _joint_keys(
+    nt1::NamedTuple{names1}, nt2::NamedTuple{names2}
+) where {names1,names2}
+    set_names1 = Set(names1)
+    set_names2 = Set(names2)
+    only_in_nt1 = tuple(setdiff(set_names1, set_names2)...)
+    only_in_nt2 = tuple(setdiff(set_names2, set_names1)...)
+    in_both = tuple(intersect(set_names1, set_names2)...)
+    return :($only_in_nt1, $only_in_nt2, $in_both)
+end
+
+"""
+    merge(at1::AccumulatorTuple, at2::AccumulatorTuple)
+
+Merge two `AccumulatorTuple`s.
+
+For any `accumulator_name` that exists in both `at1` and `at2`, we call `merge` on the two
+accumulators themselves. Other accumulators are copied.
+"""
+function Base.merge(at1::AccumulatorTuple, at2::AccumulatorTuple)
+    keys_in_at1, keys_in_at2, keys_in_both = _joint_keys(at1.nt, at2.nt)
+    accs_in_at1 = (getfield(at1.nt, key) for key in keys_in_at1)
+    accs_in_at2 = (getfield(at2.nt, key) for key in keys_in_at2)
+    accs_in_both = (
+        merge(getfield(at1.nt, key), getfield(at2.nt, key)) for key in keys_in_both
+    )
+    return AccumulatorTuple(accs_in_at1..., accs_in_at2..., accs_in_both...)
+end
+
 """
     setacc!!(at::AccumulatorTuple, acc::AbstractAccumulator)
 

src/default_accumulators.jl

@@ -266,3 +266,19 @@ function default_accumulators(
         VariableOrderAccumulator{IntT}(),
     )
 end
+
+function subset(acc::VariableOrderAccumulator, vns::AbstractVector{<:VarName})
+    order = filter(pair -> any(subsumes(vn, first(pair)) for vn in vns), acc.order)
+    return VariableOrderAccumulator(acc.num_produce, order)
+end
+
+"""
+    merge(acc1::VariableOrderAccumulator, acc2::VariableOrderAccumulator)
+
+Merge two `VariableOrderAccumulator` instances.
+
+The `num_produce` field of the return value is the `num_produce` of `acc2`.
+"""
+function Base.merge(acc1::VariableOrderAccumulator, acc2::VariableOrderAccumulator)
+    return VariableOrderAccumulator(acc2.num_produce, merge(acc1.order, acc2.order))
+end

src/simple_varinfo.jl

@@ -417,7 +417,9 @@ Base.eltype(::SimpleOrThreadSafeSimple{<:Any,V}) where {V} = V
 
 # `subset`
 function subset(varinfo::SimpleVarInfo, vns::AbstractVector{<:VarName})
-    return Accessors.@set varinfo.values = _subset(varinfo.values, vns)
+    return SimpleVarInfo(
+        _subset(varinfo.values, vns), subset(getaccs(varinfo), vns), varinfo.transformation
+    )
 end
 
 function _subset(x::AbstractDict, vns::AbstractVector{VN}) where {VN<:VarName}
@@ -454,7 +456,7 @@ _subset(x::VarNamedVector, vns) = subset(x, vns)
 # `merge`
 function Base.merge(varinfo_left::SimpleVarInfo, varinfo_right::SimpleVarInfo)
     values = merge(varinfo_left.values, varinfo_right.values)
-    accs = copy(getaccs(varinfo_right))
+    accs = merge(getaccs(varinfo_left), getaccs(varinfo_right))
     transformation = merge_transformations(
         varinfo_left.transformation, varinfo_right.transformation
     )

src/varinfo.jl

@@ -447,7 +447,7 @@ end
 
 function subset(varinfo::VarInfo, vns::AbstractVector{<:VarName})
     metadata = subset(varinfo.metadata, vns)
-    return VarInfo(metadata, copy(varinfo.accs))
+    return VarInfo(metadata, subset(getaccs(varinfo), vns))
 end
 
 function subset(metadata::NamedTuple, vns::AbstractVector{<:VarName})
@@ -528,7 +528,8 @@ end
 
 function _merge(varinfo_left::VarInfo, varinfo_right::VarInfo)
     metadata = merge_metadata(varinfo_left.metadata, varinfo_right.metadata)
-    return VarInfo(metadata, copy(varinfo_right.accs))
+    accs = merge(getaccs(varinfo_left), getaccs(varinfo_right))
+    return VarInfo(metadata, accs)
 end
 
 function merge_metadata(vnv_left::VarNamedVector, vnv_right::VarNamedVector)

test/accumulators.jl

@@ -108,6 +108,30 @@ using DynamicPPL:
             @test accumulate_observe!!(VariableOrderAccumulator(1), right, left, vn) ==
                 VariableOrderAccumulator(2)
         end
+
+        @testset "merge and subset" begin
+            @test merge(LogPriorAccumulator(1.0), LogPriorAccumulator(2.0)) ==
+                LogPriorAccumulator(3.0)
+            @test merge(LogJacobianAccumulator(1.0), LogJacobianAccumulator(2.0)) ==
+                LogJacobianAccumulator(3.0)
+            @test merge(LogLikelihoodAccumulator(1.0), LogLikelihoodAccumulator(2.0)) ==
+                LogLikelihoodAccumulator(3.0)
+
+            @test merge(
+                VariableOrderAccumulator(1, Dict{VarName,Int}()),
+                VariableOrderAccumulator(2, Dict{VarName,Int}()),
+            ) == VariableOrderAccumulator(2, Dict{VarName,Int}())
+            @test merge(
+                VariableOrderAccumulator(
+                    2, Dict{VarName,Int}((@varname(a) => 1, @varname(b) => 2))
+                ),
+                VariableOrderAccumulator(
+                    1, Dict{VarName,Int}((@varname(a) => 2, @varname(c) => 3))
+                ),
+            ) == VariableOrderAccumulator(
+                1, Dict{VarName,Int}((@varname(a) => 2, @varname(b) => 2, @varname(c) => 3))
+            )
+        end
     end
 
     @testset "accumulator tuples" begin