revert dot_assume to not explicitly resolve components of sum

Author: bgctw

src/context_implementations.jl

@@ -425,15 +425,6 @@ function dot_assume(
     var::AbstractMatrix,
     vns::AbstractVector{<:VarName},
     vi::AbstractVarInfo,
-)
-    r, lp, vi = dot_assume_vec(dist, var, vns, vi)
-    return r, sum(lp), vi
-end
-function dot_assume_vec(
-    dist::MultivariateDistribution,
-    var::AbstractMatrix,
-    vns::AbstractVector{<:VarName},
-    vi::AbstractVarInfo,
 )
     @assert length(dist) == size(var, 1) "dimensionality of `var` ($(size(var, 1))) is incompatible with dimensionality of `dist` $(length(dist))"
     # NOTE: We cannot work with `var` here because we might have a model of the form
@@ -443,7 +434,7 @@ function dot_assume_vec(
     #
     # in which case `var` will have `undef` elements, even if `m` is present in `vi`.
     r = vi[vns, dist]
-    lp = map(zip(vns, eachcol(r))) do (vn, ri)
+    lp = sum(zip(vns, eachcol(r))) do (vn, ri)
         return Bijectors.logpdf_with_trans(dist, ri, istrans(vi, vn))
     end
     return r, lp, vi
@@ -464,29 +455,21 @@ function dot_assume(
 end
 
 function dot_assume(
-    dist::Union{Distribution,AbstractArray{<:Distribution}}, var::AbstractArray, vns::AbstractArray{<:VarName}, vi
-)
-    # possibility to acesss the single logpriors
-    r, lp, vi = dot_assume_vec(dist, var, vns, vi)
-    return r, sum(lp), vi
-end
-
-function dot_assume_vec(
     dist::Distribution, var::AbstractArray, vns::AbstractArray{<:VarName}, vi
 )
     r = getindex.((vi,), vns, (dist,))
-    lp = Bijectors.logpdf_with_trans.((dist,), r, istrans.((vi,), vns))
+    lp = sum(Bijectors.logpdf_with_trans.((dist,), r, istrans.((vi,), vns)))
     return r, lp, vi
 end
 
-function dot_assume_vec(
+function dot_assume(
     dists::AbstractArray{<:Distribution},
     var::AbstractArray,
     vns::AbstractArray{<:VarName},
     vi,
 )
     r = getindex.((vi,), vns, dists)
-    lp = Bijectors.logpdf_with_trans.(dists, r, istrans.((vi,), vns))
+    lp = sum(Bijectors.logpdf_with_trans.(dists, r, istrans.((vi,), vns)))
     return r, lp, vi
 end