Add custom rrule for _with_ladj_on_mapped

Speeds up AD of with_logabsdet_jacobian with mapped/broadcasted
functions significantly.

Author: oschulz

Project.toml

@@ -3,10 +3,12 @@ uuid = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0"
 version = "0.1.1"
 
 [deps]
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]
+ChainRulesCore = "1"
 julia = "1"
 
 [extras]

src/ChangesOfVariables.jl

@@ -9,6 +9,7 @@ transformations).
 """
 module ChangesOfVariables
 
+using ChainRulesCore
 using LinearAlgebra
 using Test
 

src/with_ladj.jl

@@ -84,9 +84,19 @@ _with_ladj_on_mapped(map_or_bc::Function, y_with_ladj::Tuple{Any,Real}) = y_with
 function _with_ladj_on_mapped(map_or_bc::Function, y_with_ladj)
     y = map_or_bc(_get_y, y_with_ladj)
     ladj = sum(map_or_bc(_get_ladj, y_with_ladj))
+    #ladj = sum(_get_ladj, y_with_ladj)
     (y, ladj)
 end
 
+function _with_ladj_on_mapped_pullback(thunked_ΔΩ)
+    ys, ladj = ChainRulesCore.unthunk(thunked_ΔΩ)
+    NoTangent(), NoTangent(), broadcast(x -> (x, ladj), ys)
+end
+
+function ChainRulesCore.rrule(::typeof(ChangesOfVariables._with_ladj_on_mapped), map_or_bc::Function, y_with_ladj)
+    return ChangesOfVariables._with_ladj_on_mapped(map_or_bc, y_with_ladj), _with_ladj_on_mapped_pullback
+end
+
 function with_logabsdet_jacobian(mapped_f::Base.Fix1{<:Union{typeof(map),typeof(broadcast)}}, X)
     map_or_bc = mapped_f.f
     f = mapped_f.x

test/test_with_ladj.jl

@@ -5,7 +5,8 @@ using Test
 
 using LinearAlgebra
 
-using ChangesOfVariables: test_with_logabsdet_jacobian
+using ChangesOfVariables: test_with_logabsdet_jacobian, _with_ladj_on_mapped
+using ChainRulesCore
 
 include("getjacobian.jl")
 
@@ -59,4 +60,13 @@ include("getjacobian.jl")
             test_with_logabsdet_jacobian(f, x, getjacobian)
         end
     end
+
+    @testset "rrules" begin
+        for map_or_bc in (map, broadcast)
+            x = [(1, 2), (3, 4), (5, 6)]
+            y, back = rrule(_with_ladj_on_mapped, map_or_bc, x)
+            @test y == ([1, 3, 5], 12) == _with_ladj_on_mapped(map_or_bc, x)
+            @test back(@thunk ([7, 8, 9], 12)) == (ChainRulesCore.NoTangent(), ChainRulesCore.NoTangent(), [(7, 12), (8, 12), (9, 12)])
+        end
+    end
 end