Fix overflow in the eigh_trunc_pullback! Sylvester iteration

[leburgel]

eigh_trunc_pullback! returns NaN for a class of well-conditioned inputs.

eigh_trunc_pullback! and svd_trunc_pullback! sum the same Neumann series by the same
doubling iteration, but they normalize it differently:

# svd.jl                            # eigh.jl
minS = @view S[end:end]             dabsmax = maximum(abs, D)
AP ./= minS                         AP ./= dabsmax
S⁻¹ = minS ./ S                     D⁻¹ = dabsmax ./ D

S is sorted descending, so S[end] is the smallest retained singular value and
max|S⁻¹| = 1 exactly — squaring can only shrink it. eigh scales by the largest |λ|
instead, leaving max|D⁻¹| = |λ|max / |λ_retained|min > 1, which squares its way to Inf
while APₖ underflows to 0. The next iterate is 0 * Inf = NaN.

The failure needs the truncation to be marginal — ρ = |λ_disc|max / |λ_kept|min close to
1 — so that the series needs more doublings than the factors survive. A clean truncation gap
converges in two or three doublings, long before anything leaves range, which is why CI does
not catch it.

The failing input

A 270×270 block from a CTMRG gradient (PEPSKit.jl, C4v enlarged corner, D=3, χ=30). Exactly
Hermitian, ‖A‖ = 1.000000, V isometric to 1.5e-13:

30 eigenvalues retained, 240 discarded

  retained, largest 8 by |λ|:          retained, smallest 4 by |λ|:
    +5.85210958e-01                      -2.83099015e-02
    -5.85210958e-01                      +2.83099015e-02
    +2.37318800e-01                      +2.64251649e-02
    -2.37318800e-01                      -2.64251649e-02
    -2.20979036e-01
    +2.20979036e-01                    discarded, largest 6 by |λ|:
    +1.10288295e-01                      +2.41798930e-02
    -1.10288295e-01                      -2.41798930e-02
                                         +2.22598811e-02
                                         -2.22598811e-02
                                         +2.01417888e-02
                                         -2.01417888e-02

  |λ|max / |λ_retained|min              = 22.15    <- sets how fast D⁻¹ₖ overflows
  rho = |λ_disc|max / |λ_kept|min       = 0.915033 <- sets how fast the series converges
  smallest gap within retained spectrum = 1.885e-03  (no degeneracy)
  cond of the Sylvester operator        = 2.714e+02 (well conditioned)

22.15^(2^k) passes 1.8e308 at k=8, while ρ = 0.915 needs 338 terms — 9 doublings — to
reach 1e-13. One doubling short. (The ± pairs are a consequence of the C4v symmetry; the
overflow depends only on |λ|.)

It does not throw. It warns (Sylvester iteration did not converge ... final norm of X: NaN) and returns the NaN, which surfaces far from the cause — for us as
ArgumentError: cannot set off-diagonal entry (2, 1) to a nonzero value.

Self-contained reproducer — no data files, runs against this branch and against main

Needs only MatrixAlgebraKit, StableRNGs and stdlibs. Calls eigh_trunc_pullback!, then
runs the loop transcribed with the normalizer as a flag so the two choices can be compared.

# Self-contained reproducer for the eigh_trunc_pullback! overflow. No data files.
#   julia --project=<env with MatrixAlgebraKit> repro_standalone.jl
using MatrixAlgebraKit, LinearAlgebra, StableRNGs, Printf
const MAK = MatrixAlgebraKit
using MatrixAlgebraKit: eigh_pullback!, eigh_trunc_pullback!, qr_null, diagview

# --- an adversarial but perfectly well-conditioned truncation ------------------------
# |λ|max / |λ_kept|min = 30, so with the current normalizer max|D⁻¹| = 30 and D⁻¹^(2^k)
# overflows Float64 at 2^k > 208 (k = 8); rho = 0.95, so rho^m < 1e-13 needs m > 583
# (k = 10). The series therefore needs more doublings than the factors survive.
n, p = 24, 6
rng = StableRNG(12345)
λ = vcat([30.0, -30.0, 5.0, -5.0, 1.0, -1.0], collect(range(0.95, 0.1; length = n - p)))
Q = Matrix(qr(randn(rng, Float64, n, n)).Q)
A = Q * Diagonal(λ) * Q'; A = (A + A') / 2
V, D = Q[:, 1:p], λ[1:p]; Dmat = Diagonal(D)

