Compute logvarexp/logstdexp gradients in the log domain - #5
Conversation
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
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #5 +/- ##
=========================================
Coverage 100.00% 100.00%
=========================================
Files 2 2
Lines 59 59
=========================================
Hits 59 59 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: f361cae057
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
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
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 00ba70d471
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
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
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 924eaca3f1
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
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
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 1b8d41114a
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
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
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: aec6c513d1
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
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
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: bde39f95c0
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
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
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
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 320f816220
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
Follow-up to the post-merge Codex review of #3. Of its four comments, two were real and are fixed here; two were false positives.
Fixed
_∂x_logvarexpsquaredexp.(x .- logmean) .- 1directly, so for nearly equal entries the denominator underflowed to zero and thelogvarexp/logstdexprules returnedInf/NaNeven though the primal and the true gradient are representable. Reproduced withx = [-1e-200, 1e-200]: primal ≈ −920.3, gradient came out[-Inf, Inf]instead of≈ ±1e200. The gradient ratio is now formed in the log domain vialogsubexp/logsumexp. Regression tests added.unthunkon the incoming cotangent before broadcasting. (Broadcasting already unthunks implicitly viaBroadcast.broadcastable(::AbstractThunk), so this is hardening rather than a bug fix.)AbstractArray{<:Real}, whilelogmeanexpitself accepts complex arrays. The manual and changelog now state the real-array restriction instead of advertising unqualified AD support.Not fixed (false positive)
all(d ≤ ndims(x) for d in dims)withdims::Intis fine — numbers are iterable in Julia (all(d <= 2 for d in 1) == true), and the 330 tests in that loop already ran green in CI.Verification
Pkg.test()passes on Julia 1.12.6, including new "extreme values" regression tests (finite gradients with correct magnitude atx = [-1e-200, 1e-200]for bothrruleandfrule).--checkpasses.🤖 Generated with Claude Code
https://claude.ai/code/session_01KmceQuDSTogkuoRwSd48aj