Fix with_logabsdet_jacobian for mapped/broadcasted functions

Handle NoLogAbsDetJacobian properly.

Author: oschulz

src/with_ladj.jl

@@ -117,37 +117,61 @@ with_logabsdet_jacobian(f, x) = NoLogAbsDetJacobian(f, x)
 end
 
 
-function _with_ladj_on_mapped(@nospecialize(map_or_bc::F), y_with_ladj::NoLogAbsDetJacobian) where {F<:Union{typeof(map),typeof(broadcast)}}
-    return y_with_ladj
-end
+_with_ladj_on_mapped(y_with_ladj::NoLogAbsDetJacobian) = y_with_ladj
 
-function _with_ladj_on_mapped(map_or_bc::F, y_with_ladj::Tuple{Any,Real})  where {F<:Union{typeof(map),typeof(broadcast)}}
-    return y_with_ladj
-end
+_with_ladj_on_mapped(y_with_ladj::Tuple{Any,Real}) = y_with_ladj
+
+_with_ladj_on_mapped(::AbstractArray{T})  where {T<:NoLogAbsDetJacobian} = T()
 
 _get_all_first(x) = map(first, x)
-# Use x -> x[2] instead of last, using last causes horrible performance in Zygote here:
-_sum_over_second(x) = sum(x -> x[2], x)
+_get_second(x) = x[2]
+# Use _get_second instead of last, using last causes horrible performance in Zygote here:
+_sum_over_second(x) = sum(_get_second, x)
 
-function _with_ladj_on_mapped(map_or_bc::F, y_with_ladj) where {F<:Union{typeof(map),typeof(broadcast)}}
+function _with_ladj_on_mapped(y_with_ladj)
     y = _get_all_first(y_with_ladj)
     ladj = _sum_over_second(y_with_ladj)
     (y, ladj)
 end
 
+_combine_y_ladj(l::NoLogAbsDetJacobian, @nospecialize(r::NoLogAbsDetJacobian)) = l
+_combine_y_ladj(l::NoLogAbsDetJacobian, @nospecialize(r::Tuple{Any,Real})) = l
+_combine_y_ladj(@nospecialize(l::Tuple{Any,Real}), r::NoLogAbsDetJacobian) = r
+_combine_y_ladj(l::Tuple{Any,Real}, r::Tuple{Any,Real}) = ((l[1]..., r[1]), l[2] + r[2])
+
+function _with_ladj_on_mapped(y_with_ladj::Tuple{NoLogAbsDetJacobian, Vararg{Union{NoLogAbsDetJacobian, Tuple{Any,Real}}}})
+    return first(y_with_ladj)
+end
+
+@static if VERSION < v"1.6"
+    # Hand-rolled foldl, foldl over tuples is not inferrable on Julia < 1.6:
+    _foldl_combine_y_ladj(acc) = acc
+    _foldl_combine_y_ladj(acc, b, bs...) = _foldl_combine_y_ladj(_combine_y_ladj(acc, b), bs...)
+
+    function _with_ladj_on_mapped(y_with_ladj::Tuple{Tuple{Any,Real}, Vararg{Union{NoLogAbsDetJacobian, Tuple{Any,Real}}}})
+        a = first(y_with_ladj)
+        return _foldl_combine_y_ladj(((a[1],), a[2]), Base.tail(y_with_ladj)...)
+    end
+else
+    function _with_ladj_on_mapped(y_with_ladj::Tuple{Tuple{Any,Real}, Vararg{Union{NoLogAbsDetJacobian, Tuple{Any,Real}}}})
+        a, bs = first(y_with_ladj), Base.tail(y_with_ladj)
+        return foldl(_combine_y_ladj, bs, init = ((a[1],), a[2]))
+    end
+end
+
 @static if VERSION >= v"1.6"
     function with_logabsdet_jacobian(mapped_f::Base.Broadcast.BroadcastFunction, X)
         f = mapped_f.f
         y_with_ladj = broadcast(Base.Fix1(with_logabsdet_jacobian, f), X)
-        _with_ladj_on_mapped(broadcast, y_with_ladj)
+        _with_ladj_on_mapped(y_with_ladj)
     end
 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
     y_with_ladj = map_or_bc(Base.Fix1(with_logabsdet_jacobian, f), X)
-    _with_ladj_on_mapped(map_or_bc, y_with_ladj)
+    _with_ladj_on_mapped(y_with_ladj)
 end
 
 

test/getjacobian.jl

@@ -43,3 +43,12 @@ if !isdefined(Main, :foo)
 foo(x) = inv(exp(-x) + 1)
 
 end # !isdefined(Main, :foo)
+
+
+if !isdefined(Main, :bar)
+
+# with_logabsdet_jacobian for floats but not for integers 
+bar(x) = x
+ChangesOfVariables.with_logabsdet_jacobian(::typeof(bar), x::AbstractFloat) = (x, zero(x))
+
+end # !isdefined(Main, :bar)

test/test_with_ladj.jl

@@ -55,12 +55,26 @@ include("getjacobian.jl")
 
     @testset "with_logabsdet_jacobian on mapped and broadcasted" begin
         for f in (_bc_func(foo), Base.Fix1(map, foo), Base.Fix1(broadcast, foo))
-            for arg in (x, fill(x,), Ref(x), (x,), X)
+            for arg in (x, fill(x,), Ref(x), (x,), (x, 2*x, 3*x), X)
                 test_with_logabsdet_jacobian(f, arg, getjacobian, compare = isaprx)
             end
         end
     end
 
+    @testset "with_logabsdet_jacobian on mapped and broadcasted without ladj" begin
+        for f in (_bc_func(sin), Base.Fix1(map, sin), Base.Fix1(broadcast, sin))
+            for arg in (x, (x,), (x, x), X)
+                @test with_logabsdet_jacobian(f, arg) isa NoLogAbsDetJacobian{typeof(sin)}
+            end
+        end
+    end
+
+    @testset "with_logabsdet_jacobian on mapped and broadcasted with mixed ladj" begin
+        @test with_logabsdet_jacobian(_bc_func(bar), (1.0, 2.0)) == ((1.0, 2.0), 0.0)
+        @test with_logabsdet_jacobian(_bc_func(bar), (1.0, 2)) isa NoLogAbsDetJacobian
+        @test with_logabsdet_jacobian(Base.Fix1(map, bar), (1, 2.0)) isa NoLogAbsDetJacobian
+    end
+
     @testset "with_logabsdet_jacobian on identity, adjoint and transpose" begin
         for f in (identity, adjoint, transpose)
             for arg in (x, A)