From f361cae0572578bfcd2628f519adfc55590c8c44 Mon Sep 17 00:00:00 2001
From: cossio <cossio@users.noreply.github.com>
Date: Wed, 29 Jul 2026 12:34:05 +0200
Subject: [PATCH 1/8] Compute logvarexp/logstdexp gradients in the log domain
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Post-merge review of #3 flagged that squaring exp(xᵢ) - m outside the
log domain can underflow for nearly equal entries, returning Inf/NaN
gradients where the true gradient is representable (e.g.
x = [-1e-200, 1e-200] has gradient ±1e200). Rewrite _∂x_logvarexp with
logsubexp/logsumexp so the ratio is formed in the log domain.

Also unthunk the incoming cotangent in the shared pullback, and
document that the rules cover real arrays only (complex arrays,
supported by logmeanexp itself, are not covered).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01KmceQuDSTogkuoRwSd48aj
---
 CHANGELOG.md                             |  4 ++--
 docs/src/index.md                        |  5 +++--
 ext/LogStatFunctionsChainRulesCoreExt.jl | 16 ++++++++++------
 test/Project.toml                        |  2 ++
 test/chainrules.jl                       | 19 ++++++++++++++++++-
 5 files changed, 35 insertions(+), 11 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 16650fe..2b07e06 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,8 +5,8 @@ All notable changes to this project will be documented in this file. The format
 ## Unreleased
 
 - Add a ChainRulesCore extension defining `frule`s and `rrule`s for
-  `logmeanexp`, `logvarexp` and `logstdexp`, making them differentiable with
-  ChainRules-based AD packages (e.g. Zygote).
+  `logmeanexp`, `logvarexp` and `logstdexp` on real arrays, making them
+  differentiable with ChainRules-based AD packages (e.g. Zygote).
 - Support immutable arrays (e.g. `StaticArrays`) and complex arrays.
 - Results preserve the input eltype (e.g. `Float32` in → `Float32` out), and
   the `log(N)` normalization is computed in the result's precision (so e.g.
diff --git a/docs/src/index.md b/docs/src/index.md
index 874a37d..3a0acf0 100644
--- a/docs/src/index.md
+++ b/docs/src/index.md
@@ -30,8 +30,9 @@ logstdexp(A; corrected=false)
 
 All three functions have
 [ChainRules](https://github.com/JuliaDiff/ChainRulesCore.jl) derivative rules
-(loaded automatically when ChainRulesCore is in the environment), so they can be
-differentiated with ChainRules-based AD packages such as Zygote.
+for real arrays (loaded automatically when ChainRulesCore is in the
+environment), so they can be differentiated with ChainRules-based AD packages
+such as Zygote. Complex arrays are not covered by these rules.
 
 ## Reference
 
diff --git a/ext/LogStatFunctionsChainRulesCoreExt.jl b/ext/LogStatFunctionsChainRulesCoreExt.jl
index 98d1e03..df7b20e 100644
--- a/ext/LogStatFunctionsChainRulesCoreExt.jl
+++ b/ext/LogStatFunctionsChainRulesCoreExt.jl
@@ -1,6 +1,7 @@
 module LogStatFunctionsChainRulesCoreExt
 
 using LogStatFunctions: logmeanexp, logvarexp, logstdexp
+using LogExpFunctions: logsumexp, logsubexp
 import ChainRulesCore
 
 function ChainRulesCore.frule((_, Δx), ::typeof(logmeanexp), x::AbstractArray{<:Real}; dims = :)
@@ -52,19 +53,22 @@ end
 
 # ∂/∂xⱼ log(var(exp.(x))) = 2 exp(xⱼ) (exp(xⱼ) - m) / Σᵢ (exp(xᵢ) - m)², with m = exp(logmean).
 # The m-dependence on x drops out because Σᵢ (exp(xᵢ) - m) = 0, and `corrected` only shifts
-# the result by a constant, so the gradient is the same either way.
+# the result by a constant, so the gradient is the same either way. Computed in the log
+# domain: squaring exp(xᵢ) - m directly can under/overflow even when the gradient itself is
+# representable (e.g. nearly equal entries around zero).
 function _∂x_logvarexp(x::AbstractArray{<:Real}, logmean, dims)
-    d = x .- logmean
-    e = expm1.(d)
-    return (2 .* exp.(d) .* e) ./ sum(abs2, e; dims)
+    l = logsubexp.(x, logmean)
+    S = logsumexp(2 .* l; dims)
+    return sign.(x .- logmean) .* 2 .* exp.(x .+ l .- S)
 end
 
 function _∂x_pullback(∂x, x)
     project_x = ChainRulesCore.ProjectTo(x)
     function pullback(Ω̄)
+        ΔΩ = ChainRulesCore.unthunk(Ω̄)
         x̄ = ChainRulesCore.InplaceableThunk(
-            Δ -> Δ .+= Ω̄ .* ∂x,
-            ChainRulesCore.@thunk(project_x(Ω̄ .* ∂x)),
+            Δ -> Δ .+= ΔΩ .* ∂x,
+            ChainRulesCore.@thunk(project_x(ΔΩ .* ∂x)),
         )
         return ChainRulesCore.NoTangent(), x̄
     end
diff --git a/test/Project.toml b/test/Project.toml
index f0d671d..3d5d33f 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -1,5 +1,6 @@
 [deps]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
 ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7"
 LogStatFunctions = "1b3c2205-da69-4893-b2df-717075b12251"
@@ -9,6 +10,7 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]
 Aqua = "0.8"
+ChainRulesCore = "1"
 ChainRulesTestUtils = "1"
 ExplicitImports = "1.15"
 StaticArrays = "1"
diff --git a/test/chainrules.jl b/test/chainrules.jl
index be46b76..30bc03b 100644
--- a/test/chainrules.jl
+++ b/test/chainrules.jl
@@ -1,5 +1,6 @@
-using Test: @testset
+using Test: @test, @testset
 using ChainRulesTestUtils: test_frule, test_rrule
+using ChainRulesCore: NoTangent, frule, rrule, unthunk
 using LogStatFunctions: logmeanexp, logvarexp, logstdexp
 
 @testset "chainrules" begin
