diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 4bcf789..e0b9221 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -24,6 +24,7 @@ jobs: - 'min' - 'lts' - '1' + - 'pre' os: - ubuntu-latest include: diff --git a/Project.toml b/Project.toml index 37a0a82..d3fecb5 100644 --- a/Project.toml +++ b/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" diff --git a/src/ChangesOfVariables.jl b/src/ChangesOfVariables.jl index db62a89..4951619 100644 --- a/src/ChangesOfVariables.jl +++ b/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 diff --git a/src/setladj.jl b/src/setladj.jl index 377c8e7..e4f4fc4 100644 --- a/src/setladj.jl +++ b/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 diff --git a/src/with_ladj.jl b/src/with_ladj.jl index 2fcb4ce..8402819 100644 --- a/src/with_ladj.jl +++ b/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 diff --git a/test/getjacobian.jl b/test/getjacobian.jl index ff4c480..718d339 100644 --- a/test/getjacobian.jl +++ b/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) diff --git a/test/test_setladj.jl b/test/test_setladj.jl index 02a6bd8..d867bd1 100644 --- a/test/test_setladj.jl +++ b/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") diff --git a/test/test_with_ladj.jl b/test/test_with_ladj.jl index ff307f7..b6caeac 100644 --- a/test/test_with_ladj.jl +++ b/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)