From 8797e102bb1896c2ffaf118ac32535f1f5f006d6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Tue, 10 Nov 2026 13:34:02 +0100 Subject: [PATCH 1/4] Improve accuracy and performance of ChainRules derivatives MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - logistic: use e / (1 + e)^2 with e = exp(-|x|) instead of Ω * (1 - Ω), which suffers from cancellation for large x (returns 0 at x = 40) - logabstanh: use 2 / sinh(2x) (one transcendental call instead of two) - log1psq: avoid overflow of x^2 for large |x| - log1mexp, logexpm1: use -inv(expm1(-x)) instead of exp(x - Ω), which amplifies the error of Ω and returns NaN for logexpm1 at Inf - logaddexp, logsubexp: exploit ∂x + ∂y = 1 to compute both partials with a single logistic/expm1 call; fixes NaN for infinite arguments and loss of accuracy for arguments of large magnitude or close to each other - add rules for loglogistic, log1mlogistic, logitexp, logit1mexp - restrict cloglog and cexpexp rules to Real arguments Co-Authored-By: Claude Opus 5.5 --- ext/LogExpFunctionsChainRulesCoreExt.jl | 51 +++++++++++++++++---- test/chainrules.jl | 61 +++++++++++++++++++++++++ 2 files changed, 103 insertions(+), 9 deletions(-) diff --git a/ext/LogExpFunctionsChainRulesCoreExt.jl b/ext/LogExpFunctionsChainRulesCoreExt.jl index 28fe8fe7..f25b986b 100644 --- a/ext/LogExpFunctionsChainRulesCoreExt.jl +++ b/ext/LogExpFunctionsChainRulesCoreExt.jl @@ -114,23 +114,56 @@ function ChainRulesCore.rrule(::typeof(xexpy), x::Real, y::Real) return Ω, xexpy_pullback end -ChainRulesCore.@scalar_rule(logistic(x::Real), (Ω * (1 - Ω),)) +# logistic'(x) = logistic(x) * logistic(-x) = e / (1 + e)^2 with e = exp(-|x|) +# (avoids the cancellation in `1 - logistic(x)` for large `x`) +function _∂_logistic(x::Real) + e = exp(-abs(x)) + return e / (1 + e)^2 +end +# 2x / (1 + x^2), avoiding overflow of `x^2` for large `|x|` +_∂_log1psq(x::Real) = abs(x) < 1 ? 2 * x / (1 + x^2) : 2 / (x + inv(x)) + +ChainRulesCore.@scalar_rule(logistic(x::Real), (_∂_logistic(x),)) ChainRulesCore.@scalar_rule(logit(x::Real), (inv(x * (1 - x)),)) ChainRulesCore.@scalar_rule(logcosh(x::Real), tanh(x)) ChainRulesCore.@scalar_rule(logabssinh(x::Real), coth(x)) -ChainRulesCore.@scalar_rule(logabstanh(x::Real), inv(cosh(x) * sinh(x))) -ChainRulesCore.@scalar_rule(log1psq(x::Real), (2 * x / (1 + x^2),)) +ChainRulesCore.@scalar_rule(logabstanh(x::Real), 2 / sinh(2 * x)) +ChainRulesCore.@scalar_rule(log1psq(x::Real), (_∂_log1psq(x),)) ChainRulesCore.@scalar_rule(log1pexp(x::Real), (logistic(x),)) -ChainRulesCore.@scalar_rule(log1mexp(x::Real), (-exp(x - Ω),)) +ChainRulesCore.@scalar_rule(log1mexp(x::Real), (-inv(expm1(-x)),)) ChainRulesCore.@scalar_rule(log2mexp(x::Real), (-exp(x - Ω),)) -ChainRulesCore.@scalar_rule(logexpm1(x::Real), (exp(x - Ω),)) +ChainRulesCore.@scalar_rule(logexpm1(x::Real), (-inv(expm1(-x)),)) ChainRulesCore.@scalar_rule(log1pmx(x::Real), (-x / (x + 1),)) ChainRulesCore.@scalar_rule(logmxp1(x::Real), ((1 - x) / x,)) +ChainRulesCore.@scalar_rule(loglogistic(x::Real), (logistic(-x),)) +ChainRulesCore.@scalar_rule(log1mlogistic(x::Real), (-logistic(x),)) +ChainRulesCore.@scalar_rule(logitexp(x::Real), (-inv(expm1(x)),)) +ChainRulesCore.@scalar_rule(logit1mexp(x::Real), (inv(expm1(x)),)) -ChainRulesCore.@scalar_rule(logaddexp(x::Real, y::Real), (exp(x - Ω), exp(y - Ω))) +# Both `logaddexp` and `logsubexp` satisfy f(x + c, y + c) = f(x, y) + c, hence ∂x + ∂y = 1. +# The partial derivative `t` with respect to the smaller argument is computed directly, +# the one with respect to the larger argument as `1 - t` (no cancellation since +# 0 ≤ t ≤ 1/2 for `logaddexp` and t ≤ 0 for `logsubexp`). +# The handling of `x == y` mirrors the primal implementations. +function _∂_logaddexp(x::Real, y::Real) + Δ = x == y ? zero(x - y) : abs(x - y) + t = logistic(-Δ) + return x >= y ? (1 - t, t) : (t, 1 - t) +end +function _∂_logsubexp(x::Real, y::Real) + Δ = x == y && (isfinite(x) || x < 0) ? zero(x - y) : abs(x - y) + t = -inv(expm1(Δ)) + return x >= y ? (1 - t, t) : (t, 1 - t) +end +ChainRulesCore.@scalar_rule( + logaddexp(x::Real, y::Real), + @setup((∂x, ∂y) = _∂_logaddexp(x, y)), + (∂x, ∂y), +) ChainRulesCore.@scalar_rule( logsubexp(x::Real, y::Real), - (x > y ? exp(x - Ω) : -exp(x - Ω), x > y ? -exp(y - Ω) : exp(y - Ω)), + @setup((∂x, ∂y) = _∂_logsubexp(x, y)), + (∂x, ∂y), ) function ChainRulesCore.frule((_, Δx), ::typeof(logsumexp), x::AbstractArray{<:Real}; dims=:) @@ -185,7 +218,7 @@ function ChainRulesCore.rrule(::typeof(softmax), x::AbstractArray{<:Real}; dims= return Ω, softmax_pullback end -ChainRulesCore.@scalar_rule(cloglog(x), (-inv((1 - x) * log1p(-x)),)) -ChainRulesCore.@scalar_rule(cexpexp(x), (-xexpx(-exp(x)),)) +ChainRulesCore.@scalar_rule(cloglog(x::Real), (-inv((1 - x) * log1p(-x)),)) +ChainRulesCore.@scalar_rule(cexpexp(x::Real), (-xexpx(-exp(x)),)) end # module diff --git a/test/chainrules.jl b/test/chainrules.jl index c19ae936..8011b525 100644 --- a/test/chainrules.jl +++ b/test/chainrules.jl @@ -48,6 +48,11 @@ test_frule(log1psq, x) test_rrule(log1psq, x) end + # no overflow of `x^2` + for x in (-1e200, 1e200) + @test last(frule((NoTangent(), 1.0), log1psq, x)) ≈ 2 / x rtol=4*eps() + @test last(last(rrule(log1psq, x))(1.0)) ≈ 2 / x rtol=4*eps() + end # test all `Float64` and `Float32` branches of `logistic` for x in (-821.4, -23.5, 12.3, 41.2) @@ -58,6 +63,12 @@ test_frule(logistic, x; rtol=1f-3, atol=1f-3) test_rrule(logistic, x; rtol=1f-3, atol=1f-3) end + # derivative is accurate for large `|x|` in a relative sense (no cancellation) + for x in (40.0, -40.0, 30f0, -30f0) + ∂x = exp(-abs(x)) + @test last(frule((NoTangent(), one(x)), logistic, x)) ≈ ∂x rtol=4*eps(typeof(x)) + @test last(last(rrule(logistic, x))(one(x))) ≈ ∂x rtol=4*eps(typeof(x)) + end for x in (-randexp(), randexp()) test_frule(logcosh, x) @@ -94,6 +105,25 @@ test_frule(logexpm1, x; rtol=1f-3, atol=1f-3) test_rrule(logexpm1, x; rtol=1f-3, atol=1f-3) end + # special case (manually since FiniteDifferences/ChainRulesTestUtils fails at Inf) + @test frule((NoTangent(), 1.0), logexpm1, Inf) === (Inf, 1.0) + @test last(last(rrule(logexpm1, Inf))(1.0)) === 1.0 + + # accuracy for small `|x|` (where `Ω` is large in magnitude) + for (f, x) in ((log1mexp, -1e-12), (logexpm1, 6.039463507942533e-8)) + ∂x = -inv(expm1(-big(x))) + @test last(frule((NoTangent(), 1.0), f, x)) ≈ ∂x rtol=4*eps() + @test last(last(rrule(f, x))(1.0)) ≈ ∂x rtol=4*eps() + end + + for x in (-randexp(), randexp()) + test_scalar(loglogistic, x) + test_scalar(log1mlogistic, x) + end + for x in (-10.2, -3.3, -0.3) + test_scalar(logitexp, x) + test_scalar(logit1mexp, x) + end test_scalar(log1pmx, rand()) @@ -106,6 +136,37 @@ test_frule(logsubexp, x, y) test_rrule(logsubexp, x, y) end + # nearly equal arguments (manually since finite differences are inaccurate close to + # the singularity at `x == y`) + for (x, y) in ((1.0, 1.0 + 1e-6), (1.0 + 1e-6, 1.0)) + d = big(x) - big(y) + ∂x, ∂y = -inv(expm1(-d)), -inv(expm1(d)) + @test last(frule((NoTangent(), 1.0, 0.0), logsubexp, x, y)) ≈ ∂x rtol=4*eps() + @test last(frule((NoTangent(), 0.0, 1.0), logsubexp, x, y)) ≈ ∂y rtol=4*eps() + _, back = rrule(logsubexp, x, y) + _, x̄, ȳ = back(1.0) + @test x̄ ≈ ∂x rtol=4*eps() + @test ȳ ≈ ∂y rtol=4*eps() + end + # accuracy for arguments of large magnitude + let (x, y) = (300.0, 301.0) + d = big(x) - big(y) + ∂x, ∂y = logistic(d), logistic(-d) + _, back = rrule(logaddexp, x, y) + _, x̄, ȳ = back(1.0) + @test x̄ ≈ ∂x rtol=4*eps() + @test ȳ ≈ ∂y rtol=4*eps() + end + # special cases (manually since FiniteDifferences/ChainRulesTestUtils fails at Inf) + for (f, ∂x, ∂y) in ((logaddexp, 1.0, 0.0), (logsubexp, 1.0, -0.0)) + Ω, back = rrule(f, Inf, 1.0) + @test Ω === Inf + @test back(1.0) === (NoTangent(), ∂x, ∂y) + Ω, back = rrule(f, 1.0, Inf) + @test Ω === Inf + @test back(1.0) === (NoTangent(), ∂y, ∂x) + @test last(frule((NoTangent(), 1.0, 0.0), f, Inf, 1.0)) === ∂x + end for x in (randn(10), randn(10, 8)), dims in (:, 1, 1:2, 2) dims isa Colon || all(d <= ndims(x) for d in dims) || continue From be04233b8cf718743d44fb418af1ca37a406c586 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Tue, 10 Nov 2026 13:38:43 +0100 Subject: [PATCH 2/4] Bump version 1.0.5 --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 9ca4f16e..b7f781c3 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "LogExpFunctions" uuid = "2ab3a3ac-af41-5b50-aa03-7779005ae688" authors = ["StatsFun.jl contributors, Tamas K. Papp "] -version = "1.0.3" +version = "1.0.5" [deps] ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" From ad431f530ce4d9538c23992b7a8e54e2c17d7887 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Tue, 10 Nov 2026 23:10:41 +0100 Subject: [PATCH 3/4] Narrow `x == y` guards of `logaddexp`/`logsubexp` derivatives and test equal arguments MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace `x - y` with zero only if it is `NaN` (`x == y == ±Inf` for `logaddexp`, `x == y == -Inf` for `logsubexp`): for finite `x == y` it is already zero, and the replacement would drop higher-order derivatives in nested AD. Add tests for equal and `NaN` arguments (#61), compare the `logistic` derivative with an exact reference, and shorten comments. Co-Authored-By: Claude Opus 5.5 --- ext/LogExpFunctionsChainRulesCoreExt.jl | 17 +++++++---------- test/chainrules.jl | 23 +++++++++++++++++++++-- 2 files changed, 28 insertions(+), 12 deletions(-) diff --git a/ext/LogExpFunctionsChainRulesCoreExt.jl b/ext/LogExpFunctionsChainRulesCoreExt.jl index f25b986b..99ce05b2 100644 --- a/ext/LogExpFunctionsChainRulesCoreExt.jl +++ b/ext/LogExpFunctionsChainRulesCoreExt.jl @@ -114,13 +114,12 @@ function ChainRulesCore.rrule(::typeof(xexpy), x::Real, y::Real) return Ω, xexpy_pullback end -# logistic'(x) = logistic(x) * logistic(-x) = e / (1 + e)^2 with e = exp(-|x|) -# (avoids the cancellation in `1 - logistic(x)` for large `x`) +# logistic(x) * logistic(-x), avoiding cancellation in `1 - logistic(x)` function _∂_logistic(x::Real) e = exp(-abs(x)) return e / (1 + e)^2 end -# 2x / (1 + x^2), avoiding overflow of `x^2` for large `|x|` +# avoid overflow of `x^2` _∂_log1psq(x::Real) = abs(x) < 1 ? 2 * x / (1 + x^2) : 2 / (x + inv(x)) ChainRulesCore.@scalar_rule(logistic(x::Real), (_∂_logistic(x),)) @@ -140,18 +139,16 @@ ChainRulesCore.@scalar_rule(log1mlogistic(x::Real), (-logistic(x),)) ChainRulesCore.@scalar_rule(logitexp(x::Real), (-inv(expm1(x)),)) ChainRulesCore.@scalar_rule(logit1mexp(x::Real), (inv(expm1(x)),)) -# Both `logaddexp` and `logsubexp` satisfy f(x + c, y + c) = f(x, y) + c, hence ∂x + ∂y = 1. -# The partial derivative `t` with respect to the smaller argument is computed directly, -# the one with respect to the larger argument as `1 - t` (no cancellation since -# 0 ≤ t ≤ 1/2 for `logaddexp` and t ≤ 0 for `logsubexp`). -# The handling of `x == y` mirrors the primal implementations. +# ∂x + ∂y = 1: `t` is the partial derivative with respect to the smaller argument +# (0 ≤ t ≤ 1/2 for `logaddexp` and t ≤ 0 for `logsubexp`, so `1 - t` is accurate). +# `x - y` is `NaN` if `x == y == ±Inf`. function _∂_logaddexp(x::Real, y::Real) - Δ = x == y ? zero(x - y) : abs(x - y) + Δ = x == y && isinf(x) ? zero(x - y) : abs(x - y) t = logistic(-Δ) return x >= y ? (1 - t, t) : (t, 1 - t) end function _∂_logsubexp(x::Real, y::Real) - Δ = x == y && (isfinite(x) || x < 0) ? zero(x - y) : abs(x - y) + Δ = x == y == -Inf ? zero(x - y) : abs(x - y) t = -inv(expm1(Δ)) return x >= y ? (1 - t, t) : (t, 1 - t) end diff --git a/test/chainrules.jl b/test/chainrules.jl index 8011b525..9ee20598 100644 --- a/test/chainrules.jl +++ b/test/chainrules.jl @@ -63,9 +63,10 @@ test_frule(logistic, x; rtol=1f-3, atol=1f-3) test_rrule(logistic, x; rtol=1f-3, atol=1f-3) end - # derivative is accurate for large `|x|` in a relative sense (no cancellation) + # accuracy for large `|x|` for x in (40.0, -40.0, 30f0, -30f0) - ∂x = exp(-abs(x)) + e = exp(-abs(big(x))) + ∂x = e / (1 + e)^2 @test last(frule((NoTangent(), one(x)), logistic, x)) ≈ ∂x rtol=4*eps(typeof(x)) @test last(last(rrule(logistic, x))(one(x))) ≈ ∂x rtol=4*eps(typeof(x)) end @@ -167,6 +168,24 @@ @test back(1.0) === (NoTangent(), ∂y, ∂x) @test last(frule((NoTangent(), 1.0, 0.0), f, Inf, 1.0)) === ∂x end + # equal and `NaN` arguments (`x == y == ±Inf`: #61) + for T in (Float32, Float64) + for x in (T(-Inf), one(T), T(Inf)) + _, back = rrule(logaddexp, x, x) + @test back(one(T)) === (NoTangent(), T(0.5), T(0.5)) + end + # limits as `x` approaches `y` from above + for x in (T(-Inf), one(T)) + _, back = rrule(logsubexp, x, x) + @test back(one(T)) === (NoTangent(), T(Inf), T(-Inf)) + end + _, back = rrule(logsubexp, T(Inf), T(Inf)) + @test all(isnan, back(one(T))[2:3]) + for f in (logaddexp, logsubexp) + _, back = rrule(f, T(NaN), one(T)) + @test all(isnan, back(one(T))[2:3]) + end + end for x in (randn(10), randn(10, 8)), dims in (:, 1, 1:2, 2) dims isa Colon || all(d <= ndims(x) for d in dims) || continue From 3b06dac36d539a41e7fc43ef91c4e42c76f92346 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Wed, 11 Nov 2026 10:21:07 +0100 Subject: [PATCH 4/4] Restore comments MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: David Müller-Widmann --- ext/LogExpFunctionsChainRulesCoreExt.jl | 6 ++++-- test/chainrules.jl | 2 +- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/ext/LogExpFunctionsChainRulesCoreExt.jl b/ext/LogExpFunctionsChainRulesCoreExt.jl index 99ce05b2..b8fc2bd3 100644 --- a/ext/LogExpFunctionsChainRulesCoreExt.jl +++ b/ext/LogExpFunctionsChainRulesCoreExt.jl @@ -119,7 +119,7 @@ function _∂_logistic(x::Real) e = exp(-abs(x)) return e / (1 + e)^2 end -# avoid overflow of `x^2` +# 2x / (1 + x^2), avoiding overflow of `x^2` for large `|x|` _∂_log1psq(x::Real) = abs(x) < 1 ? 2 * x / (1 + x^2) : 2 / (x + inv(x)) ChainRulesCore.@scalar_rule(logistic(x::Real), (_∂_logistic(x),)) @@ -139,7 +139,9 @@ ChainRulesCore.@scalar_rule(log1mlogistic(x::Real), (-logistic(x),)) ChainRulesCore.@scalar_rule(logitexp(x::Real), (-inv(expm1(x)),)) ChainRulesCore.@scalar_rule(logit1mexp(x::Real), (inv(expm1(x)),)) -# ∂x + ∂y = 1: `t` is the partial derivative with respect to the smaller argument +# the one with respect to the larger argument as `1 - t` (no cancellation since +# 0 ≤ t ≤ 1/2 for `logaddexp` and t ≤ 0 for `logsubexp`). +# The handling of `x == y` mirrors the primal implementations. # (0 ≤ t ≤ 1/2 for `logaddexp` and t ≤ 0 for `logsubexp`, so `1 - t` is accurate). # `x - y` is `NaN` if `x == y == ±Inf`. function _∂_logaddexp(x::Real, y::Real) diff --git a/test/chainrules.jl b/test/chainrules.jl index 9ee20598..101d8855 100644 --- a/test/chainrules.jl +++ b/test/chainrules.jl @@ -63,7 +63,7 @@ test_frule(logistic, x; rtol=1f-3, atol=1f-3) test_rrule(logistic, x; rtol=1f-3, atol=1f-3) end - # accuracy for large `|x|` + # derivative is accurate for large `|x|` in a relative sense (no cancellation) for x in (40.0, -40.0, 30f0, -30f0) e = exp(-abs(big(x))) ∂x = e / (1 + e)^2