@@ -15,3 +16,19 @@ using LogStatFunctions: logmeanexp, logvarexp, logstdexp
         end
     end
 end
+
+@testset "chainrules extreme values" begin
+    # Nearly equal entries: the primal and its gradient are representable, but squaring
+    # exp(xᵢ) - m outside the log domain underflows and used to yield Inf/NaN gradients.
+    x = [-1.0e-200, 1.0e-200]
+    for (f, h) in ((logvarexp, 1), (logstdexp, 2))
+        Ω, pb = rrule(f, x)
+        @test isfinite(Ω)
+        x̄ = unthunk(pb(1.0)[2])
+        @test all(isfinite, x̄)
+        @test x̄[1] ≈ -1.0e200 / h
+        Ω, ΔΩ = frule((NoTangent(), [1.0, 0.0]), f, x)
+        @test isfinite(Ω)
+        @test ΔΩ ≈ -1.0e200 / h
+    end
+end

From 00ba70d4719a081d384d2ac7871ae4ccd3d5d49d Mon Sep 17 00:00:00 2001
From: cossio <cossio@users.noreply.github.com>
Date: Wed, 29 Jul 2026 12:50:15 +0200
Subject: [PATCH 2/8] Make ChainRules gradients invariant to large common
 offsets

Codex review flagged that x .+ l .- S carries the input's common offset
through each term, losing low bits at large magnitudes. Centering on
logmean is not enough: logmean itself carries an ulp(offset)-scale
rounding that shifts all centered values, and against a BigFloat
reference both formulations lose all accuracy at Float64 offsets ~1e15
(and ~27% error at Float32 offsets ~1e5).

Instead center on maximum(x; dims), which subtracts exactly, and form
the log-mean from the centered values in O(1) arithmetic. The log(n)
term is folded into lm before subtracting from t so tiny spreads near
zero are not absorbed by O(1) intermediates. Apply the same idea to
the logmeanexp gradient by normalizing max-centered exponentials
(softmax) instead of dividing exp(x - logmean) by n.

Regression tests compare gradients at offsets 1e12 and 1e15 (chosen so
the shift is exactly representable) against the unshifted gradients.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01KmceQuDSTogkuoRwSd48aj
---
 ext/LogStatFunctionsChainRulesCoreExt.jl | 45 ++++++++++++++++--------
 test/chainrules.jl                       | 15 ++++++++
 2 files changed, 45 insertions(+), 15 deletions(-)

diff --git a/ext/LogStatFunctionsChainRulesCoreExt.jl b/ext/LogStatFunctionsChainRulesCoreExt.jl
index df7b20e..8410148 100644
--- a/ext/LogStatFunctionsChainRulesCoreExt.jl
+++ b/ext/LogStatFunctionsChainRulesCoreExt.jl
@@ -1,20 +1,18 @@
 module LogStatFunctionsChainRulesCoreExt
 
 using LogStatFunctions: logmeanexp, logvarexp, logstdexp
-using LogExpFunctions: logsumexp, logsubexp
+using LogExpFunctions: logsumexp
 import ChainRulesCore
 
 function ChainRulesCore.frule((_, Δx), ::typeof(logmeanexp), x::AbstractArray{<:Real}; dims = :)
     Ω = logmeanexp(x; dims)
-    n = length(x) ÷ length(Ω)
-    ΔΩ = sum(exp.(x .- Ω) .* Δx; dims) ./ n
+    ΔΩ = sum(_softmax(x, dims) .* Δx; dims)
     return Ω, ΔΩ
 end
 
 function ChainRulesCore.rrule(::typeof(logmeanexp), x::AbstractArray{<:Real}; dims = :)
     Ω = logmeanexp(x; dims)
