Use Tangent in _with_ladj_on_mapped pullback

Author: oschulz

Project.toml

@@ -12,8 +12,9 @@ ChainRulesCore = "1"
 julia = "1"
 
 [extras]
+ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 
 [targets]
-test = ["Documenter", "ForwardDiff"]
+test = ["ChainRulesTestUtils", "Documenter", "ForwardDiff"]

src/with_ladj.jl

@@ -86,13 +86,18 @@ function _with_ladj_on_mapped(map_or_bc::F, y_with_ladj) where {F<:Union{typeof(
     (y, ladj)
 end
 
-function _with_ladj_on_mapped_pullback(thunked_ΔΩ)
+
+# Need to use a type for this, type inference fails when using a pullback
+# closure over YLT in the rrule, resulting in bad performance:
+struct WithLadjOnMappedPullback{YLT} <: Function end
+function (::WithLadjOnMappedPullback{YLT})(thunked_ΔΩ) where YLT
     ys, ladj = unthunk(thunked_ΔΩ)
-    return NoTangent(), NoTangent(), tuple.(ys, ladj)
+    return NoTangent(), NoTangent(), broadcast((y, l) -> Tangent{YLT}(y, l), ys, ladj)
 end
 
 function ChainRulesCore.rrule(::typeof(_with_ladj_on_mapped), map_or_bc::F, y_with_ladj) where {F<:Union{typeof(map),typeof(broadcast)}}
-    return _with_ladj_on_mapped(map_or_bc, y_with_ladj), _with_ladj_on_mapped_pullback
+    YLT = eltype(y_with_ladj)
+    return _with_ladj_on_mapped(map_or_bc, y_with_ladj), WithLadjOnMappedPullback{YLT}()
 end
 
 function with_logabsdet_jacobian(mapped_f::Base.Fix1{<:Union{typeof(map),typeof(broadcast)}}, X)

test/test_with_ladj.jl

@@ -5,8 +5,9 @@ using Test
 
 using LinearAlgebra
 
-using ChangesOfVariables: test_with_logabsdet_jacobian, _with_ladj_on_mapped
-using ChainRulesCore
+using ChangesOfVariables
+using ChangesOfVariables: test_with_logabsdet_jacobian
+using ChainRulesTestUtils
 
 include("getjacobian.jl")
 
@@ -63,10 +64,7 @@ include("getjacobian.jl")
 
     @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)])
+            test_rrule(ChangesOfVariables._with_ladj_on_mapped, map_or_bc, [(1.0, 2.0), (3.0, 4.0), (5.0, 6.0)])
         end
     end
 end