Keep derivatives and imaginary parts at ties in logsumexp

.github/dependabot.yml

@@ -10,6 +10,7 @@ updates:
       - "/"
       - "/docs"
       - "/test"
+      - "/test/ad"
     schedule:
       interval: "daily"
     groups:

.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

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

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"

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

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]