Fix overflow in the eigh_trunc_pullback! Sylvester iteration

src/pullbacks/eigh.jl

@@ -153,9 +153,11 @@ function eigh_trunc_pullback!(
     if !iszerotangent(ΔV₊)
         X₀ = rdiv!(ΔV₊, Diagonal(D))
         AP = mul!(copy(A), V * Dmat, V', -1, 1)
-        dabsmax = maximum(abs, D)
-        AP ./= dabsmax
-        D⁻¹ = dabsmax ./ D
+        # Normalize by the smallest retained |eigenvalue|, as `svd_trunc_pullback!` does
+        # with `S[end]`. That caps `max|D⁻¹|` at 1, so squaring can only shrink it.
+        dabsmin = minimum(abs, D)
+        AP ./= dabsmin
+        D⁻¹ = dabsmin ./ D
         X₁ = rmul!(AP * X₀, Diagonal(D⁻¹))
         X₁ .+= X₀
         Xₖ, Xₖ₊₁ = X₁, X₀
@@ -182,7 +184,8 @@ function eigh_trunc_pullback!(
         # take the Hermitian part, and cannot apply project_hermitian! to
         # the current contents of ΔA
         # TODO: add an `add_project_hermitian!`
-        ΔA′ = project_hermitian!(mul!(AP, Z, V', 1, 1)) # recycle AP
+        # recycle AP's storage, but overwrite it: the loop leaves APₖ in that buffer
+        ΔA′ = project_hermitian!(mul!(AP, Z, V'))
         ΔA .+= ΔA′
     else
         # in this case, Z * V' is automatically Hermitian, so we can directly add it to ΔA

test/common/eigh_trunc_pullback_overflow.jl

@@ -0,0 +1,41 @@
+using MatrixAlgebraKit
+using MatrixAlgebraKit: eigh_pullback!, eigh_trunc_pullback!, diagview
+using Test
+using StableRNGs
+using LinearAlgebra: Diagonal, norm, qr
+
+# Regression test for the Sylvester doubling iteration in `eigh_trunc_pullback!`,
+# which previously gave rise to overflow and underflow issues when the gap between the
+# absolute values of the largest discarded and the smallest retained eigenvalues
+# was small, but the gap between the absolute values of the largest and smallest retained
+# eigenvalues was large.
+
+@testset "Regression test for overflow in eigh_trunc_pullback! ($T)" for T in (Float64, ComplexF64)
+    rng = StableRNG(12345)
+    n, p = 24, 6
+    # construct artifical spectrum with |λ_discarded|max / |λ_kept|min close to 1, but
+    # |λ_kept|max / |λ_kept|min large
+    λ = 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, T, n, n)).Q)
+    A = Q * Diagonal(T.(λ)) * Q'
+    A = (A + A') / 2
+
+    V, D = Q[:, 1:p], real(T).(λ[1:p])       # truncated factors taken straight from the
+    Dmat = Diagonal(D)                        # construction: no phase/ordering ambiguity
+    @test 0.9 < maximum(abs, λ[(p + 1):end]) / minimum(abs, D) < 1
+
+    ΔD = randn(rng, real(T), p)
+    ΔV = randn(rng, T, n, p)
+    ΔV .-= V * Diagonal(diagview(V' * ΔV))    # drop the gauge-sensitive component
+
+    ΔA = eigh_trunc_pullback!(zeros(T, n, n), A, (Dmat, V), (Diagonal(ΔD), copy(ΔV)))
+    @test all(isfinite, ΔA)
+    @test ΔA ≈ ΔA'
+
+    # reference: the same cotangents through the full pullback, which has no series in it
+    ΔA_full = eigh_pullback!(
+        zeros(T, n, n), A, (Diagonal(real(T).(λ)), Q),
+        (Diagonal(vcat(ΔD, zeros(real(T), n - p))), hcat(copy(ΔV), zeros(T, n, n - p)))
+    )
+    @test ΔA ≈ ΔA_full rtol = 1.0e-8
+end