Add curried pushfwd and pullbck

Author: oschulz

src/combinators/transformedmeasure.jl

@@ -222,25 +222,52 @@ To manually specify an inverse, call
 function pushfwd end
 export pushfwd
 
-@inline pushfwd(f, μ) = _pushfwd_impl(f, μ, AdaptRootMeasure())
-@inline pushfwd(f, μ, style::AdaptRootMeasure) = _pushfwd_impl(f, μ, style)
-@inline pushfwd(f, μ, style::PushfwdRootMeasure) = _pushfwd_impl(f, μ, style)
+@inline pushfwd(f, μ) = _pushfwd_impl1(f, μ, AdaptRootMeasure())
+@inline pushfwd(f, μ, style::AdaptRootMeasure) = _pushfwd_impl1(f, μ, style)
+@inline pushfwd(f, μ, style::PushfwdRootMeasure) = _pushfwd_impl1(f, μ, style)
 
-_pushfwd_impl(f, μ, style) = PushforwardMeasure(f, inverse(f), μ, style)
+_pushfwd_impl1(f, μ, style::PushFwdStyle) = _pushfwd_impl2(f, inverse(f), μ, style)
+_pushfwd_impl1(::typeof(identity), μ, ::AdaptRootMeasure) = μ
+_pushfwd_impl1(::typeof(identity), μ, ::PushfwdRootMeasure) = μ
 
