Fix with_logabsdet_jacobian for mapped/broadcasted functions

.github/workflows/ci.yml

@@ -24,6 +24,7 @@ jobs:
           - 'min'
           - 'lts'
           - '1'
+          - 'pre'
         os:
           - ubuntu-latest
         include:

Project.toml

@@ -1,6 +1,6 @@
 name = "ChangesOfVariables"
 uuid = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0"
-version = "0.1.10"
+version = "0.1.11"
 
 [deps]
 InverseFunctions = "3587e190-3f89-42d0-90ee-14403ec27112"
@@ -19,7 +19,7 @@ ChangesOfVariablesTestExt = "Test"
 InverseFunctions = "0.1"
 LinearAlgebra = "<0.0.1, 1"
 Test = "<0.0.1, 1"
-julia = "1"
+julia = "1.10"
 
 [extras]
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"

src/ChangesOfVariables.jl

@@ -35,24 +35,16 @@ with automatic differentiation.
 `kwargs...` are forwarded to `compare`.
 
 !!! Note
-    On Julia >= 1.9, you have to load the `Test` standard library to be able to use
-    this function.
+    You have to load the `Test` standard library to be able to use this function.
 """
 function test_with_logabsdet_jacobian end
 
-@static if !isdefined(Base, :get_extension)
-    include("../ext/ChangesOfVariablesInverseFunctionsExt.jl")
-    include("../ext/ChangesOfVariablesTestExt.jl")
-end
-
 # Better error message if users forget to load Test
-if isdefined(Base, :get_extension) && isdefined(Base.Experimental, :register_error_hint)
-    function __init__()
-        Base.Experimental.register_error_hint(MethodError) do io, exc, _, _
-            if exc.f === test_with_logabsdet_jacobian &&
-                (Base.get_extension(ChangesOfVariables, :ChangesOfVariablesTest) === nothing)
-                print(io, "\nDid you forget to load Test?")
-            end
+function __init__()
+    Base.Experimental.register_error_hint(MethodError) do io, exc, _, _
+        if exc.f === test_with_logabsdet_jacobian &&
+            (Base.get_extension(ChangesOfVariables, :ChangesOfVariablesTest) === nothing)
+            print(io, "\nDid you forget to load Test?")
         end
     end
 end

src/setladj.jl

@@ -45,20 +45,19 @@ if `f` does so.
 Example:
 
 ```jldoctest setladj
-VERSION < v"1.6" || begin # Support for ∘ requires Julia >= v1.6
-    # Increases precition before calculation exp:
-    foo = exp ∘ setladj(setinverse(Float64, Float32), _ -> 0)
+# Increases precision before calculating exp:
+foo = exp ∘ setladj(setinverse(Float64, Float32), _ -> 0)
 
-    # A log-value from some low-precision (e.g. GPU) computation:
-    log_x = Float32(100)
+# A log-value from some low-precision (e.g. GPU) computation:
+log_x = Float32(100)
 
-    # f(log_x) would return Inf32 without going to Float64:
-    y, ladj = with_logabsdet_jacobian(foo, log_x) 
+# f(log_x) would return Inf32 without going to Float64:
+y, ladj = with_logabsdet_jacobian(foo, log_x)
 
-    r_log_x, ladj_inv = with_logabsdet_jacobian(inverse(foo), y)
+r_log_x, ladj_inv = with_logabsdet_jacobian(inverse(foo), y)
+
+ladj ≈ 100 ≈ -ladj_inv && r_log_x ≈ log_x
 
-    ladj ≈ 100 ≈ -ladj_inv && r_log_x ≈ log_x
-end
 # output
 
 true

src/with_ladj.jl

@@ -43,11 +43,7 @@ true
 
 ```jldoctest a
 X = rand(10)
-broadcasted_foo = if VERSION >= v"1.6"
-    Base.Broadcast.BroadcastFunction(foo)
-else
-    Base.Fix1(broadcast, foo)
-end
+broadcasted_foo = Base.Broadcast.BroadcastFunction(foo)
 Y, ladj_Y = with_logabsdet_jacobian(broadcasted_foo, X)
 Y == broadcasted_foo(X) && ladj_Y ≈ logabsdet(ForwardDiff.jacobian(broadcasted_foo, X))[1]
 
@@ -57,10 +53,8 @@ true
 ```
 
 ```jldoctest a
-VERSION < v"1.6" || begin # Support for ∘ requires Julia >= v1.6
-    z, ladj_z = with_logabsdet_jacobian(log ∘ foo, x)
-    z == log(foo(x)) && ladj_z == ladj_y + with_logabsdet_jacobian(log, y)[2]
-end
+z, ladj_z = with_logabsdet_jacobian(log ∘ foo, x)
+z == log(foo(x)) && ladj_z == ladj_y + with_logabsdet_jacobian(log, y)[2]
 
 # output
 
@@ -98,56 +92,58 @@ export NoLogAbsDetJacobian
 with_logabsdet_jacobian(f, x) = NoLogAbsDetJacobian(f, x)
 
 
-@static if VERSION >= v"1.6"
-    function with_logabsdet_jacobian(f::Base.ComposedFunction, x)
-        y_ladj_inner = with_logabsdet_jacobian(f.inner, x)
-        if y_ladj_inner isa NoLogAbsDetJacobian
+function with_logabsdet_jacobian(f::Base.ComposedFunction, x)
+    y_ladj_inner = with_logabsdet_jacobian(f.inner, x)
+    if y_ladj_inner isa NoLogAbsDetJacobian
+        NoLogAbsDetJacobian(f, 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(f, 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(f, x)
-            else
-                y, ladj_outer = y_ladj_outer
-                (y, ladj_inner + ladj_outer)
-            end
+            y, ladj_outer = y_ladj_outer
+            (y, ladj_inner + ladj_outer)
         end
     end
 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
+_get_all_first(x) = map(first, 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)
+
+_combine_y_ladj(l, r) = ((l[1]..., r[1]), l[2] + r[2])
 
-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
+_with_ladj_on_mapped(y_with_ladj::Tuple{Any,Any}, ::NoLogAbsDetJacobian) = y_with_ladj
+
+function _with_ladj_on_mapped(y_with_ladj::Tuple{Tuple{Any,Any}, Vararg{Tuple{Any,Any}}}, ::NoLogAbsDetJacobian)
+    a, bs = first(y_with_ladj), Base.tail(y_with_ladj)
+    return foldl(_combine_y_ladj, bs, init = ((a[1],), a[2]))
 end
 
-_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)
+_with_ladj_on_mapped(y_with_ladj::Tuple{Vararg{Union{Tuple{Any,Any},NoLogAbsDetJacobian}}}, noladj::NoLogAbsDetJacobian) = noladj
 
-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::AbstractArray{<:Tuple{Any,Any}}, ::NoLogAbsDetJacobian)
     y = _get_all_first(y_with_ladj)
     ladj = _sum_over_second(y_with_ladj)
     (y, ladj)
 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)
-    end
+_with_ladj_on_mapped(::Any, noladj::NoLogAbsDetJacobian) = noladj
+
+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(y_with_ladj, NoLogAbsDetJacobian(mapped_f, X))
 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, NoLogAbsDetJacobian(mapped_f, X))
 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_setladj.jl

@@ -4,11 +4,7 @@ using Test
 using ChangesOfVariables
 using InverseFunctions
 
-const ChangesOfVariablesInverseFunctionsExt = if isdefined(Base, :get_extension)
-    Base.get_extension(ChangesOfVariables, :ChangesOfVariablesInverseFunctionsExt)
-else
-    ChangesOfVariables.ChangesOfVariablesInverseFunctionsExt
-end
+const ChangesOfVariablesInverseFunctionsExt = Base.get_extension(ChangesOfVariables, :ChangesOfVariablesInverseFunctionsExt)
 const InverseFunctionWithLADJ = ChangesOfVariablesInverseFunctionsExt.InverseFunctionWithLADJ
 
 include("getjacobian.jl")

test/test_with_ladj.jl

@@ -12,11 +12,7 @@ include("getjacobian.jl")
 
 
 @testset "with_logabsdet_jacobian" begin
-    @static if VERSION >= v"1.6"
-        _bc_func(f) = Base.Broadcast.BroadcastFunction(f)
-    else
-        _bc_func(f) = Base.Fix1(broadcast, f)
-    end
+    _bc_func(f) = Base.Broadcast.BroadcastFunction(f)
 
     @test with_logabsdet_jacobian(sum, rand(5)) === NoLogAbsDetJacobian(sum, rand(5))
     @test with_logabsdet_jacobian(log ∘ sum, 5.0f0) === NoLogAbsDetJacobian(log ∘ sum, 5.0f0)
@@ -30,7 +26,7 @@ include("getjacobian.jl")
     @test with_logabsdet_jacobian(sin ∘ log, 4.9) === NoLogAbsDetJacobian{typeof(sin ∘ log), Float64}()
     @test with_logabsdet_jacobian(log ∘ sin, 4.9) === NoLogAbsDetJacobian{typeof(log ∘ sin), Float64}()
 
-    @test with_logabsdet_jacobian(Base.Fix1(broadcast, sin), 4.9) === NoLogAbsDetJacobian{typeof(sin), Float64}()
+    @test with_logabsdet_jacobian(Base.Fix1(broadcast, sin), 4.9) === NoLogAbsDetJacobian(Base.Fix1(broadcast, sin), 4.9)
 
     function ChangesOfVariables.with_logabsdet_jacobian(::typeof(foo), x)
         y = foo(x)
@@ -49,18 +45,34 @@ include("getjacobian.jl")
 
     test_with_logabsdet_jacobian(foo, x, getjacobian)
 
-    @static if VERSION >= v"1.6"
-        test_with_logabsdet_jacobian(log ∘ foo, x, getjacobian)
-    end
+    test_with_logabsdet_jacobian(log ∘ foo, x, getjacobian)
 
     @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) === NoLogAbsDetJacobian(f, arg)
+            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.0]) == ([1.0, 2.0], 0.0)
+        @test with_logabsdet_jacobian(_bc_func(bar), (1.0, 2)) === NoLogAbsDetJacobian(_bc_func(bar), (1.0, 2))
+        @test with_logabsdet_jacobian(Base.Fix1(map, bar), (1, 2.0)) === NoLogAbsDetJacobian(Base.Fix1(map, bar), (1, 2.0))
+        @test with_logabsdet_jacobian(_bc_func(bar), Real[1.0, 2]) === NoLogAbsDetJacobian(_bc_func(bar), Real[1.0, 2])
+        @test with_logabsdet_jacobian(Base.Fix1(map, bar), Real[1, 2.0]) === NoLogAbsDetJacobian(Base.Fix1(map, bar), Real[1, 2.0])
+        @test with_logabsdet_jacobian(Base.Fix1(broadcast, bar), Real[1, 2]) === NoLogAbsDetJacobian(Base.Fix1(broadcast, bar), Real[1, 2])
+    end
+
     @testset "with_logabsdet_jacobian on identity, adjoint and transpose" begin
         for f in (identity, adjoint, transpose)
             for arg in (x, A)