From e1e77e75334fd915a63818319139531d35f2389f Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 22:52:28 +0200 Subject: [PATCH 1/2] Fix with_logabsdet_jacobian for mapped/broadcasted functions Handle NoLogAbsDetJacobian properly. --- src/with_ladj.jl | 35 ++++++++++++++++++++++++----------- test/getjacobian.jl | 9 +++++++++ test/test_with_ladj.jl | 16 +++++++++++++++- 3 files changed, 48 insertions(+), 12 deletions(-) diff --git a/src/with_ladj.jl b/src/with_ladj.jl index 2fcb4ce..979a49e 100644 --- a/src/with_ladj.jl +++ b/src/with_ladj.jl @@ -117,29 +117,42 @@ 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 + +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 + @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 @@ -147,7 +160,7 @@ function with_logabsdet_jacobian(mapped_f::Base.Fix1{<:Union{typeof(map),typeof( 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 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_with_ladj.jl b/test/test_with_ladj.jl index ff307f7..b52428d 100644 --- a/test/test_with_ladj.jl +++ b/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) From 393df8d52cf74601dd6282929d2f149ee8963487 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Wed, 8 Jul 2026 22:09:56 +0200 Subject: [PATCH 2/2] Increase package version to 0.1.11 --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 37a0a82..948fa0e 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"