logsumexp gives the wrong derivative at exact ties under Enzyme and Mooncake #128

Enzyme and Mooncake give the wrong derivative for logsumexp when the two terms being added are equal:

using Enzyme, ForwardDiff, LogExpFunctions

# Both terms equal 1 at t = 1, with derivatives 1 and 2, so the answer is 1.5
f_tuple(t) = logsumexp((t, 2t - 1))
f_vector(t) = logsumexp([t, 2t - 1])
f_generator(t) = logsumexp(x for x in (t, 2t - 1))
f_plain(t) = log(exp(t) + exp(2t - 1))

for (name, f) in (("tuple", f_tuple), ("vector", f_vector), ("generator", f_generator), ("plain", f_plain))
    println(rpad(name, 10),
        " Enzyme.Forward=", only(autodiff(Forward, f, Duplicated(1.0, 1.0))),
        "  Enzyme.Reverse=", only(only(autodiff(Reverse, f, Active, Active(1.0)))),
        "  ForwardDiff=", ForwardDiff.derivative(f, 1.0))
end

tuple      Enzyme.Forward=1.0  Enzyme.Reverse=1.0  ForwardDiff=1.5
vector     Enzyme.Forward=1.0  Enzyme.Reverse=1.0  ForwardDiff=1.5
generator  Enzyme.Forward=2.0  Enzyme.Reverse=2.0  ForwardDiff=1.5
plain      Enzyme.Forward=1.5  Enzyme.Reverse=1.5  ForwardDiff=1.5

Julia 1.13.0, LogExpFunctions 0.3.29, Enzyme 0.13.204, ForwardDiff 1.4.5.

Mooncake 0.5.60 (value_and_gradient!!) gives 1.0 for the tuple and 2.0 for the generator. The vector and plain forms give 1.5.

The three _logsumexp_onepass_op methods in src/logsumexp.jl handle a tie by replacing the difference with a constant:

x2, zero(x1 - x2)
xmax, r + exp(zero(x - xmax))
xmax2, r2 + (r1 + one(r1)) * exp(zero(xmax1 - xmax2))

The value is unchanged, but the derivative of the difference is lost. ForwardDiff never takes this branch: since #60 (the fix for #59) it runs only when neither > nor < holds, and Dual comparison breaks ties using the partials.

The constant is there so that x1 = x2 = ±Inf gives 0 rather than Inf - Inf = NaN. Replacing it only when the difference is NaN handles both cases:

d = x1 - x2
x2, ifelse(isnan(d), zero(d), d)

With the same change in all three methods, every case above gives 1.5 in both Enzyme modes. logsumexp((Inf, Inf)), logsumexp((-Inf, -Inf)) and logsumexp([Inf, Inf, 1.0]) still return Inf, -Inf and Inf.