ΔD = randn(rng, p)
ΔV = randn(rng, n, p)
ΔV .-= V * Diagonal(diagview(V' * ΔV))          # drop the gauge-sensitive component

# the Sylvester operator X -> (W'AW)X - XD, whose conditioning is what actually matters
W = qr_null(V)
Aperp = W' * A * W
M = kron(Matrix(I, p, p), Aperp) - kron(Matrix(Dmat), Matrix(I, n - p, n - p))
@printf("rho = %.4f   cond(Sylvester) = %.3e\n",
        maximum(abs, λ[(p + 1):end]) / minimum(abs, D), cond(M))

# reference: the full untruncated pullback, which contains no series at all
ΔA_full = eigh_pullback!(zeros(n, n), A, (Diagonal(λ), Q),
                         (Diagonal(vcat(ΔD, zeros(n - p))), hcat(copy(ΔV), zeros(n, n - p))))

ΔA = eigh_trunc_pullback!(zeros(n, n), A, (Dmat, V), (Diagonal(ΔD), copy(ΔV)))
@printf("eigh_trunc_pullback!: %d / %d non-finite, rel. error vs full = %.3e\n",
        count(!isfinite, ΔA), length(ΔA), norm(ΔA - ΔA_full) / norm(ΔA_full))

# --- the loop, transcribed, with the normalizer as a flag -----------------------------
# `scale = maximum` is what eigh does today; `scale = minimum` is what svd_trunc_pullback!
# does (it divides by S[end], the smallest retained singular value).
function doubling(A, Dmat, V, ΔV₊; degeneracy_atol, maxiter = 100, scale = maximum)
    D = diagview(Dmat)
    X₀ = ΔV₊ / Diagonal(D)
    AP = mul!(copy(A), V * Dmat, V', -1, 1)
    dabs = scale(abs, D); AP ./= dabs; D⁻¹ = dabs ./ D
    X₁ = rmul!(AP * X₀, Diagonal(D⁻¹)); X₁ .+= X₀
    Xₖ, Xₖ₊₁ = X₁, copy(X₀)
    APₖ, APₖ₊₁ = AP * AP, copy(AP)
    D⁻¹ₖ, D⁻¹ₖ₊₁ = D⁻¹ .^ 2, copy(D⁻¹)
    @printf("  k   terms   max|AP_k|    max|Dinv_k|   product      increment\n")
    for k in 1:maxiter
        Xₖ₊₁ = rmul!(mul!(Xₖ₊₁, APₖ, Xₖ), Diagonal(D⁻¹ₖ))
        a, d = maximum(abs, APₖ), maximum(abs, D⁻¹ₖ)
        @printf("  %2d  %6d   %.4e   %.4e   %.4e   %.4e\n", k, 2^k, a, d, a * d, maximum(abs, Xₖ₊₁))
        maximum(abs, Xₖ₊₁) < degeneracy_atol && (println("  -> converged"); break)
        isfinite(a * d) || (println("  -> NaN; the iteration is dead"); break)
        Xₖ₊₁ .+= Xₖ
        D⁻¹ₖ₊₁ .= D⁻¹ₖ .^ 2
        APₖ₊₁ = mul!(APₖ₊₁, APₖ, APₖ)
        Xₖ, Xₖ₊₁ = Xₖ₊₁, Xₖ; APₖ, APₖ₊₁ = APₖ₊₁, APₖ; D⁻¹ₖ, D⁻¹ₖ₊₁ = D⁻¹ₖ₊₁, D⁻¹ₖ
    end
    return Xₖ
end

atol = MAK.default_pullback_rank_atol(Dmat)
ΔV₊ = ΔV - V * (V' * ΔV)
println("\nnormalizing by maximum(abs, D)  -- what eigh does today:")
doubling(A, Dmat, V, copy(ΔV₊); degeneracy_atol = atol, scale = maximum)
println("\nnormalizing by minimum(abs, D)  -- what svd_trunc_pullback! does:")
doubling(A, Dmat, V, copy(ΔV₊); degeneracy_atol = atol, scale = minimum)

Output on main:

rho = 0.9500   cond(Sylvester) = 6.190e+02
┌ Warning: Sylvester iteration did not converge after 100 iterations, final norm of X: NaN)
└ @ MatrixAlgebraKit ~/.julia/packages/MatrixAlgebraKit/.../src/pullbacks/eigh.jl:171
eigh_trunc_pullback!: 576 / 576 non-finite, rel. error vs full = NaN

Output on this branch:

rho = 0.9500   cond(Sylvester) = 6.190e+02
eigh_trunc_pullback!: 0 / 576 non-finite, rel. error vs full = 6.500e-15

normalizing by maximum(abs, D)  -- what eigh does today:
  k   terms   max|AP_k|    max|Dinv_k|   product      increment
   1       2   5.4282e-04   9.0000e+02   4.8854e-01   2.4044e+00
   2       4   4.3733e-07   8.1000e+05   3.5424e-01   3.1685e+00
   3       8   3.6441e-13   6.5610e+11   2.3909e-01   3.8465e+00
   4      16   3.2010e-25   4.3047e+23   1.3779e-01   3.5950e+00
   5      32   2.8880e-49   1.8530e+47   5.3516e-02   1.9738e+00
   6      64   2.7962e-97   3.4337e+94   9.6012e-03   4.2573e-01
   7     128   3.0063e-193   1.1790e+189   3.5445e-04   1.6367e-02
   8     256   0.0000e+00   Inf   NaN   NaN
  -> NaN; the iteration is dead

normalizing by minimum(abs, D)  -- what svd_trunc_pullback! does:
  k   terms   max|AP_k|    max|Dinv_k|   product      increment
   1       2   4.8854e-01   1.0000e+00   4.8854e-01   2.4044e+00
   2       4   3.5424e-01   1.0000e+00   3.5424e-01   3.1685e+00
   3       8   2.3909e-01   1.0000e+00   2.3909e-01   3.8465e+00
   4      16   1.3779e-01   1.0000e+00   1.3779e-01   3.5950e+00
   5      32   5.3516e-02   1.0000e+00   5.3516e-02   1.9738e+00
   6      64   9.6012e-03   1.0000e+00   9.6012e-03   4.2573e-01
   7     128   3.5445e-04   1.0000e+00   3.5445e-04   1.6367e-02
   8     256   4.9882e-07   1.0000e+00   4.9882e-07   2.3068e-05
   9     512   9.8898e-13   1.0000e+00   9.8898e-13   4.5736e-11
  10    1024   3.8875e-24   1.0000e+00   3.8875e-24   1.7978e-22
  -> converged

max|Dinv_k| is pinned at exactly 1 in the second run, and the product column — the
quantity that actually governs convergence — is identical in both at every k where the
first still has numbers. The change is a rescaling, not a change of algorithm.

Changes

1. Normalize by the smallest retained |eigenvalue|, matching svd_trunc_pullback!:

-dabsmax = maximum(abs, D)
-AP ./= dabsmax
-D⁻¹ = dabsmax ./ D
+dabsmin = minimum(abs, D)
+AP ./= dabsmin
+D⁻¹ = dabsmin ./ D

2. β = 0 in the final assembly. This is a separate, pre-existing bug:

-ΔA′ = project_hermitian!(mul!(AP, Z, V', 1, 1)) # recycle AP
+ΔA′ = project_hermitian!(mul!(AP, Z, V', 1, 0))

APₖ₊₁ aliases the AP buffer (APₖ, APₖ₊₁ = AP * AP, AP), so the squaring loop writes
through it and β = 1 adds whatever it last held into the result.

This is already wrong on main, independently of change 1. When the truncation is well
separated the loop converges at k = 1, so no squaring has been written yet and the buffer
still holds the original AP — of size |λ_disc|max / |λ|max, not zero. The resulting error
tracks exactly that ratio (error vs the full pullback, main, retained spectrum
[±30, ±5, ±1]):

discarded band	|λ_disc|max / |λ|max	main	this PR
1e-6 … 1e-7	3.33e-08	1.664e-08	4.625e-13
1e-8 … 1e-9	3.33e-10	1.664e-10	1.776e-15

It stays hidden in the cases that run many doublings only because the old normalizer squares
AP to exactly 0.0 there. Change 1 sets ‖AP‖ = ρ, so APₖ no longer vanishes and the
term would become visible everywhere: with change 1 alone the three spectra below give
errors of 3.1e-9, 1.5e-15 and 1.5e-7.

Verification

Against the full untruncated pullback on the 270×270 block (eigh_pullback! with
Vfull = [V, W*Q], Dfull = [D; Λ], zero-padded cotangents; reconstruction error 8.2e-15):

	rel. error vs full pullback
before	NaN
after	1.494e-12

Hermiticity of the result is exactly 0. Same comparison on three synthetic spectra,
confirming no regression where the iteration already worked:

spectrum	ρ	before	after
well-separated	0.01	4.185e-15	1.676e-15
moderate	0.50	3.181e-15	3.207e-15
marginal	0.95	NaN	7.974e-13

The new test/common/eigh_trunc_pullback_overflow.jl covers both failure modes against the
full pullback, for Float64 and ComplexF64, and fails on main in both:

• a spectrum needing more doublings than the factors survive (|λ|max/|λ_kept|min = 30,
  ρ = 0.95) — main returns NaN;
• a well-separated one that converges at k = 1 (ρ = 1e-6) — main is off by 1.6e-08.

Not addressed here

• Both eigh.jl and svd.jl exit on norm(Xₖ₊₁, Inf) < degeneracy_atol, using a parameter
  named for degeneracy detection as an absolute tolerance on the increment.
• The TODO on maxiter looks tractable now: with this normalization ρ is exactly
  max|AP₁| · max|D⁻¹₁|, so the required doublings can be estimated up front.
• A truncation with ρ ≥ 1 still diverges, and still reports it only as a warning plus a
  NaN return.

[leburgel]

I ran into this issue while trying to properly benchmark QuantumKitHub/PEPSKit.jl#430.

[Jutho]

Nice catch. Thanks for the fix Lander & 🤖 .

[lkdvos]

Can we rewrite this into a small comment instead?

[leburgel]

I fought a bit to get the diff and description compact, but looks like I missed this one. I'll compress it.

[leburgel]

I cut it down to the most compact I could get it while still describing the original problem in full. I thought this could be useful, but I can also just say "Regression test for over/underflow issues in eigh_trunc_pullback!" if that's better.

[lkdvos]

Is there a way to add this to the existing tests for mooncake/enzyme instead, or do we really want to start testing our pullbacks directly? I'm not necessarily convinced for either way, just wanted to bring this up for discussion briefly

[leburgel]

Either is fine for me, I didn't really think about where exactly I was adding this in comparison to the other pullback tests.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

Files with missing lines	Coverage Δ
src/pullbacks/eigh.jl	86.13% <100.00%> (ø)

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.