Handle NoLogAbsDetJacobian in ComposedFunction

Author: oschulz

src/with_ladj.jl

@@ -86,9 +86,19 @@ with_logabsdet_jacobian(::F, ::T) where {F,T} = NoLogAbsDetJacobian{F,T}()
 
 @static if VERSION >= v"1.6"
     function with_logabsdet_jacobian(f::Base.ComposedFunction, x)
-        y_inner, ladj_inner = with_logabsdet_jacobian(f.inner, x)
-        y, ladj_outer = with_logabsdet_jacobian(f.outer, y_inner)
-        (y, ladj_inner + ladj_outer)
+        y_ladj_inner = with_logabsdet_jacobian(f.inner, x)
+        if y_ladj_inner isa NoLogAbsDetJacobian
+            NoLogAbsDetJacobian{typeof(f),typeof(x)}()
+        else
+            y_inner, ladj_inner = y_ladj_inner
+            y_ladj_outer = with_logabsdet_jacobian(f.outer, y_inner)
+            if y_ladj_outer isa NoLogAbsDetJacobian
+                NoLogAbsDetJacobian{typeof(f),typeof(x)}()
+            else
+                y, ladj_outer = y_ladj_outer
+                (y, ladj_inner + ladj_outer)
+            end
+        end
     end
 end
 

test/test_with_ladj.jl

@@ -14,6 +14,8 @@ include("getjacobian.jl")
 
 @testset "with_logabsdet_jacobian" begin
     @test with_logabsdet_jacobian(sum, rand(5)) == NoLogAbsDetJacobian{typeof(sum),Vector{Float64}}()
+    @test with_logabsdet_jacobian(sum ∘ log, 5.0f0) == NoLogAbsDetJacobian{typeof(sum ∘ log),Float32}()
+    @test with_logabsdet_jacobian(log ∘ sum, 5.0f0) == NoLogAbsDetJacobian{typeof(log ∘ sum),Float32}()
     @test_throws MethodError _, _ = with_logabsdet_jacobian(sum, rand(5))
 
     function ChangesOfVariables.with_logabsdet_jacobian(::typeof(foo), x)