Add collection utils

Author: oschulz

ext/MeasureBaseChainRulesCoreExt.jl

@@ -7,6 +7,53 @@ using ChainRulesCore: NoTangent, ZeroTangent
 import ChainRulesCore
 
 
+# = collection utils =========================================================
+
+using MeasureBase: _dropfront, _dropback, _rev_cumsum, _exp_cumsum_log
+
+function ChainRulesCore.rrule(::typeof(_pushfront), v::AbstractVector, x)
+    result = _pushfront(v, x)
+    function _pushfront_pullback(thunked_ΔΩ)
+        ΔΩ = ChainRulesCore.unthunk(thunked_ΔΩ)
+        (NoTangent(), ΔΩ[firstindex(ΔΩ)+1:lastindex(ΔΩ)], ΔΩ[firstindex(ΔΩ)])
+    end
+    return result, _pushfront_pullback
+end
+
+
+function ChainRulesCore.rrule(::typeof(_pushback), v::AbstractVector, x)
+    result = _pushback(v, x)
+    function _pushback_pullback(thunked_ΔΩ)
+        ΔΩ = ChainRulesCore.unthunk(thunked_ΔΩ)
+        (NoTangent(), ΔΩ[firstindex(ΔΩ):lastindex(ΔΩ)-1], ΔΩ[lastindex(ΔΩ)])
+    end
+    return result, _pushback_pullback
+end
+
+
+function ChainRulesCore.rrule(::typeof(_rev_cumsum), xs::AbstractVector)
+    result = _rev_cumsum(xs)
+    function _rev_cumsum_pullback(ΔΩ)
+        ∂xs = ChainRulesCore.@thunk cumsum(ChainRulesCore.unthunk(ΔΩ))
+        (NoTangent(), ∂xs)
+    end
+    return result, _rev_cumsum_pullback
+end
+
+
+function ChainRulesCore.rrule(::typeof(_exp_cumsum_log), xs::AbstractVector)
+    result = _exp_cumsum_log(xs)
+    function _exp_cumsum_log_pullback(ΔΩ)
+        ∂xs = inv.(xs) .* _rev_cumsum(exp.(cumsum(log.(xs))) .* ChainRulesCore.unthunk(ΔΩ))
+        (NoTangent(), ∂xs)
+    end
+    return result, _exp_cumsum_log_pullback
+end
+
+
+# = measure functions ========================================================
+
+
 @inline function ChainRulesCore.rrule(::typeof(_checksupport), cond, result)
     y = _checksupport(cond, result)
     function _checksupport_pullback(ȳ)

src/MeasureBase.jl

@@ -144,6 +144,7 @@ using Compat
 using IrrationalConstants
 
 include("static.jl")
+include("collection_utils.jl")
 include("smf.jl")
 include("getdof.jl")
 include("transport.jl")

src/collection_utils.jl

@@ -0,0 +1,24 @@
+function _pushfront(v::AbstractVector, x)
+    T = promote_type(eltype(v), typeof(x))
+    r = similar(v, T, length(eachindex(v)) + 1)
+    r[firstindex(r)] = x
+    r[firstindex(r)+1:lastindex(r)] = v
+    r
+end
+
+function _pushback(v::AbstractVector, x)
+    T = promote_type(eltype(v), typeof(x))
+    r = similar(v, T, length(eachindex(v)) + 1)
+    r[lastindex(r)] = x
+    r[firstindex(r):lastindex(r)-1] = v
+    r
+end
+
+_dropfront(v::AbstractVector) = v[firstindex(v)+1:lastindex(v)]
+
+_dropback(v::AbstractVector) = v[firstindex(v):lastindex(v)-1]
+
+_rev_cumsum(xs::AbstractVector) = reverse(cumsum(reverse(xs)))
+
+# Equivalent to `cumprod(xs)``:
+_exp_cumsum_log(xs::AbstractVector) = exp.(cumsum(log.(xs)))