Merge remote-tracking branch 'origin/breaking' into mhauru/logprobacc

Author: mhauru

src/abstract_varinfo.jl

@@ -204,7 +204,7 @@ Jacobian here is taken with respect to the forward (link) transform.
 
 See also: [`setlogjac!!`](@ref).
 """
-getlogjac(vi::AbstractVarInfo) = getacc(vi, Val(:LogJacobian)).logJ
+getlogjac(vi::AbstractVarInfo) = getacc(vi, Val(:LogJacobian)).logjac
 
 """
     getloglikelihood(vi::AbstractVarInfo)
@@ -239,14 +239,14 @@ See also: [`setloglikelihood!!`](@ref), [`setlogp!!`](@ref), [`getlogprior`](@re
 setlogprior!!(vi::AbstractVarInfo, logp) = setacc!!(vi, LogPriorAccumulator(logp))
 
 """
-    setlogjac!!(vi::AbstractVarInfo, logJ)
+    setlogjac!!(vi::AbstractVarInfo, logjac)
 
 Set the accumulated log-Jacobian term for any linked parameters in `vi`. The
 Jacobian here is taken with respect to the forward (link) transform.
 
 See also: [`getlogjac`](@ref), [`acclogjac!!`](@ref).
 """
-setlogjac!!(vi::AbstractVarInfo, logJ) = setacc!!(vi, LogJacobianAccumulator(logJ))
+setlogjac!!(vi::AbstractVarInfo, logjac) = setacc!!(vi, LogJacobianAccumulator(logjac))
 
 """
     setloglikelihood!!(vi::AbstractVarInfo, logp)
@@ -372,6 +372,19 @@ function acclogjac!!(vi::AbstractVarInfo, logJ)
     return map_accumulator!!(acc -> acclogp(acc, logJ), vi, Val(:LogJacobian))
 end
 
+"""
+    acclogjac!!(vi::AbstractVarInfo, logjac)
+
+Add `logjac` to the value of the log Jacobian in `vi`.
+
+See also: [`getlogjac`](@ref), [`setlogjac!!`](@ref).
+"""
+function acclogjac!!(vi::AbstractVarInfo, logjac)
+    return map_accumulator!!(
+        acc -> acc + LogJacobianAccumulator(logjac), vi, Val(:LogJacobian)
+    )
+end
+
 """
     accloglikelihood!!(vi::AbstractVarInfo, logp)
 
@@ -903,28 +916,29 @@ function link!!(
     x = vi[:]
     y, logjac = with_logabsdet_jacobian(b, x)
 
-    # Set parameters
-    vi_new = unflatten(vi, y)
-    # Update logjac. We can overwrite any old value since there is only
-    # a single logjac term to worry about.
-    vi_new = setlogjac!!(vi_new, logjac)
-    return settrans!!(vi_new, t)
+    # Set parameters and add the logjac term.
+    vi = unflatten(vi, y)
+    if hasacc(vi, Val(:LogJacobian))
+        vi = acclogjac!!(vi, logjac)
+    end
+    return settrans!!(vi, t)
 end
 
 function invlink!!(
     t::StaticTransformation{<:Bijectors.Transform}, vi::AbstractVarInfo, ::Model
 )
     b = t.bijector
     y = vi[:]
-    x = b(y)
+    x, inv_logjac = with_logabsdet_jacobian(b, y)
 
-    # Set parameters
-    vi_new = unflatten(vi, x)
-    # Reset logjac to 0.
-    if hasacc(vi_new, Val(:LogJacobian))
-        vi_new = map_accumulator!!(zero, vi_new, Val(:LogJacobian))
+    # Mildly confusing: we need to _add_ the logjac of the inverse transform,
+    # because we are trying to remove the logjac of the forward transform
+    # that was previously accumulated when linking.
+    vi = unflatten(vi, x)
+    if hasacc(vi, Val(:LogJacobian))
+        vi = acclogjac!!(vi, inv_logjac)
     end
-    return settrans!!(vi_new, NoTransformation())
+    return settrans!!(vi, NoTransformation())
 end
 
 """

src/accumulators.jl

@@ -16,8 +16,8 @@ An accumulator type `T <: AbstractAccumulator` must implement the following meth
 - `Base.copy(acc::T)`
 
 In these functions:
-- `val` is the new value of the random variable sampled from a new distribution (always
-  in the original unlinked space), or the value on the left-hand side of an observe
+- `val` is the new value of the random variable sampled from a distribution (always in
+  the original unlinked space), or the value on the left-hand side of an observe
   statement.
 - `dist` is the distribution on the RHS of the tilde statement.
 - `vn` is the `VarName` that is on the left-hand side of the tilde-statement. If the

src/default_accumulators.jl

@@ -129,10 +129,10 @@ $(TYPEDFIELDS)
 """
 struct LogJacobianAccumulator{T<:Real} <: LogProbAccumulator{T}
     "the logabsdet of the link transform Jacobian"
-    logJ::T
+    logjac::T
 end
 
-logp(acc::LogJacobianAccumulator) = acc.logJ
+logp(acc::LogJacobianAccumulator) = acc.logjac
 
 accumulator_name(::Type{<:LogJacobianAccumulator}) = :LogJacobian
 

src/simple_varinfo.jl

@@ -622,10 +622,8 @@ function link!!(
     x = vi.values
     y, logjac = with_logabsdet_jacobian(b, x)
     vi_new = Accessors.@set(vi.values = y)
-    # Since there's only a single transformation, we can overwrite any previous
-    # value in logjac.
     if hasacc(vi_new, Val(:LogJacobian))
-        vi_new = setlogjac!!(vi_new, logjac)
+        vi_new = acclogjac!!(vi_new, logjac)
     end
     return settrans!!(vi_new, t)
 end
@@ -637,11 +635,13 @@ function invlink!!(
 )
     b = t.bijector
     y = vi.values
-    x = b(y)
+    x, inv_logjac = with_logabsdet_jacobian(b, y)
     vi_new = Accessors.@set(vi.values = x)
-    # logjac should be zero for an unlinked VarInfo.
+    # Mildly confusing: we need to _add_ the logjac of the inverse transform,
+    # because we are trying to remove the logjac of the forward transform
+    # that was previously accumulated when linking.
     if hasacc(vi_new, Val(:LogJacobian))
-        vi_new = map_accumulator!!(zero, vi_new, Val(:LogJacobian))
+        vi_new = acclogjac!!(vi_new, inv_logjac)
     end
     return settrans!!(vi_new, NoTransformation())
 end