Fix variable order and name of map_accumulator!!

Author: mhauru

src/abstract_varinfo.jl

@@ -239,7 +239,7 @@ end
 Update all the accumulators of `vi` by calling `accumulate_assume!!` on them.
 """
 function accumulate_assume!!(vi::AbstractVarInfo, val, logjac, vn, right)
-    return map_accumulator!!(vi, accumulate_assume!!, val, logjac, vn, right)
+    return map_accumulators!!(acc -> accumulate_assume!!(acc, val, logjac, vn, right), vi)
 end
 
 """
@@ -248,37 +248,35 @@ end
 Update all the accumulators of `vi` by calling `accumulate_observe!!` on them.
 """
 function accumulate_observe!!(vi::AbstractVarInfo, right, left, vn)
-    return map_accumulator!!(vi, accumulate_observe!!, right, left, vn)
+    return map_accumulators!!(acc -> accumulate_observe!!(acc, right, left, vn), vi)
 end
 
 """
-    map_accumulator!!(vi::AbstractVarInfo, func::Function, args...) where {accname}
+    map_accumulators(vi::AbstractVarInfo, func::Function)
 
-Update all accumulators of `vi` by calling `func(acc, args...)` on them and replacing
-them with the return values.
+Update all accumulators of `vi` by calling `func` on them and replacing them with the return
+values.
 """
-function map_accumulator!!(vi::AbstractVarInfo, func::Function, args...)
-    return setaccs!!(vi, map_accumulator!!(getaccs(vi), func, args...))
+function map_accumulators!!(func::Function, vi::AbstractVarInfo)
+    return setaccs!!(vi, map(func, getaccs(vi)))
 end
 
 """
-    map_accumulator!!(vi::AbstractVarInfo, ::Val{accname}, func::Function, args...) where {accname}
+    map_accumulator!!(func::Function, vi::AbstractVarInfo, ::Val{accname}) where {accname}
 
-Update the accumulator `accname` of `vi` by calling `func(acc, args...)` on and replacing
-it with the return value.
+Update the accumulator `accname` of `vi` by calling `func` on it and replacing it with the
+return value.
 """
-function map_accumulator!!(vi::AbstractVarInfo, accname::Val, func::Function, args...)
-    return setaccs!!(vi, map_accumulator!!(getaccs(vi), accname, func, args...))
+function map_accumulator!!(func::Function, vi::AbstractVarInfo, accname::Val)
+    return setaccs!!(vi, map_accumulator(func, getaccs(vi), accname))
 end
 
-function map_accumulator!!(vi::AbstractVarInfo, accname::Symbol, func::Function, args...)
+function map_accumulator!!(func::Function, vi::AbstractVarInfo, accname::Symbol)
     return error(
         """
-        The method
-        map_accumulator!!(vi::AbstractVarInfo, accname::Symbol, func::Function, args...)
+        The method map_accumulator!!(func::Function, vi::AbstractVarInfo, accname::Symbol)
         does not exist. For type stability reasons use
-        map_accumulator!!(vi::AbstractVarInfo, accname::Val, func::Function, args...)
-        instead.
+        map_accumulator!!(func::Function, vi::AbstractVarInfo, ::Val{accname}) instead.
         """
     )
 end
@@ -291,7 +289,7 @@ Add `logp` to the value of the log of the prior probability in `vi`.
 See also: [`accloglikelihood!!`](@ref), [`acclogp!!`](@ref), [`getlogprior`](@ref), [`setlogprior!!`](@ref).
 """
 function acclogprior!!(vi::AbstractVarInfo, logp)
-    return map_accumulator!!(vi, Val(:LogPrior), +, LogPrior(logp))
+    return map_accumulator!!(acc -> acc + LogPrior(logp), vi, Val(:LogPrior))
 end
 
 """
@@ -302,7 +300,7 @@ Add `logp` to the value of the log of the likelihood in `vi`.
 See also: [`accloglikelihood!!`](@ref), [`acclogp!!`](@ref), [`getloglikelihood`](@ref), [`setloglikelihood!!`](@ref).
 """
 function accloglikelihood!!(vi::AbstractVarInfo, logp)