-function _pushfwd_impl(
+_pushfwd_impl2(f, finv, μ, style::PushFwdStyle) = PushforwardMeasure(f, finv, μ, style)
+
+function _pushfwd_impl2(
     f,
+    finv,
     μ::PushforwardMeasure{F,I,M,S},
     style::S,
 ) where {F,I,M,S<:PushFwdStyle}
     orig_μ = μ.origin
     new_f = fcomp(f, μ.f)
-    new_f_inv = fcomp(μ.finv, inverse(f))
+    new_f_inv = fcomp(μ.finv, finv)
     PushforwardMeasure(new_f, new_f_inv, orig_μ, style)
 end
 
-_pushfwd_impl(::typeof(identity), μ, ::AdaptRootMeasure) = μ
-_pushfwd_impl(::typeof(identity), μ, ::PushfwdRootMeasure) = μ
+struct _CurriedPushfwd{F,I,S<:PushFwdStyle} <: Function
+    f::F
+    finv::I
+    style::S
+
+    function _CurriedPushfwd{F,I,S}(f::F, finv::I, style::S) where {F,I,S<:PushFwdStyle}
+        new{F,I,S}(f, finv, style)
+    end
+
+    function _CurriedPushfwd(f, finv, style::S) where {S<:PushFwdStyle}
+        new{Core.Typeof(f),Core.Typeof(finv),S}(f, finv, style)
+    end
+end
+
+@inline (cf::_CurriedPushfwd{F,FI})(μ) where {F,FI} =
+    _pushfwd_impl2(cf.f, cf.finv, μ, cf.style)
+
+@inline pushfwd(f) = _curried_pushfwd_impl(f, AdaptRootMeasure())
+@inline pushfwd(f, style::AdaptRootMeasure) = _curried_pushfwd_impl(f, style)
+@inline pushfwd(f, style::PushfwdRootMeasure) = _curried_pushfwd_impl(f, style)
+
+_curried_pushfwd_impl(f, style::PushFwdStyle) = _CurriedPushfwd(f, inverse(f), style)
+@inline _curried_pushfwd_impl(::typeof(identity), ::AdaptRootMeasure) = identity
+@inline _curried_pushfwd_impl(::typeof(identity), ::PushfwdRootMeasure) = identity
 
 ###############################################################################
 # pullback
@@ -267,8 +294,16 @@ export pullbck
 @inline pullbck(f, μ, style::AdaptRootMeasure) = _pullback_impl(f, μ, style)
 @inline pullbck(f, μ, style::PushfwdRootMeasure) = _pullback_impl(f, μ, style)
 
-function _pullback_impl(f, μ, style = AdaptRootMeasure())
-    pushfwd(inverse(f), μ, style)
-end
+_pullback_impl(f, μ, style::PushFwdStyle) = _pushfwd_impl2(inverse(f), f, μ, style)
+_pullback_impl(::typeof(identity), μ, ::AdaptRootMeasure) = μ
+_pullback_impl(::typeof(identity), μ, ::PushfwdRootMeasure) = μ
+
+@inline pullbck(f) = _curried_pullbck_impl(f, AdaptRootMeasure())
+@inline pullbck(f, style::AdaptRootMeasure) = _curried_pullbck_impl(f, style)
+@inline pullbck(f, style::PushfwdRootMeasure) = _curried_pullbck_impl(f, style)
+
+_curried_pullbck_impl(f, style::PushFwdStyle) = _CurriedPushfwd(inverse(f), f, style)
+@inline _curried_pullbck_impl(::typeof(identity), ::AdaptRootMeasure) = identity
+@inline _curried_pullbck_impl(::typeof(identity), ::PushfwdRootMeasure) = identity
 
 @deprecate pullback(f, μ, style::PushFwdStyle = AdaptRootMeasure()) pullbck(f, μ, style)

test/Project.toml

@@ -7,6 +7,7 @@ DensityInterface = "b429d917-457f-4dbc-8f4c-0cc954292b1d"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
+FunctionChains = "8e6b2b91-af83-483e-ba35-d00930e4cf9b"
 InverseFunctions = "3587e190-3f89-42d0-90ee-14403ec27112"
 IrrationalConstants = "92d709cd-6900-40b7-9082-c6be49f344b6"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

test/combinators/transformedmeasure.jl

@@ -12,6 +12,7 @@ import InverseFunctions: inverse, FunctionWithInverse, setinverse
 using IrrationalConstants: invsqrt2, sqrt2
 import ChangesOfVariables: with_logabsdet_jacobian
 using MeasureBase.Interface: transport_to, test_transport
+using FunctionChains: fchain
 
 Φ(z) = erfc(-z * invsqrt2) / 2
 Φinv(p) = -erfcinv(2 * p) * sqrt2
@@ -106,7 +107,7 @@ using ChangesOfVariables
     # Test basic pushforward construction
     μ = StdNormal()
     f = exp
-    ν = pushfwd(f, μ)
+    ν = @inferred pushfwd(f, μ)
 
     @test ν isa PushforwardMeasure
 
@@ -167,10 +168,29 @@ using ChangesOfVariables
     # Test pullback
     pb = pullbck(f, ν)
     @test pb isa PushforwardMeasure
+    @test pb.origin === μ
+    @test pb.f === fchain(exp, log)
     @test logdensityof(pb, y) ≈ logdensityof(μ, y)
 
     # Test deprecated pullback
     @test_deprecated pullback(f, μ)
+
+    # Test identity specializations
+    for stylearg in [(), (AdaptRootMeasure(),), (PushfwdRootMeasure(),)]
+        @test @inferred(pushfwd(identity, μ, stylearg...)) === μ
+        @test @inferred(pullbck(identity, ν, stylearg...)) === ν
+    end
+
+    # Test curried pushfwd and pullback
+    for stylearg in [(), (AdaptRootMeasure(),), (PushfwdRootMeasure(),)]
+        @test @inferred(pushfwd(f, stylearg...)(μ)) === pushfwd(f, μ, stylearg...)
+        @test @inferred(pullbck(f, stylearg...)(ν)) === pullbck(f, ν, stylearg...)
+
+        @test @inferred(pushfwd(identity, stylearg...)) === identity
+        @test @inferred(pullbck(identity, stylearg...)) === identity
+    end
+
+    @test pushfwd(identity) === identity
 end
 
 @testset "PushFwdStyle types" begin