diff --git a/.github/dependabot.yml b/.github/dependabot.yml index c5cd4b68..6ebc5818 100644 --- a/.github/dependabot.yml +++ b/.github/dependabot.yml @@ -10,6 +10,7 @@ updates: - "/" - "/docs" - "/test" + - "/test/ad" schedule: interval: "daily" groups: diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 03c8f641..7e6da014 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -40,6 +40,27 @@ jobs: files: lcov.info token: ${{ secrets.CODECOV_TOKEN }} fail_ci_if_error: true + ad: + name: AD - Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ github.event_name }} + runs-on: ${{ matrix.os }} + strategy: + fail-fast: false + matrix: + version: + - '1' + os: + - ubuntu-latest + steps: + - uses: actions/checkout@v7 + - uses: julia-actions/setup-julia@v3 + with: + version: ${{ matrix.version }} + - uses: julia-actions/cache@v3 + - run: | + julia --project=test/ad -e ' + using Pkg + Pkg.instantiate()' + - run: julia --project=test/ad --check-bounds=yes test/ad/runtests.jl docs: name: Documentation runs-on: ubuntu-latest diff --git a/src/logsumexp.jl b/src/logsumexp.jl index 0c655f4e..7c3eb454 100644 --- a/src/logsumexp.jl +++ b/src/logsumexp.jl @@ -98,7 +98,9 @@ function _logsumexp_onepass_op(x1::T, x2::T) where {T<:Number} else # handle `x1 = x2 = ±Inf` correctly # checking inequalities above instead of equality fixes issue #59 - x2, zero(x1 - x2) + # replacing only `NaN` keeps derivatives and imaginary parts at ties (#128) + d = x1 - x2 + x2, isnan(d) ? zero(d) : d end end r = exp(a) @@ -160,7 +162,9 @@ function _logsumexp_onepass_op(xmax1::T, xmax2::T, r1::R, r2::R) where {T<:Numbe else # handle `xmax1 = xmax2 = ±Inf` correctly # checking inequalities above instead of equality fixes issue #59 - xmax2, r2 + (r1 + one(r1)) * exp(zero(xmax1 - xmax2)) + # replacing only `NaN` keeps derivatives and imaginary parts at ties (#128) + d = xmax1 - xmax2 + xmax2, r2 + (r1 + one(r1)) * exp(isnan(d) ? zero(d) : d) end end return xmax, r diff --git a/test/ad/Project.toml b/test/ad/Project.toml new file mode 100644 index 00000000..eee02c61 --- /dev/null +++ b/test/ad/Project.toml @@ -0,0 +1,14 @@ +[deps] +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" +LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688" +Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" + +[sources] +LogExpFunctions = {path = "../.."} + +[compat] +Enzyme = "0.13" +ForwardDiff = "1" +Mooncake = "0.5" diff --git a/test/ad/runtests.jl b/test/ad/runtests.jl new file mode 100644 index 00000000..0b660ea3 --- /dev/null +++ b/test/ad/runtests.jl @@ -0,0 +1,42 @@ +using LogExpFunctions +using Enzyme +using ForwardDiff +using Mooncake +using Test + +# issue #128 +@testset "logsumexp at ties" begin + # `t` and `2t - 0.37` are exactly equal at `t = 0.37` + x0 = 0.37 + dref = 1.5 + fs = ( + tuple = t -> logsumexp((t, 2t - 0.37)), + vector = t -> logsumexp([t, 2t - 0.37]), + generator = t -> logsumexp(x for x in (t, 2t - 0.37)), + dims = t -> logsumexp([t 2t - 0.37]; dims=2)[1], + # abstract eltype and > 1024 elements: combines partial sums + abstract = t -> logsumexp(Number[fill(t, 1024); fill(2t - 0.37, 1024)]), + ) + @testset "$name" for (name, f) in pairs(fs) + y = f(x0) + + @test ForwardDiff.derivative(f, x0) ≈ dref + + df, val = autodiff(ForwardWithPrimal, f, Duplicated(x0, 1.0)) + @test val ≈ y + @test df ≈ dref + (df,), val = autodiff(ReverseWithPrimal, f, Active, Active(x0)) + @test val ≈ y + @test df ≈ dref + + @testset "Mooncake $mode" for (mode, prepare) in pairs(( + forward = Mooncake.prepare_derivative_cache, + reverse = Mooncake.prepare_gradient_cache, + )) + cache = prepare(f, x0) + val, (_, df) = Mooncake.value_and_gradient!!(cache, f, x0) + @test val ≈ y + @test df ≈ dref + end + end +end diff --git a/test/basicfuns.jl b/test/basicfuns.jl index cc5bccd2..f36c06f2 100644 --- a/test/basicfuns.jl +++ b/test/basicfuns.jl @@ -373,6 +373,10 @@ end @test isnan(logsumexp!(Complex{Float64}[1.0], Complex{Float64}[NaN * im, 9.0])[1]) @test isnan(logsumexp!(Complex{Float64}[1.0], Complex{Float64}[NaN * im, Inf])[1]) @test isnan(logsumexp!(Complex{Float64}[1.0], Complex{Float64}[NaN * im, -Inf])[1]) + @test isnan(logsumexp(x for x in (NaN, 9.0))) + @test isnan(logsumexp(x for x in (Inf, NaN))) + @test isnan(logsumexp(x for x in (NaN, NaN))) + @test isnan(logsumexp(x for x in (NaN * im, 9.0))) # logsumexp with general iterables (issue #63) xs = range(-500, stop = 10, length = 1000) @@ -388,6 +392,22 @@ end @test @inferred(logsumexp(xs; dims=[1, 2])) ≈ log(sum(exp.(xs); dims=[1, 2])) @test @inferred(logsumexp(x for x in xs)) == logsumexp(xs) + # issue #128 + @testset "complex ties" begin + @test logsumexp(z for z in (0.0, im)) ≈ log(1 + exp(im)) + # abstract eltype and > 1024 elements: combines partial sums + zs = Number[fill(0.0 + 0.0im, 1024); fill(im, 1024)] + @test logsumexp(zs) ≈ log(1024) + log(1 + exp(im)) + end + + @testset "infinite ties" begin + for x in (Inf, -Inf) + @test logsumexp(y for y in (x, x)) ≡ x + @test logsumexp(y for y in (x, x, 1.0)) ≡ max(x, 1.0) + @test logsumexp(Number[fill(x, 1024); fill(x, 1024)]) ≡ x + end + end + # output arrays with abstract eltype xs = randn(2, 4) out = [missing, 1.0]