-    return map_accumulator!!(vi, Val(:LogLikelihood), +, LogLikelihood(logp))
+    return map_accumulator!!(acc -> acc + LogLikelihood(logp), vi, Val(:LogLikelihood))
 end
 
 """
@@ -326,10 +324,10 @@ Reset the values of the log probabilities (prior and likelihood) in `vi`
 """
 function resetlogp!!(vi::AbstractVarInfo)
     if hasacc(vi, Val(:LogPrior))
-        vi = map_accumulator!!(vi, Val(:LogPrior), zero)
+        vi = map_accumulator!!(zero, vi, Val(:LogPrior))
     end
     if hasacc(vi, Val(:LogLikelihood))
-        vi = map_accumulator!!(vi, Val(:LogLikelihood), zero)
+        vi = map_accumulator!!(zero, vi, Val(:LogLikelihood))
     end
     return vi
 end

src/accumulators.jl

@@ -89,15 +89,15 @@ function combine end
 # TODO(mhauru) The existence of this function makes me sad. See comment in unflatten in
 # src/varinfo.jl.
 """
-    convert_eltype(acc::AbstractAccumulator, ::Type{T})
+    convert_eltype(::Type{T}, acc::AbstractAccumulator)
 
 Convert `acc` to use element type `T`.
 
 What "element type" means depends on the type of `acc`. By default this function does
 nothing. Accumulator types that need to hold differentiable values, such as dual numbers
 used by various AD backends, should implement a method for this function.
 """
-convert_eltype(acc::AbstractAccumulator, ::Type) = acc
+convert_eltype(::Type, acc::AbstractAccumulator) = acc
 
 # END ABSTRACT ACCUMULATOR, BEGIN ACCUMULATOR TUPLE
 
@@ -167,36 +167,25 @@ function getacc(at::AccumulatorTuple, ::Val{accname}) where {accname}
     return at[accname]
 end
 
-"""
-    map_accumulator!!(at::AccumulatorTuple, func::Function, args...)
-
-Update the accumulators in `at` by calling `func(acc, args...)` on them and replacing them
-with the return values.
-
-Returns a new `AccumulatorTuple`. The `!!` in the name is for consistency with the
-corresponding function for `AbstractVarInfo`.
-"""
-function map_accumulator!!(at::AccumulatorTuple, func::Function, args...)
-    return AccumulatorTuple(map(acc -> func(acc, args...), at.nt))
+function Base.map(func::Function, at::AccumulatorTuple)
+    return AccumulatorTuple(map(func, at.nt))
 end
 
 """
-    map_accumulator!!(at::AccumulatorTuple, ::Val{accname}, func::Function, args...)
+    map_accumulator(func::Function, at::AccumulatorTuple, ::Val{accname})
 
-Update the accumulator with name `accname` in `at` by calling `func(acc, args...)` on it
-and replacing it with the return value.
+Update the accumulator with name `accname` in `at` by calling `func` on it.
 
-Returns a new `AccumulatorTuple`. The `!!` in the name is for consistency with the
-corresponding function for `AbstractVarInfo`.
+Returns a new `AccumulatorTuple`.
 """
