Improve _with_ladj_on_mapped pullback implementation

Co-authored-by: David Widmann <[email protected]>

Author: oschulz

src/with_ladj.jl

@@ -81,20 +81,20 @@ end
 
 _with_ladj_on_mapped(map_or_bc::Function, y_with_ladj::Tuple{Any,Real}) = y_with_ladj
 
-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))
+function _with_ladj_on_mapped(map_or_bc::F, y_with_ladj) where {F<:Union{typeof(map),typeof(broadcast)}}
+    y = map_or_bc(first, y_with_ladj)
+    ladj = sum(map_or_bc(last, 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)
+    ys, ladj = unthunk(thunked_ΔΩ)
+    return NoTangent(), NoTangent(), tuple.(ys, ladj)
 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
+function ChainRulesCore.rrule(::typeof(_with_ladj_on_mapped), map_or_bc::Function, y_with_ladj)
+    return _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)