-    n = length(x) ÷ length(Ω)
-    return Ω, _∂x_pullback(exp.(x .- Ω) ./ n, x)
+    return Ω, _∂x_pullback(_softmax(x, dims), x)
 end
 
 function ChainRulesCore.frule(
@@ -22,7 +20,7 @@ function ChainRulesCore.frule(
         dims = :, corrected::Bool = true, logmean = logmeanexp(x; dims)
     )
     Ω = logvarexp(x; dims, corrected, logmean)
-    ΔΩ = sum(_∂x_logvarexp(x, logmean, dims) .* Δx; dims)
+    ΔΩ = sum(_∂x_logvarexp(x, dims) .* Δx; dims)
     return Ω, ΔΩ
 end
 
@@ -31,7 +29,7 @@ function ChainRulesCore.rrule(
         dims = :, corrected::Bool = true, logmean = logmeanexp(x; dims)
     )
     Ω = logvarexp(x; dims, corrected, logmean)
-    return Ω, _∂x_pullback(_∂x_logvarexp(x, logmean, dims), x)
+    return Ω, _∂x_pullback(_∂x_logvarexp(x, dims), x)
 end
 
 function ChainRulesCore.frule(
@@ -39,7 +37,7 @@ function ChainRulesCore.frule(
         dims = :, corrected::Bool = true, logmean = logmeanexp(x; dims)
     )
     Ω = logstdexp(x; dims, corrected, logmean)
-    ΔΩ = sum(_∂x_logvarexp(x, logmean, dims) ./ 2 .* Δx; dims)
+    ΔΩ = sum(_∂x_logvarexp(x, dims) ./ 2 .* Δx; dims)
     return Ω, ΔΩ
 end
 
@@ -48,18 +46,35 @@ function ChainRulesCore.rrule(
         dims = :, corrected::Bool = true, logmean = logmeanexp(x; dims)
     )
     Ω = logstdexp(x; dims, corrected, logmean)
-    return Ω, _∂x_pullback(_∂x_logvarexp(x, logmean, dims) / 2, x)
+    return Ω, _∂x_pullback(_∂x_logvarexp(x, dims) / 2, x)
+end
+
+# ∂/∂xⱼ log(mean(exp.(x))) = exp(xⱼ) / Σᵢ exp(xᵢ), i.e. softmax(x). Normalizing by the
+# actual sum of the max-centered exponentials (rather than dividing exp(x - logmean) by n)
+# makes a large common offset in x cancel exactly instead of leaking ulp-level errors of
+# the offset's magnitude into the gradient.
+function _softmax(x::AbstractArray{<:Real}, dims)
+    y = exp.(x .- maximum(x; dims))
+    return y ./ sum(y; dims)
 end
 
 # ∂/∂xⱼ log(var(exp.(x))) = 2 exp(xⱼ) (exp(xⱼ) - m) / Σᵢ (exp(xᵢ) - m)², with m = exp(logmean).
 # The m-dependence on x drops out because Σᵢ (exp(xᵢ) - m) = 0, and `corrected` only shifts
-# the result by a constant, so the gradient is the same either way. Computed in the log
-# domain: squaring exp(xᵢ) - m directly can under/overflow even when the gradient itself is
-# representable (e.g. nearly equal entries around zero).
-function _∂x_logvarexp(x::AbstractArray{<:Real}, logmean, dims)
-    l = logsubexp.(x, logmean)
+# the result by a constant, so the gradient is the same either way. The gradient is
+# translation-invariant, so everything is computed from max-centered values (a large common
+# offset cancels exactly in t and never enters the arithmetic), and in the log domain
+# (squaring expm1(d) directly can underflow even when the gradient itself is representable,
+# e.g. for nearly equal entries). lm is formed before subtracting from t so that log(n)
+# cancels against logsumexp's log(n) content instead of absorbing tiny d values.
+function _∂x_logvarexp(x::AbstractArray{<:Real}, dims)
+    t = x .- maximum(x; dims)
+    lse = logsumexp(t; dims)
+    n = length(x) ÷ length(lse)
+    lm = lse .- log(convert(eltype(lse), n))
+    d = t .- lm
+    l = log.(abs.(expm1.(d)))
     S = logsumexp(2 .* l; dims)
-    return sign.(x .- logmean) .* 2 .* exp.(x .+ l .- S)
+    return sign.(d) .* 2 .* exp.(d .+ l .- S)
 end
 
 function _∂x_pullback(∂x, x)
diff --git a/test/chainrules.jl b/test/chainrules.jl
index 30bc03b..753e70a 100644
--- a/test/chainrules.jl
+++ b/test/chainrules.jl
@@ -32,3 +32,18 @@ end
         @test ΔΩ ≈ -1.0e200 / h
     end
 end
+
+@testset "chainrules large common offset" begin
+    # The gradients are translation-invariant. The spread is a multiple of ulp(c) for every
+    # offset c below, so x .+ c is exact and the gradients must agree to machine precision.
+    x = collect((1:10) .* 0.125)
+    for f in (logmeanexp, logvarexp, logstdexp)
+        g = unthunk(rrule(f, x)[2](1.0)[2])
+        Δx = [1.0; zeros(9)]
+        ΔΩ = frule((NoTangent(), Δx), f, x)[2]
+        for c in (1.0e12, 1.0e15)
+            @test unthunk(rrule(f, x .+ c)[2](1.0)[2]) ≈ g rtol = 1.0e-10
+            @test frule((NoTangent(), Δx), f, x .+ c)[2] ≈ ΔΩ rtol = 1.0e-10
+        end
+    end
+end

From 924eaca3f11517dd8329b953f6aa5cfc3ef28092 Mon Sep 17 00:00:00 2001
From: cossio <cossio@users.noreply.github.com>
Date: Wed, 29 Jul 2026 12:59:15 +0200
Subject: [PATCH 3/8] Retain sub-epsilon mean offsets in the centered gradient
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Compute the centered log-mean as log1p(mean(expm1(t))) instead of
logsumexp(t) - log(n): the latter rounds through O(1) intermediates and
absorbs offsets below their epsilon, zeroing the gradient of the
maximal entry for nearly equal inputs (Codex review on 00ba70d). With
log1p the regression case x = [-1e-200, 1e-200] now yields both
components ±1e200; the test asserts the full gradient.

Also document why the gradients differentiate under the assumption
logmean == logmeanexp(x; dims) rather than honoring an
arbitrary supplied logmean: ChainRules pullbacks cannot attribute
tangents to keyword arguments, so treating the cache kwarg as an
independent constant would silently drop the mean's x-dependence from
end-to-end gradients.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01KmceQuDSTogkuoRwSd48aj
---
 ext/LogStatFunctionsChainRulesCoreExt.jl | 22 +++++++++++++---------
 test/chainrules.jl                       |  4 ++--
 2 files changed, 15 insertions(+), 11 deletions(-)

diff --git a/ext/LogStatFunctionsChainRulesCoreExt.jl b/ext/LogStatFunctionsChainRulesCoreExt.jl
index 8410148..393dd1c 100644
--- a/ext/LogStatFunctionsChainRulesCoreExt.jl
+++ b/ext/LogStatFunctionsChainRulesCoreExt.jl
@@ -60,17 +60,21 @@ end
 
 # ∂/∂xⱼ log(var(exp.(x))) = 2 exp(xⱼ) (exp(xⱼ) - m) / Σᵢ (exp(xᵢ) - m)², with m = exp(logmean).
 # The m-dependence on x drops out because Σᵢ (exp(xᵢ) - m) = 0, and `corrected` only shifts
-# the result by a constant, so the gradient is the same either way. The gradient is
-# translation-invariant, so everything is computed from max-centered values (a large common
-# offset cancels exactly in t and never enters the arithmetic), and in the log domain
-# (squaring expm1(d) directly can underflow even when the gradient itself is representable,
-# e.g. for nearly equal entries). lm is formed before subtracting from t so that log(n)
-# cancels against logsumexp's log(n) content instead of absorbing tiny d values.
+# the result by a constant, so the gradient is the same either way. The `logmean` keyword is
+# a cache of logmeanexp(x; dims), and the rules differentiate under that assumption (a
+# ChainRules pullback cannot attribute tangents to keyword arguments, so treating logmean as
+# an independent constant would silently drop the mean's own x-dependence); it is therefore
+# recomputed here rather than taken from the caller. The gradient is translation-invariant,
+# so everything is computed from max-centered values (a large common offset cancels exactly
+# in t and never enters the arithmetic), and in the log domain (squaring expm1(d) directly
+# can underflow even when the gradient itself is representable, e.g. for nearly equal
+# entries). The centered log-mean uses log1p(mean(expm1(t))), which retains offsets below
+# the epsilon of logsumexp(t) - log(n) (e.g. lm = -1e-200 for x = [-1e-200, 1e-200]).
 function _∂x_logvarexp(x::AbstractArray{<:Real}, dims)
     t = x .- maximum(x; dims)
-    lse = logsumexp(t; dims)
-    n = length(x) ÷ length(lse)
-    lm = lse .- log(convert(eltype(lse), n))
+    s = sum(expm1.(t); dims)
+    n = length(x) ÷ length(s)
+    lm = log1p.(s ./ n)
     d = t .- lm
     l = log.(abs.(expm1.(d)))
     S = logsumexp(2 .* l; dims)
diff --git a/test/chainrules.jl b/test/chainrules.jl
index 753e70a..0379bd0 100644
--- a/test/chainrules.jl
+++ b/test/chainrules.jl
@@ -25,11 +25,11 @@ end
         Ω, pb = rrule(f, x)
         @test isfinite(Ω)
         x̄ = unthunk(pb(1.0)[2])
-        @test all(isfinite, x̄)
-        @test x̄[1] ≈ -1.0e200 / h
+        @test collect(x̄) ≈ [-1.0e200, 1.0e200] ./ h
         Ω, ΔΩ = frule((NoTangent(), [1.0, 0.0]), f, x)
         @test isfinite(Ω)
         @test ΔΩ ≈ -1.0e200 / h
+        @test frule((NoTangent(), [0.0, 1.0]), f, x)[2] ≈ 1.0e200 / h
     end
 end
 

From 1b8d41114a611a89bb69449234e4307e51a27621 Mon Sep 17 00:00:00 2001
From: cossio <cossio@users.noreply.github.com>
Date: Wed, 29 Jul 2026 13:06:47 +0200
Subject: [PATCH 4/8] Floor the centered log-mean at -log(n)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

In low precision with large reductions, sum(expm1.(t)) can round to
exactly -n (e.g. Float16 with 4095 entries of -1 pairwise-summing to
-4096), driving log1p(s / n) to -Inf and the gradients to NaN. Since
the max entry contributes exp(0) = 1, mean(exp.(t)) >= 1/n and
-log(n) is a true lower bound for lm — and in the rounding regime the
max dominates, so the floor is also the correct value. Regression
tests cover the Float16 case for rrule and frule.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01KmceQuDSTogkuoRwSd48aj
---
 ext/LogStatFunctionsChainRulesCoreExt.jl |  5 ++++-
 test/chainrules.jl                       | 17 +++++++++++++++++
 2 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/ext/LogStatFunctionsChainRulesCoreExt.jl b/ext/LogStatFunctionsChainRulesCoreExt.jl
index 393dd1c..22902a8 100644
--- a/ext/LogStatFunctionsChainRulesCoreExt.jl
+++ b/ext/LogStatFunctionsChainRulesCoreExt.jl
@@ -74,7 +74,10 @@ function _∂x_logvarexp(x::AbstractArray{<:Real}, dims)
     t = x .- maximum(x; dims)
     s = sum(expm1.(t); dims)
     n = length(x) ÷ length(s)
-    lm = log1p.(s ./ n)
+    # mean(exp.(t)) ≥ 1/n since the max entry contributes exp(0) = 1, so lm ≥ -log(n); the
+    # floor also catches low-precision sums rounding s to exactly -n (log1p(-1) == -Inf),
+    # which happens precisely when the max dominates and -log(n) is the correct value.
+    lm = max.(log1p.(s ./ n), -log(convert(eltype(s), n)))
     d = t .- lm
     l = log.(abs.(expm1.(d)))
     S = logsumexp(2 .* l; dims)
diff --git a/test/chainrules.jl b/test/chainrules.jl
index 0379bd0..c814bbf 100644
--- a/test/chainrules.jl
+++ b/test/chainrules.jl
@@ -33,6 +33,23 @@ end
     end
 end
 
+@testset "chainrules low-precision large reduction" begin
+    # In Float16 the sum of 4095 expm1(-20) ≈ -1 terms rounds to -4096 == -n, which used to
+    # drive the centered log-mean to -Inf and the gradients to NaN. The max entry dominates,
+    # so its logvarexp gradient is ≈ 2 (and half that for logstdexp).
+    x = Float16[0; fill(Float16(-20), 4095)]
+    for (f, h) in ((logvarexp, 1), (logstdexp, 2))
+        Ω, pb = rrule(f, x)
+        @test isfinite(Ω)
+        x̄ = unthunk(pb(one(Float16))[2])
+        @test all(isfinite, x̄)
+        @test x̄[1] ≈ 2 / h rtol = 0.05
+        Ω, ΔΩ = frule((NoTangent(), [1; zeros(4095)]), f, x)
+        @test isfinite(Ω)
+        @test ΔΩ ≈ 2 / h rtol = 0.05
+    end
+end
+
 @testset "chainrules large common offset" begin
     # The gradients are translation-invariant. The spread is a multiple of ulp(c) for every
     # offset c below, so x .+ c is exact and the gradients must agree to machine precision.

From aec6c513d187f0dd972008dc5a272080c985696b Mon Sep 17 00:00:00 2001
From: cossio <cossio@users.noreply.github.com>
Date: Wed, 29 Jul 2026 13:16:42 +0200
Subject: [PATCH 5/8] Harden the centered gradient for half-precision
 reductions
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Two more Codex-flagged Float16 corner cases:

- expm1(d) can overflow (d ≤ log(n) reaches 11.09 at n = 65504 while
  Float16 tops out at e^11.04), turning the gradient into NaN. Compute
  log(abs(expm1(d))) directly as max(d, 0) + log1mexp(-abs(d)) without
  materializing expm1(d).
- Individually tiny expm1 terms round to -1 at the element level, so a
  collectively significant tail (e.g. 4095 entries of -8.5) vanished
  from the centered mean and tail gradients came out zero. Accumulate
  the expm1 terms in at least Float64 (BigFloat stays BigFloat).

Regression tests cover both cases; the -log(n) floor stays as a
belt-and-braces guard.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01KmceQuDSTogkuoRwSd48aj
---
 ext/LogStatFunctionsChainRulesCoreExt.jl | 22 +++++++++++------
 test/chainrules.jl                       | 30 ++++++++++++++++--------
 2 files changed, 35 insertions(+), 17 deletions(-)

diff --git a/ext/LogStatFunctionsChainRulesCoreExt.jl b/ext/LogStatFunctionsChainRulesCoreExt.jl
index 22902a8..99ec1b4 100644
--- a/ext/LogStatFunctionsChainRulesCoreExt.jl
+++ b/ext/LogStatFunctionsChainRulesCoreExt.jl
@@ -1,7 +1,7 @@
 module LogStatFunctionsChainRulesCoreExt
 
 using LogStatFunctions: logmeanexp, logvarexp, logstdexp
-using LogExpFunctions: logsumexp
+using LogExpFunctions: log1mexp, logsumexp
 import ChainRulesCore
 
 function ChainRulesCore.frule((_, Δx), ::typeof(logmeanexp), x::AbstractArray{<:Real}; dims = :)
@@ -69,21 +69,29 @@ end
 # in t and never enters the arithmetic), and in the log domain (squaring expm1(d) directly
 # can underflow even when the gradient itself is representable, e.g. for nearly equal
 # entries). The centered log-mean uses log1p(mean(expm1(t))), which retains offsets below
-# the epsilon of logsumexp(t) - log(n) (e.g. lm = -1e-200 for x = [-1e-200, 1e-200]).
+# the epsilon of logsumexp(t) - log(n) (e.g. lm = -1e-200 for x = [-1e-200, 1e-200]); the
+# expm1 terms are accumulated in at least Float64, since in half precision individually
+# tiny terms round to -1 and their collective contribution to the mean is lost.
 function _∂x_logvarexp(x::AbstractArray{<:Real}, dims)
     t = x .- maximum(x; dims)
-    s = sum(expm1.(t); dims)
+    s = sum(_wexpm1, t; dims)
     n = length(x) ÷ length(s)
+    T = float(eltype(t))
     # mean(exp.(t)) ≥ 1/n since the max entry contributes exp(0) = 1, so lm ≥ -log(n); the
-    # floor also catches low-precision sums rounding s to exactly -n (log1p(-1) == -Inf),
-    # which happens precisely when the max dominates and -log(n) is the correct value.
-    lm = max.(log1p.(s ./ n), -log(convert(eltype(s), n)))
+    # floor catches sums rounding s to exactly -n (log1p(-1) == -Inf), which can only
+    # happen when the max dominates and -log(n) is the correct value.
+    lm = max.(T.(log1p.(s ./ n)), -log(convert(T, n)))
     d = t .- lm
-    l = log.(abs.(expm1.(d)))
+    # log(abs(expm1(d))) without materializing expm1(d), which can overflow in half
+    # precision even though d ≤ log(n) keeps the final gradient representable.
+    l = max.(d, 0) .+ log1mexp.(-abs.(d))
     S = logsumexp(2 .* l; dims)
     return sign.(d) .* 2 .* exp.(d .+ l .- S)
 end
 
+# expm1 with the accumulation widened to at least Float64 (BigFloat stays BigFloat)
+_wexpm1(u::Real) = expm1(convert(promote_type(Float64, typeof(u)), u))
+
 function _∂x_pullback(∂x, x)
     project_x = ChainRulesCore.ProjectTo(x)
     function pullback(Ω̄)
diff --git a/test/chainrules.jl b/test/chainrules.jl
index c814bbf..f6005da 100644
--- a/test/chainrules.jl
+++ b/test/chainrules.jl
@@ -37,17 +37,27 @@ end
     # In Float16 the sum of 4095 expm1(-20) ≈ -1 terms rounds to -4096 == -n, which used to
     # drive the centered log-mean to -Inf and the gradients to NaN. The max entry dominates,
     # so its logvarexp gradient is ≈ 2 (and half that for logstdexp).
-    x = Float16[0; fill(Float16(-20), 4095)]
-    for (f, h) in ((logvarexp, 1), (logstdexp, 2))
-        Ω, pb = rrule(f, x)
-        @test isfinite(Ω)
-        x̄ = unthunk(pb(one(Float16))[2])
-        @test all(isfinite, x̄)
-        @test x̄[1] ≈ 2 / h rtol = 0.05
-        Ω, ΔΩ = frule((NoTangent(), [1; zeros(4095)]), f, x)
-        @test isfinite(Ω)
-        @test ΔΩ ≈ 2 / h rtol = 0.05
+    for m in (4095, 65503)
+        x = Float16[0; fill(Float16(-20), m)]
+        for (f, h) in ((logvarexp, 1), (logstdexp, 2))
+            Ω, pb = rrule(f, x)
+            @test isfinite(Ω)
+            x̄ = unthunk(pb(one(Float16))[2])
+            @test all(isfinite, x̄)
+            @test x̄[1] ≈ 2 / h rtol = 0.05
+            Ω, ΔΩ = frule((NoTangent(), [1; zeros(m)]), f, x)
+            @test isfinite(Ω)
+            @test ΔΩ ≈ 2 / h rtol = 0.05
+        end
     end
+    # Individually tiny but collectively significant tail: expm1(-8.5) rounds to -1 in
+    # Float16, yet the 4095 tail entries contribute most of the mean. The tail gradients
+    # are ≈ -1e-7 (subnormal in Float16, hence the loose tolerance).
+    x = Float16[0; fill(Float16(-8.5), 4095)]
+    x̄ = unthunk(rrule(logvarexp, x)[2](one(Float16))[2])
+    @test all(isfinite, x̄)
+    @test x̄[1] ≈ 2 rtol = 0.05
+    @test x̄[2] ≈ -9.9e-8 rtol = 0.5
 end
 
 @testset "chainrules large common offset" begin

From bde39f95c01ab4af9805c300b397a4a0641ab866 Mon Sep 17 00:00:00 2001
From: cossio <cossio@users.noreply.github.com>
Date: Wed, 29 Jul 2026 13:26:21 +0200
Subject: [PATCH 6/8] Support abstract-eltype arrays in the centered gradient
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

float(eltype(t)) throws for arrays with abstract element type (e.g.
Vector{Real}), which match the rules' signature and worked with the
earlier fully elementwise implementation. Narrow the centered values
back to input precision per element with oftype.(float.(t), ...) —
deriving types from values, as the package's own helpers do — instead
of computing a type from the declared eltype.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01KmceQuDSTogkuoRwSd48aj
---
 ext/LogStatFunctionsChainRulesCoreExt.jl | 7 ++++---
 test/chainrules.jl                       | 9 +++++++++
 2 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/ext/LogStatFunctionsChainRulesCoreExt.jl b/ext/LogStatFunctionsChainRulesCoreExt.jl
index 99ec1b4..654aec9 100644
--- a/ext/LogStatFunctionsChainRulesCoreExt.jl
+++ b/ext/LogStatFunctionsChainRulesCoreExt.jl
@@ -76,12 +76,13 @@ function _∂x_logvarexp(x::AbstractArray{<:Real}, dims)
     t = x .- maximum(x; dims)
     s = sum(_wexpm1, t; dims)
     n = length(x) ÷ length(s)
-    T = float(eltype(t))
     # mean(exp.(t)) ≥ 1/n since the max entry contributes exp(0) = 1, so lm ≥ -log(n); the
     # floor catches sums rounding s to exactly -n (log1p(-1) == -Inf), which can only
     # happen when the max dominates and -log(n) is the correct value.
-    lm = max.(T.(log1p.(s ./ n)), -log(convert(T, n)))
-    d = t .- lm
+    lm = max.(log1p.(s ./ n), -log(n))
+    # narrow back to the input's precision elementwise (via values, not eltype, so arrays
+    # with abstract element types still work)
+    d = oftype.(float.(t), t .- lm)
     # log(abs(expm1(d))) without materializing expm1(d), which can overflow in half
     # precision even though d ≤ log(n) keeps the final gradient representable.
     l = max.(d, 0) .+ log1mexp.(-abs.(d))
diff --git a/test/chainrules.jl b/test/chainrules.jl
index f6005da..df99438 100644
--- a/test/chainrules.jl
+++ b/test/chainrules.jl
@@ -60,6 +60,15 @@ end
     @test x̄[2] ≈ -9.9e-8 rtol = 0.5
 end
 
+@testset "chainrules abstract eltype" begin
+    x = Real[0.0, 1.0]
+    for f in (logmeanexp, logvarexp, logstdexp)
+        @test unthunk(rrule(f, x)[2](1.0)[2]) ≈ unthunk(rrule(f, [0.0, 1.0])[2](1.0)[2])
+        @test frule((NoTangent(), [1.0, 0.0]), f, x)[2] ≈
+            frule((NoTangent(), [1.0, 0.0]), f, [0.0, 1.0])[2]
+    end
+end
+
 @testset "chainrules large common offset" begin
     # The gradients are translation-invariant. The spread is a multiple of ulp(c) for every
     # offset c below, so x .+ c is exact and the gradients must agree to machine precision.

From eac882aff7071a59dd19e339d93b65696f947b4b Mon Sep 17 00:00:00 2001
From: cossio <cossio@users.noreply.github.com>
Date: Wed, 29 Jul 2026 13:36:24 +0200
Subject: [PATCH 7/8] Narrow only the final gradient to input precision

Narrowing d = t - lm to input precision rounded centered offsets below
the type's resolution to zero (Float16 nearly adjacent subnormals gave
zero gradients where +-2^14 is exactly representable). Keep the whole
ratio widened and apply the oftype-on-values narrowing only to the
final result. Regression test covers the subnormal Float16 case for
logvarexp and logstdexp.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01KmceQuDSTogkuoRwSd48aj
---
 ext/LogStatFunctionsChainRulesCoreExt.jl | 10 ++++++----
 test/chainrules.jl                       | 10 ++++++++++
 2 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/ext/LogStatFunctionsChainRulesCoreExt.jl b/ext/LogStatFunctionsChainRulesCoreExt.jl
index 654aec9..38c3bf2 100644
--- a/ext/LogStatFunctionsChainRulesCoreExt.jl
+++ b/ext/LogStatFunctionsChainRulesCoreExt.jl
@@ -80,14 +80,16 @@ function _∂x_logvarexp(x::AbstractArray{<:Real}, dims)
     # floor catches sums rounding s to exactly -n (log1p(-1) == -Inf), which can only
     # happen when the max dominates and -log(n) is the correct value.
     lm = max.(log1p.(s ./ n), -log(n))
-    # narrow back to the input's precision elementwise (via values, not eltype, so arrays
-    # with abstract element types still work)
-    d = oftype.(float.(t), t .- lm)
+    # d stays widened: narrowing here can round centered offsets below the input type's
+    # resolution to zero even when the final gradient is representable
+    d = t .- lm
     # log(abs(expm1(d))) without materializing expm1(d), which can overflow in half
     # precision even though d ≤ log(n) keeps the final gradient representable.
     l = max.(d, 0) .+ log1mexp.(-abs.(d))
     S = logsumexp(2 .* l; dims)
-    return sign.(d) .* 2 .* exp.(d .+ l .- S)
+    # narrow only the final result back to the input's precision, elementwise via values
+    # rather than eltype so arrays with abstract element types still work
+    return oftype.(float.(t), sign.(d) .* 2 .* exp.(d .+ l .- S))
 end
 
 # expm1 with the accumulation widened to at least Float64 (BigFloat stays BigFloat)
diff --git a/test/chainrules.jl b/test/chainrules.jl
index df99438..41318a2 100644
--- a/test/chainrules.jl
+++ b/test/chainrules.jl
@@ -58,6 +58,16 @@ end
     @test all(isfinite, x̄)
     @test x̄[1] ≈ 2 rtol = 0.05
     @test x̄[2] ≈ -9.9e-8 rtol = 0.5
+    # Nearly adjacent subnormals: the centered offsets (±2^-25) sit below the Float16
+    # subnormal spacing, but the gradients (±2^14, and half that for logstdexp) are
+    # exactly representable.
+    x = Float16[zeros(2048); fill(nextfloat(Float16(0)), 2048)]
+    for (f, h) in ((logvarexp, 1), (logstdexp, 2))
+        x̄ = unthunk(rrule(f, x)[2](one(Float16))[2])
+        @test all(isfinite, x̄)
+        @test x̄[1] ≈ -16384 / h rtol = 0.05
+        @test x̄[end] ≈ 16384 / h rtol = 0.05
+    end
 end
 
 @testset "chainrules abstract eltype" begin

From 320f816220248b9b0fd7bb6fd6a9914c68a3c948 Mon Sep 17 00:00:00 2001
From: cossio <cossio@users.noreply.github.com>
Date: Wed, 29 Jul 2026 13:53:10 +0200
Subject: [PATCH 8/8] Simplify: drop the low-precision hardening, use
 LogExpFunctions.softmax
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Review pass against unneeded complexity. The max-centering, Float64
widening, -log(n) floor, log1mexp overflow guard, and oftype narrowing
defended Float16 arrays with tens of thousands of elements and Float64
log-values around 1e15 — regimes that don't occur in practice and where
the primal functions are at their own representational limits anyway.

Keep the one load-bearing fix: the log-domain ratio (the original
Inf/NaN underflow for nearly equal entries), centered on the logmean
the rules already compute. Use LogExpFunctions.softmax for the
logmeanexp gradient instead of a hand-rolled helper, and trim comments.
Drop the tests that only exercised the removed machinery.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01KmceQuDSTogkuoRwSd48aj
---
 ext/LogStatFunctionsChainRulesCoreExt.jl | 64 ++++++------------------
 test/chainrules.jl                       | 52 -------------------
 2 files changed, 14 insertions(+), 102 deletions(-)

diff --git a/ext/LogStatFunctionsChainRulesCoreExt.jl b/ext/LogStatFunctionsChainRulesCoreExt.jl
index 38c3bf2..e013f9d 100644
--- a/ext/LogStatFunctionsChainRulesCoreExt.jl
+++ b/ext/LogStatFunctionsChainRulesCoreExt.jl
@@ -1,18 +1,18 @@
 module LogStatFunctionsChainRulesCoreExt
 
 using LogStatFunctions: logmeanexp, logvarexp, logstdexp
-using LogExpFunctions: log1mexp, logsumexp
+using LogExpFunctions: logsumexp, softmax
 import ChainRulesCore
 
 function ChainRulesCore.frule((_, Δx), ::typeof(logmeanexp), x::AbstractArray{<:Real}; dims = :)
     Ω = logmeanexp(x; dims)
-    ΔΩ = sum(_softmax(x, dims) .* Δx; dims)
+    ΔΩ = sum(softmax(x; dims) .* Δx; dims)
     return Ω, ΔΩ
 end
 
 function ChainRulesCore.rrule(::typeof(logmeanexp), x::AbstractArray{<:Real}; dims = :)
     Ω = logmeanexp(x; dims)
-    return Ω, _∂x_pullback(_softmax(x, dims), x)
+    return Ω, _∂x_pullback(softmax(x; dims), x)
 end
 
 function ChainRulesCore.frule(
@@ -20,7 +20,7 @@ function ChainRulesCore.frule(
         dims = :, corrected::Bool = true, logmean = logmeanexp(x; dims)
     )
     Ω = logvarexp(x; dims, corrected, logmean)
-    ΔΩ = sum(_∂x_logvarexp(x, dims) .* Δx; dims)
+    ΔΩ = sum(_∂x_logvarexp(x, logmean, dims) .* Δx; dims)
     return Ω, ΔΩ
 end
 
@@ -29,7 +29,7 @@ function ChainRulesCore.rrule(
         dims = :, corrected::Bool = true, logmean = logmeanexp(x; dims)
     )
     Ω = logvarexp(x; dims, corrected, logmean)
-    return Ω, _∂x_pullback(_∂x_logvarexp(x, dims), x)
+    return Ω, _∂x_pullback(_∂x_logvarexp(x, logmean, dims), x)
 end
 
 function ChainRulesCore.frule(
@@ -37,7 +37,7 @@ function ChainRulesCore.frule(
         dims = :, corrected::Bool = true, logmean = logmeanexp(x; dims)
     )
     Ω = logstdexp(x; dims, corrected, logmean)
-    ΔΩ = sum(_∂x_logvarexp(x, dims) ./ 2 .* Δx; dims)
+    ΔΩ = sum(_∂x_logvarexp(x, logmean, dims) ./ 2 .* Δx; dims)
     return Ω, ΔΩ
 end
 
@@ -46,55 +46,19 @@ function ChainRulesCore.rrule(
         dims = :, corrected::Bool = true, logmean = logmeanexp(x; dims)
     )
     Ω = logstdexp(x; dims, corrected, logmean)
-    return Ω, _∂x_pullback(_∂x_logvarexp(x, dims) / 2, x)
+    return Ω, _∂x_pullback(_∂x_logvarexp(x, logmean, dims) / 2, x)
 end
 
-# ∂/∂xⱼ log(mean(exp.(x))) = exp(xⱼ) / Σᵢ exp(xᵢ), i.e. softmax(x). Normalizing by the
-# actual sum of the max-centered exponentials (rather than dividing exp(x - logmean) by n)
-# makes a large common offset in x cancel exactly instead of leaking ulp-level errors of
-# the offset's magnitude into the gradient.
-function _softmax(x::AbstractArray{<:Real}, dims)
-    y = exp.(x .- maximum(x; dims))
-    return y ./ sum(y; dims)
-end
-
-# ∂/∂xⱼ log(var(exp.(x))) = 2 exp(xⱼ) (exp(xⱼ) - m) / Σᵢ (exp(xᵢ) - m)², with m = exp(logmean).
-# The m-dependence on x drops out because Σᵢ (exp(xᵢ) - m) = 0, and `corrected` only shifts
-# the result by a constant, so the gradient is the same either way. The `logmean` keyword is
-# a cache of logmeanexp(x; dims), and the rules differentiate under that assumption (a
-# ChainRules pullback cannot attribute tangents to keyword arguments, so treating logmean as
-# an independent constant would silently drop the mean's own x-dependence); it is therefore
-# recomputed here rather than taken from the caller. The gradient is translation-invariant,
-# so everything is computed from max-centered values (a large common offset cancels exactly
-# in t and never enters the arithmetic), and in the log domain (squaring expm1(d) directly
-# can underflow even when the gradient itself is representable, e.g. for nearly equal
-# entries). The centered log-mean uses log1p(mean(expm1(t))), which retains offsets below
-# the epsilon of logsumexp(t) - log(n) (e.g. lm = -1e-200 for x = [-1e-200, 1e-200]); the
-# expm1 terms are accumulated in at least Float64, since in half precision individually
-# tiny terms round to -1 and their collective contribution to the mean is lost.
-function _∂x_logvarexp(x::AbstractArray{<:Real}, dims)
-    t = x .- maximum(x; dims)
-    s = sum(_wexpm1, t; dims)
-    n = length(x) ÷ length(s)
-    # mean(exp.(t)) ≥ 1/n since the max entry contributes exp(0) = 1, so lm ≥ -log(n); the
-    # floor catches sums rounding s to exactly -n (log1p(-1) == -Inf), which can only
-    # happen when the max dominates and -log(n) is the correct value.
-    lm = max.(log1p.(s ./ n), -log(n))
-    # d stays widened: narrowing here can round centered offsets below the input type's
-    # resolution to zero even when the final gradient is representable
-    d = t .- lm
-    # log(abs(expm1(d))) without materializing expm1(d), which can overflow in half
-    # precision even though d ≤ log(n) keeps the final gradient representable.
-    l = max.(d, 0) .+ log1mexp.(-abs.(d))
+# ∂/∂xⱼ log(var(exp.(x))) = 2 exp(xⱼ) (exp(xⱼ) - m) / Σᵢ (exp(xᵢ) - m)², with m = exp(logmean);
+# the dependence of m on x drops out since Σᵢ (exp(xᵢ) - m) = 0, and `corrected` only shifts
+# the result by a constant. Log-domain to avoid under/overflow of the squared terms.
+function _∂x_logvarexp(x::AbstractArray{<:Real}, logmean, dims)
+    d = x .- logmean
+    l = log.(abs.(expm1.(d)))
     S = logsumexp(2 .* l; dims)
-    # narrow only the final result back to the input's precision, elementwise via values
-    # rather than eltype so arrays with abstract element types still work
-    return oftype.(float.(t), sign.(d) .* 2 .* exp.(d .+ l .- S))
+    return sign.(d) .* 2 .* exp.(d .+ l .- S)
 end
 
-# expm1 with the accumulation widened to at least Float64 (BigFloat stays BigFloat)
-_wexpm1(u::Real) = expm1(convert(promote_type(Float64, typeof(u)), u))
-
 function _∂x_pullback(∂x, x)
     project_x = ChainRulesCore.ProjectTo(x)
     function pullback(Ω̄)
diff --git a/test/chainrules.jl b/test/chainrules.jl
index 41318a2..ee19ea6 100644
--- a/test/chainrules.jl
+++ b/test/chainrules.jl
@@ -33,43 +33,6 @@ end
     end
 end
 
-@testset "chainrules low-precision large reduction" begin
-    # In Float16 the sum of 4095 expm1(-20) ≈ -1 terms rounds to -4096 == -n, which used to
-    # drive the centered log-mean to -Inf and the gradients to NaN. The max entry dominates,
-    # so its logvarexp gradient is ≈ 2 (and half that for logstdexp).
-    for m in (4095, 65503)
-        x = Float16[0; fill(Float16(-20), m)]
-        for (f, h) in ((logvarexp, 1), (logstdexp, 2))
-            Ω, pb = rrule(f, x)
-            @test isfinite(Ω)
-            x̄ = unthunk(pb(one(Float16))[2])
-            @test all(isfinite, x̄)
-            @test x̄[1] ≈ 2 / h rtol = 0.05
-            Ω, ΔΩ = frule((NoTangent(), [1; zeros(m)]), f, x)
-            @test isfinite(Ω)
-            @test ΔΩ ≈ 2 / h rtol = 0.05
-        end
-    end
-    # Individually tiny but collectively significant tail: expm1(-8.5) rounds to -1 in
-    # Float16, yet the 4095 tail entries contribute most of the mean. The tail gradients
-    # are ≈ -1e-7 (subnormal in Float16, hence the loose tolerance).
-    x = Float16[0; fill(Float16(-8.5), 4095)]
-    x̄ = unthunk(rrule(logvarexp, x)[2](one(Float16))[2])
-    @test all(isfinite, x̄)
-    @test x̄[1] ≈ 2 rtol = 0.05
-    @test x̄[2] ≈ -9.9e-8 rtol = 0.5
-    # Nearly adjacent subnormals: the centered offsets (±2^-25) sit below the Float16
-    # subnormal spacing, but the gradients (±2^14, and half that for logstdexp) are
-    # exactly representable.
-    x = Float16[zeros(2048); fill(nextfloat(Float16(0)), 2048)]
-    for (f, h) in ((logvarexp, 1), (logstdexp, 2))
-        x̄ = unthunk(rrule(f, x)[2](one(Float16))[2])
-        @test all(isfinite, x̄)
-        @test x̄[1] ≈ -16384 / h rtol = 0.05
-        @test x̄[end] ≈ 16384 / h rtol = 0.05
-    end
-end
-
 @testset "chainrules abstract eltype" begin
     x = Real[0.0, 1.0]
     for f in (logmeanexp, logvarexp, logstdexp)
@@ -78,18 +41,3 @@ end
             frule((NoTangent(), [1.0, 0.0]), f, [0.0, 1.0])[2]
     end
 end
-
-@testset "chainrules large common offset" begin
-    # The gradients are translation-invariant. The spread is a multiple of ulp(c) for every
-    # offset c below, so x .+ c is exact and the gradients must agree to machine precision.
-    x = collect((1:10) .* 0.125)
-    for f in (logmeanexp, logvarexp, logstdexp)
-        g = unthunk(rrule(f, x)[2](1.0)[2])
-        Δx = [1.0; zeros(9)]
-        ΔΩ = frule((NoTangent(), Δx), f, x)[2]
-        for c in (1.0e12, 1.0e15)
-            @test unthunk(rrule(f, x .+ c)[2](1.0)[2]) ≈ g rtol = 1.0e-10
-            @test frule((NoTangent(), Δx), f, x .+ c)[2] ≈ ΔΩ rtol = 1.0e-10
-        end
-    end
-end