-function map_accumulator!!(
-    at::AccumulatorTuple, ::Val{accname}, func::Function, args...
+function map_accumulator(
+    func::Function, at::AccumulatorTuple, ::Val{accname}
 ) where {accname}
     # Would like to write this as
     # return Accessors.@set at.nt[accname] = func(at[accname], args...)
     # for readability, but that one isn't type stable due to
     # https://github.com/JuliaObjects/Accessors.jl/issues/198
-    new_val = func(at[accname], args...)
+    new_val = func(at[accname])
     new_nt = merge(at.nt, NamedTuple{(accname,)}((new_val,)))
     return AccumulatorTuple(new_nt)
 end
@@ -318,7 +307,7 @@ end
 # convert_eltype(::AbstractAccumulator, ::Type). This is because they are only used to
 # deal with dual number types of AD backends, which shouldn't concern NumProduce. This is
 # horribly hacky and should be fixed. See also comment in `unflatten` in `src/varinfo.jl`.
-convert_eltype(acc::LogPrior, ::Type{T}) where {T} = LogPrior(convert(T, acc.logp))
-function convert_eltype(acc::LogLikelihood, ::Type{T}) where {T}
+convert_eltype(::Type{T}, acc::LogPrior) where {T} = LogPrior(convert(T, acc.logp))
+function convert_eltype(::Type{T}, acc::LogLikelihood) where {T}
     return LogLikelihood(convert(T, acc.logp))
 end

src/simple_varinfo.jl

@@ -251,7 +251,7 @@ function unflatten(svi::SimpleVarInfo, x::AbstractVector)
     # TODO(mhauru) See comment in unflatten in src/varinfo.jl for why this conversion is
     # required but undesireable.
     et = float_type_with_fallback(eltype(x))
-    accs = map_accumulator!!(svi.accs, convert_eltype, et)
+    accs = map(acc -> convert_eltype(et, acc), svi.accs)
     return SimpleVarInfo(vals, accs, svi.transformation)
 end
 

src/threadsafe.jl

@@ -46,19 +46,17 @@ function getaccs(vi::ThreadSafeVarInfo)
     return AccumulatorTuple(map(anv -> getacc(vi, anv), accname_vals))
 end
 
-# Calls to map_accumulator!! are thread-specific by default. For any use of them that should
-# _not_ be thread-specific a specific method has to be written.
-function map_accumulator!!(vi::ThreadSafeVarInfo, accname::Val, func::Function, args...)
+# Calls to map_accumulator(s)!! are thread-specific by default. For any use of them that
+# should _not_ be thread-specific a specific method has to be written.
+function map_accumulator!!(func::Function, vi::ThreadSafeVarInfo, accname::Val)
     tid = Threads.threadid()
-    vi.accs_by_thread[tid] = map_accumulator!!(
-        vi.accs_by_thread[tid], accname, func, args...
-    )
+    vi.accs_by_thread[tid] = map_accumulator(func, vi.accs_by_thread[tid], accname)
     return vi
 end
 
-function map_accumulator!!(vi::ThreadSafeVarInfo, func::Function, args...)
+function map_accumulators!!(func::Function, vi::ThreadSafeVarInfo)
     tid = Threads.threadid()
-    vi.accs_by_thread[tid] = map_accumulator!!(vi.accs_by_thread[tid], func, args...)
+    vi.accs_by_thread[tid] = map(func, vi.accs_by_thread[tid])
     return vi
 end
 
@@ -186,9 +184,9 @@ end
 function resetlogp!!(vi::ThreadSafeVarInfo)
     vi = Accessors.@set vi.varinfo = resetlogp!!(vi.varinfo)
     for i in eachindex(vi.accs_by_thread)
-        vi.accs_by_thread[i] = map_accumulator!!(vi.accs_by_thread[i], Val(:LogPrior), zero)
-        vi.accs_by_thread[i] = map_accumulator!!(
-            vi.accs_by_thread[i], Val(:LogLikelihood), zero
+        vi.accs_by_thread[i] = map_accumulator(zero, vi.accs_by_thread[i], Val(:LogPrior))
+        vi.accs_by_thread[i] = map_accumulator(
+            zero, vi.accs_by_thread[i], Val(:LogLikelihood)
         )
     end
     return vi

src/transforming.jl

@@ -45,7 +45,7 @@ function invlink!!(::DynamicTransformation, vi::AbstractVarInfo, model::Model)
     return _transform!!(NoTransformation(), DynamicTransformationContext{true}(), vi, model)
 end
 
-function _transform(
+function _transform!!(
     t::AbstractTransformation,
     ctx::DynamicTransformationContext,
     vi::AbstractVarInfo,
@@ -54,8 +54,8 @@ function _transform(
     # To transform using DynamicTransformationContext, we evaluate the model, but we do not
     # need to use any accumulators other than LogPrior (which is affected by the Jacobian of
     # the transformation).
-    accs = getaccs(vi.accs)
-    has_logprior = hasacc(accs, Val(:LogPrior))
+    accs = getaccs(vi)
+    has_logprior = haskey(accs, Val(:LogPrior))
     if has_logprior
         old_logprior = getacc(accs, Val(:LogPrior))
         vi = setaccs!!(vi, (old_logprior,))

src/varinfo.jl

@@ -450,7 +450,7 @@ function unflatten(vi::VarInfo, x::AbstractVector)
     # messes with cases like using Float32 of logprobs and Float64 for x. Also, this is just
     # plain ugly and hacky.
     et = float_type_with_fallback(eltype(x))
-    accs = map_accumulator!!(deepcopy(vi.accs), convert_eltype, et)
+    accs = map(acc -> convert_eltype(et, acc), deepcopy(getaccs(vi)))
     return VarInfo(md, accs)
 end
 
@@ -1032,15 +1032,15 @@ set_num_produce!!(vi::VarInfo, n::Int) = setacc!!(vi, NumProduce(n))
 
 Add 1 to `num_produce` in `vi`.
 """
-increment_num_produce!!(vi::VarInfo) = map_accumulator!!(vi, Val(:NumProduce), increment)
+increment_num_produce!!(vi::VarInfo) = map_accumulator!!(increment, vi, Val(:NumProduce))
 
 """
     reset_num_produce!!(vi::VarInfo)
 
 Reset the value of `num_produce` the log of the joint probability of the observed data
 and parameters sampled in `vi` to 0.
 """
-reset_num_produce!!(vi::VarInfo) = map_accumulator!!(vi, Val(:NumProduce), zero)
+reset_num_produce!!(vi::VarInfo) = map_accumulator!!(zero, vi, Val(:NumProduce))
 
 # Need to introduce the _isempty to avoid type piracy of isempty(::NamedTuple).
 isempty(vi::VarInfo) = _isempty(vi.metadata)

test/accumulators.jl

@@ -14,7 +14,7 @@ using DynamicPPL:
     convert_eltype,
     getacc,
     increment,
-    map_accumulator!!,
+    map_accumulator,
     setacc!!,
     split
 
@@ -66,8 +66,8 @@ using DynamicPPL:
                 LogLikelihood{Float32}(1.0f0)
             @test convert(NumProduce{UInt8}, NumProduce(1)) == NumProduce{UInt8}(1)
 
-            @test convert_eltype(LogPrior(1.0), Float32) == LogPrior{Float32}(1.0f0)
-            @test convert_eltype(LogLikelihood(1.0), Float32) ==
+            @test convert_eltype(Float32, LogPrior(1.0)) == LogPrior{Float32}(1.0f0)
+            @test convert_eltype(Float32, LogLikelihood(1.0)) ==
                 LogLikelihood{Float32}(1.0f0)
         end
 
@@ -137,23 +137,23 @@ using DynamicPPL:
             @test getacc(at_all64, Val(:LogPrior)) == lp_f64
         end
 
-        @testset "map_accumulator!!" begin
+        @testset "map_accumulator(s)!!" begin
             # map over all accumulators
             accs = AccumulatorTuple(lp_f32, ll_f32)
-            @test map_accumulator!!(accs, zero) ==
-                AccumulatorTuple(LogPrior(0.0f0), LogLikelihood(0.0f0))
+            @test map(zero, accs) == AccumulatorTuple(LogPrior(0.0f0), LogLikelihood(0.0f0))
             # Test that the original wasn't modified.
             @test accs == AccumulatorTuple(lp_f32, ll_f32)
 
-            # A map with extra arguments that changes the types of the accumulators.
-            @test map_accumulator!!(accs, convert_eltype, Float64) ==
+            # A map with a closure that changes the types of the accumulators.
+            @test map(acc -> convert_eltype(Float64, acc), accs) ==
                 AccumulatorTuple(LogPrior(1.0), LogLikelihood(1.0))
 
             # only apply to a particular accumulator
-            @test map_accumulator!!(accs, Val(:LogLikelihood), zero) ==
+            @test map_accumulator(zero, accs, Val(:LogLikelihood)) ==
                 AccumulatorTuple(lp_f32, LogLikelihood(0.0f0))
-            @test map_accumulator!!(accs, Val(:LogLikelihood), convert_eltype, Float64) ==
-                AccumulatorTuple(lp_f32, LogLikelihood(1.0))
+            @test map_accumulator(
+                acc -> convert_eltype(Float64, acc), accs, Val(:LogLikelihood)
+            ) == AccumulatorTuple(lp_f32, LogLikelihood(1.0))
         end
     end
 end