Differentiate general eigen and eigvals

[JoshuaLampert]

Currently, ForwardDiff.jl only supports differentiating through eigen and eigvals of Symmetric or Hermitian matrices, but not for a general StridedMatrix. In JuliaManifolds/Manifolds.jl#27 (comment), there is a working version for differentiating through eigen and eigvals for general matrices, which also works for complex eigenvalues (e.g., DifferentiableEigen.jl does not support complex eigenvalues). I just copied that version here for higher visibility, more convenient usage, and to avoid type piracy in my (and maybe other's) code. So credit goes to @mateuszbaran. The algorithm is from https://people.maths.ox.ac.uk/gilesm/files/NA-08-01.pdf.
I open this as a draft PR because I am not sure if there is something additional that needs to be considered. If this PR is wanted, I can add tests and do a bit of cleanup.

cc @dlfivefifty since he added support for eigen for symmetric and Hermitian matrices.
Closes #111.

[codecov]

Codecov Report

❌ Patch coverage is 97.05882% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 91.05%. Comparing base (569af35) to head (f30acb9).

Files with missing lines	Patch %	Lines
src/dual.jl	97.05%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #788      +/-   ##
==========================================
+ Coverage   90.68%   91.05%   +0.37%     
==========================================
  Files          11       11              
  Lines        1052     1118      +66     
==========================================
+ Hits          954     1018      +64     
- Misses         98      100       +2     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

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

[JoshuaLampert]

I added some tests. I would be happy for a review.

[JoshuaLampert]

@devmotion, could you do a review of this PR?

[JoshuaLampert]

Another gentle bump on this PR. Is there anything I can do to speed up the review process?

[devmotion]

Sorry for the long delay. I went through this with the help of Claude, since the interaction between the eigen rules and nested Duals isn't obvious from reading the diff. First derivatives are correct, but second derivatives aren't, and the eigenvector derivatives use a different normalization than the one eigen actually returns. Both are silent wrong answers rather than errors.

The formula itself matches Giles and is transcribed correctly, and the general approach is fine. What goes wrong is specific to how the computation is arranged.

The failing CI jobs on pre are the pre-existing JET failures in test/QATest.jl, not this PR. The suite passes for me on the branch (9311/9311).

Second derivatives are wrong

sum(eigvals(B(w))) is tr(B(w)), so for a B that is affine in w it is linear and its Hessian has to be exactly zero:

using ForwardDiff, LinearAlgebra
B(w) = [3.0+w[1] 1.0+w[2]; 0.4+w[2] 2.0-2*w[1]]     # tr(B(w)) == 5 - w[1]

julia> ForwardDiff.hessian(w -> sum(eigvals(B(w))), [0.11, -0.07])
2×2 Matrix{Float64}:
  1.50425  -0.467918
 -2.27002   0.767292

julia> ForwardDiff.hessian(w -> tr(B(w)), [0.11, -0.07])
2×2 Matrix{Float64}:
 0.0  0.0
 0.0  0.0

One degree up, sum(eigvals(B(w)).^2) is tr(B^2) and quadratic, so the Hessian is [10 0; 0 4]:

julia> ForwardDiff.hessian(w -> sum(eigvals(B(w)) .^ 2), [0.11, -0.07])
2×2 Matrix{Float64}:
  16.434    0.0194723
 -11.7302   7.28187

It isn't even symmetric. Computing the same thing as jacobian of gradient gives [10 0; 0 4] to 5e-10, as does the Symmetric path.

This is the case in #111, which asks for a Hessian:

using FiniteDifferences
const fdm = central_fdm(5, 1)
fd_hessian(f, w) = FiniteDifferences.jacobian(fdm, x -> FiniteDifferences.grad(fdm, f, x)[1], w)[1]

S(w) = [w[1]^2 w[1]*w[2]*w[3]; w[1]*w[2]*w[3] w[2]^2]
g(w) = sum(log, eigvals(S(w)))

julia> maximum(abs, ForwardDiff.hessian(g, [0.9,1.4,0.3]) .- fd_hessian(g, [0.9,1.4,0.3]))
1.9378524327776427

versus 6.4e-10 if S(w) is wrapped in Symmetric. So the PR doesn't close #111 in its current form.

The cause is that A_values is one Dual level below A, and the two are then combined:

A_values     = map(d -> d.value, A)
A_values_eig = eigen(A_values)
UinvAU       = A_values_eig.vectors \ A * A_values_eig.vectors

ForwardDiff.hessian seeds Dual{T,Dual{T,V,N},N} with the same tag on both levels, so when U::Dual{T,Float64,N} meets A::Dual{T,Dual{T,Float64,N},N} the same-tag * treats them as one level and applies the product rule at the outer level. U's inner partials end up in the outer slot. Reduced to one multiplication:

julia> u   # Dual{T,Float64,2}, from eigen(A_values)
Dual{T}(-0.5192863992964954, 0.7957065830762764, -0.2753367363366924)

julia> a   # Dual{T,Dual{T,Float64,2},2}, an entry of A
Dual{T}(Dual{T}(3.11,1.0,0.0), Dual{T}(1.0,0.0,0.0), Dual{T}(0.0,0.0,0.0))

julia> value(partials(a * u, 1))
1.9553610740707241        # should be value(partials(a,1)) * value(u) == -0.5192863992964954

U \ A is still correct because \ goes through promote/ldiv!; it's the subsequent * U that breaks, since generic matmul calls scalar * directly. Inserting convert.(eltype(A), U) before the multiplication makes the Hessian exact, which confirms the mechanism, but it's a workaround rather than a fix.

The tag ordering can't catch this, because both levels carry the same tag by construction.

The existing Symmetric methods don't have the problem because they never mix levels:

parts = ntuple(j -> diag(Q' * getindex.(partials.(A), j) * Q), N)
Dual{Tg}.(λ, tuple.(parts...))

Q and getindex.(partials.(A), j) are both at the value level, and the result is reassembled with Dual{Tg}.. Written that way the general case works too:

function _eigvals_general(A::StridedMatrix{Dual{Tg,T,N}}) where {Tg,T<:Real,N}
    A0 = value.(A)
    λ, U = eigen(A0)
    luU = lu(U)
    parts = ntuple(j -> diag(luU \ (getindex.(partials.(A), j) * U)), N)
    return Dual{Tg}.(λ, tuple.(parts...))
end

This agrees with the current implementation to 4.4e-16 on first derivatives (n = 4, 8, 16), gives [10 0; 0 4] to 1.8e-15 above, and matches finite differences on #111 to 6.4e-10. It only covers the real-eigenvalue case as written — the complex branch needs the same treatment, and eigen needs the eigenvector part — but the structure is what matters.

Eigenvector derivatives use a different normalization than eigen

U̇ = U (F ∘ U⁻¹ȦU) fixes the gauge by diag(U⁻¹U̇) == 0. LAPACK's geevx! returns eigenvectors with unit 2-norm and the largest-magnitude element real. For non-normal A these are different conventions, so .vectors carries LAPACK values with partials belonging to a different parameterization.

A0 = [1.0 1.0; 0.0 2.0]
Adot = [1.0 0.0; 0.0 0.0]
V = eigen(ForwardDiff.Dual.(A0, Adot)).vectors

julia> partials.(V, 1)
2×2 Matrix{Float64}:
 0.0  0.7071067811865476
 0.0  0.0

# central differences of the same expression:
2×2 Matrix{Float64}:
 5.43878e-14   0.353553
 0.0          -0.353553

dot(value.(V[:,2]), partials.(V[:,2], 1)) is 0.4999999999999999, which can't happen if the norm is held at 1.

For a random non-symmetric 3×3 with real eigenvalues, ForwardDiff.jacobian(x -> vec(eigen(reshape(x,3,3)).vectors), x) differs from central differences by 0.322 where the largest entry is 0.711. The same test on the Symmetric path gives 3.7e-13, and eigvals on the general path gives 4.4e-12, so it's the gauge and not the tolerance.

ChainRules handles this with _eigen_norm_phase_fwd! in src/rulesets/LinearAlgebra/factorization.jl, applied right after the F ∘ step:

∂c_norm  = -realdot(v, ∂v)
∂c_phase = -imag(∂v[k]) / real(v[k])          # k = argmax abs2, complex case only
∂v     .+= v .* (eltype(V) <: Real ? ∂c_norm : complex(∂c_norm, ∂c_phase))

It would be good to match that, so the two systems agree.

The eigenvector test can't fail

h(x) = begin
    v = eigen(reshape(x, 2, 2)).vectors[:,1]
    v = v / norm(v)
end

v / norm(v) is a no-op on the values, since LAPACK already returns unit-norm vectors, so it reads as cosmetic. It isn't: two eigenvector conventions differ by a scalar c(t) with c(0) == 1, and dividing by the norm removes it. With the line, the maximum difference against Calculus.finite_difference_jacobian is 6.9e-12; without it, 0.0876.

All three tests are first-order jacobians of 2×2 matrices, so nothing exercises the nesting either.

Smaller things

• No keyword arguments. eigen(A; sortby=nothing), eigvals(A; sortby=nothing) and eigen(A; permute=false) all throw MethodError, while the internal eigen(A_values) silently applies the default sortby. Forwarding kwargs... would fix it. eigen! and eigvals! are undefined as well; eigvecs happens to work through the generic fallback.
• Repeated eigenvalues give NaN in .vectors without any error, e.g. for [2.0 0.0; 0.0 2.0] and [2.0 1.0; 0.0 2.0]. eigvals is unaffected, since vals_diff is taken before F is built. It's the same weakness as _lyap_div!, so at least it's consistent, but it deserves a note somewhere.
• _lyap_div!! already does what F is used for, including the diagonal special case, with an in-place Matrix method and an MMatrix method in the StaticArrays extension. Q*_lyap_div!!(Q'ȦQ - Diagonal(λ̇), λ) in _eigen is the same expression. Reusing it removes the F loop and an n² allocation, and divides once instead of inv followed by a multiplication. The loop is also written for i in axes(A_values,1), j in axes(A_values,2), which traverses a column-major array row by row.
• eigvals(A) = eigen(A).values computes the eigenvector derivatives and throws them away. Measured, eigvals takes 1.04–1.05x the time of eigen at n = 20 and n = 40, i.e. essentially all of it. _eigen calls eigvals, not the other way round.
• A symmetric matrix passed as a plain Matrix now goes through the general path rather than _eigen(Symmetric(...)). That's the situation in Compatibility with Base linear algebra functions #111. ChainRules dispatches on ishermitian(A); doing the same would also keep eltype stable there.
• value.(A) instead of map(d -> d.value, A) and partials(partial) instead of partial.partials would match the rest of the file, and make_eigen_dual would fit better as _make_eigen_dual next to _eigvals, _eigen and _lyap_div!!. tagtype(real(partial)) and tagtype(imag(partial)) are the same type; elsewhere the tag is carried as a where {Tg,T<:Real,N} parameter instead. The section header says # General eigvals # although it covers eigen too, and is missing the #---# line the other sections have.
• Nothing in docs/src mentions eigen at all. Which matrix types are supported, the eigenvector normalization, and the repeated-eigenvalue caveat would be worth a paragraph.

Separately, the performance argument for the rewrite above is smaller than I first assumed — for one eigvals call with 8 partials I measured 0.9x at n = 10, 1.3x at n = 20, 2.0x at n = 40 and 3.8x at n = 80. It's worth doing for correctness, not for speed.

Tests

• Nothing covers nested Duals. test/ConfusionTest.jl would be the natural place for a Hessian test on a non-symmetric matrix, and Compatibility with Base linear algebra functions #111's example belongs in the suite given that the PR closes it.
• A .vectors test without the renormalization, compared against finite differences, would have caught the gauge issue.
• Everything is 2×2. A 3×3 or larger would exercise the F broadcast structurally.
• The complex branch only tests eigvals, so make_eigen_dual(::Complex, ::Complex) and the value grafting in it are untested.
• No eltype assertions, so the return type of the complex case isn't pinned.
• x2 = [0.0, -1.0, 1.0, 0.0] is a rotation with eigenvalues ±i. Both have real part 0, so the eigsortby order depends on the imaginary tie-break surviving the perturbation, which makes the finite-difference comparison more fragile than it needs to be.
• The SVector/MVector tests that exist for the symmetric case have no counterpart here, though StaticArrays doesn't support eigen for non-Hermitian matrices at all, so that may be out of scope.

[JoshuaLampert]

Thanks for the review! I added some commits developed together with my coding agent that should fix the issues. The commits also add more tests.

[devmotion]

One should basically never unconditionally extract primal values of Duals - it should basically always be based on a tag to avoid stripping partials from dual numbers that actually are primal values themselves in the context of the current differentiation.

[JoshuaLampert]

It seems extracting the primal values was not necessary.

[devmotion]

I'd suggest changing this to a lazy error message.

[JoshuaLampert]

Done.

[devmotion]

Why NaN? _lyap_div!! explicitly handles the case of zero denominator in the code below.

[JoshuaLampert]

I adapted the comment.

[devmotion]

Thanks for the updates. Another pass, again with the help of Claude. Everything from the last round is fixed: the levels are no longer mixed, the eigenvector normalization matches eigen, keyword arguments are forwarded, repeated eigenvalues error instead of returning NaN, and _lyap_div!! is reused. I checked tag safety and the nesting specifically and they hold up.

One problem is left, in the normalization commit. First derivatives are correct everywhere I tried (2.5e-9 or better against finite differences). Second derivatives through complex eigenvectors are wrong from n = 4 on.

The failing pre jobs are the pre-existing JET failures in test/QATest.jl, not this PR.

The phase gauge is fixed at the wrong index one level up

eigen returns eigenvectors of unit 2-norm whose largest-magnitude entry is real, so the derivative has to satisfy imag(u̇[k]) == 0 for that entry. Taking the maximum of |imag(u̇[k])| over all columns and seed directions of

A = [0.5 -1.3 0.2 0.9; 1.1 0.4 -0.6 0.2; 0.3 0.7 2.0 -1.1; 0.4 -0.2 1.3 0.8]

gives 5.551115123125783e-17 for a Dual matrix, and 9.165426697038974 once a second level is added the way ForwardDiff.hessian does.

The cause is isreal. _findrealmaxabs2 filters with isreal(vi), which for a Complex{<:Dual} is iszero(imag(vi)), and iszero(::Dual) requires the partials to vanish too (src/dual.jl:417). The 5.55e-17 is not zero, so at the second level no entry passes the filter at all and the function falls through to imax = firstindex(v). With V2 the eigenvectors of the doubly-nested matrix:

julia> [any(isreal, view(V2, :, i)) for i in 1:4]
4-element Vector{Bool}:
 0
 0
 0
 0

julia> [ForwardDiff._findrealmaxabs2(view(V2, :, i)) for i in 1:4]
4-element Vector{Int64}:
 1
 1
 1
 1

The largest-magnitude entry per column, i.e. the one LAPACK actually made real, is [2, 2, 4, 4].

n = 2 and n = 3 pass by accident: at n = 2 the residue comes out exactly zero, and at n = 3 the firstindex fallback happens to equal the right index. Against sign-aligned central differences the Hessian error is 5.2e-6 at n = 2, 1.03e-5 at n = 3, then 83.6, 134 and 145 at n = 4, 5, 6. Third derivatives are garbage as well.

The 5.55e-17 is itself avoidable. imag(u[k]) == 0 holds identically along the curve, so imag(u̇[k]) is exactly zero, but it is currently computed as the cancellation q + fl(r * fl(-q/r)) with q = imag(u̇[k]) and r = real(u[k]). Asserting it instead fixes the residue and the gauge at every level:

u̇ .+= u .* ċ
# `imag(u[k]) == 0` holds identically, so `imag(u̇[k])` is exactly zero; computing it as a
# cancellation leaves a rounding residue that breaks `isreal` -- and hence `_findrealmaxabs2`
# -- one differentiation level up
u̇[k] = complex(real(u̇[k]), zero(real(u̇[k])))

With that, the residual is 0.0 at both levels for n = 2 through 6, and the Hessian errors drop to 2.5e-5. Being exact rather than approximate it also fixes third and higher order. Taking k = argmax(abs2, u) and dropping the isreal filter works too, since LAPACK makes the largest-magnitude entry the real one, but it leaves 1e-16 residues, so I'd prefer the above.

This is not inherited from ChainRules. _eigen_norm_phase_fwd! is transcribed faithfully, but ChainRules only ever applies it to StridedMatrix{<:BlasFloat}, where isreal is exact. _findrealmaxabs2 here is a reimplementation and is in fact slightly more careful than the original, which seeds amax = abs2(first(x)) and can therefore return a non-real index. What both share is the silent fallback, and that is what breaks once the eltype is a Dual.

While you're in there, the imax = firstindex(v) fallback shouldn't be silent. Reaching it means the assumed LAPACK convention doesn't hold.

Finite differences on eigenvectors need sign alignment

Not a problem with the PR, but it cost me time and will cost whoever writes the regression test. eigen(::Matrix{Float64}).vectors is sign-discontinuous: LAPACK makes the largest-magnitude entry real but doesn't pin its sign. For the A above, column 3 flips between A[1] ± 1e-6, so plain central differences are off by O(1) and look like an AD bug.

Two baselines that work: assert the gauge invariants directly (imag(u̇[k]) ≈ 0 and real(u' * u̇) ≈ 0), which needs no FD and is sign-invariant; or flip each FD column so that real(V[k,i]) keeps its sign.

The test for #111 can't fail

S(w) = [w[1]^2 w[1]*w[2]*w[3]; w[1]*w[2]*w[3] w[2]^2] is Hermitian in its values and its partials, so g and gsym both dispatch to _eigvals(Symmetric(...)). With a counter in _eigvals_general it is called 0 times for either, and the two Hessians are bitwise identical. The test checks the Hermitian dispatch, not the general path that #111 needs.

A non-symmetric log-det Hessian does exercise it and matches central differences to 1.8e-9:

h(w) = sum(log, eigvals([w[1]^2 w[1]*w[2]; 0.5*w[1]*w[2] w[2]^2+1]))

HessianTest.jl nests through eigenvectors only for a real spectrum, and the complex eigenvector tests in JacobianTest.jl are all first-order and 2×2/3×3, i.e. exactly the sizes that pass by accident.

Smaller things

• _check_distinct_eigvals runs once per partial direction, 8 times per eigen at N = 8, and in _eigen_general it runs after lu(U), so a singular U can surface as SingularException rather than the intended ArgumentError. Latent in practice, every Jordan block and repeated-diagonal case I tried still threw ArgumentError. Hoisting it into the _eigen* methods right after the decomposition fixes the count and the ordering, and keeps the _lyap_div kernel purely numeric.
• The comment that eigvals "never divides by the gaps and keeps working" is too generous. It keeps running, but diag(U⁻¹ȦU) gives the diagonal of the compressed perturbation on the degenerate eigenspace, not its eigenvalues. For A(t) = [2 t; t 2] it reports [0.0, 0.0] where the analytic branches 2 ∓ t have derivatives ∓1, and the sorted vector eigvals returns, 2-|t|, 2+|t|, isn't differentiable at all. For A(t) = [2 1; t 2] the eigenvalues are 2 ± √t, so the derivative is infinite and ±2.25e15 comes back. Not a regression, ChainRules computes the same formula and returns the same numbers and tracks it as an open TODO (factorization.jl:302), but the PR is now inconsistent, with eigen erroring loudly and eigvals staying quiet.
• The check makes eigen(::Symmetric{<:Dual}), eigen(::SymTridiagonal{<:Dual}) and the StaticArrays MMatrix path throw where they previously returned NaN vectors next to good .values. Worth a release note.
• For eigvals, ChainRules gets the diagonal without the second n³ matmul: U \ Ȧ, then row-column dot products (factorization.jl:424-431). -24% time and -36% allocations at n = 40, N = 8, -44% allocations at n = 150, N = 12.
• getindex.(partials.(A), j) is allocation-identical to partials.(A, j), since the broadcast fuses and no Matrix{Partials} is materialized. Still worth switching for readability, but not for speed.

Things I checked that are fine

• value.(A) strips exactly the level named by Tg, because the signature pins the eltype invariantly as Dual{Tg,T,N}. Independent nesting (derivative ∘ derivative with two different tags) is correct for eigvals, eigen and products and powers of them, including a nonzero mixed second derivative, where det² gives -1.6.
• The ishermitian dispatch is right, and right because == on a Dual compares partials (src/dual.jl:419-425): Hermitian values with non-Hermitian partials correctly stay on the general path, A_asym gives [-0.5, 0.5] rather than the symmetrized [-1.0, 1.0].
• Keyword forwarding matches Base, which also ignores permute/scale once it routes to the symmetric algorithm and forwards sortby. sortby = nothing on the shortcut agrees with Base numerically.
• Both gauge constraint derivations are correct, _eigen!!'s fallback works for Float16 and for nested Dual value types, 0×0 and 1×1 work, and eigvecs/eigmax/eigmin pick the new methods up.

This should go in after #757

#757 renames four internals this PR calls: _lyap_div!! and _lyap_div!, which this PR inserts _check_distinct_eigvals into, and _eigvals/_eigen, which _use_symmetric forwards to. Also the _lyap_div!!(::MMatrix) override in the StaticArrays extension that this PR relies on. A trial merge conflicts in src/dual.jl and test/JacobianTest.jl.

Rebasing on #757 makes this PR smaller:

• _make_eigen_dual and its four map((val, p) -> ..., ...) call sites can go. #757's _to_duals does the same job array-at-a-time and already handles all three shapes needed here: real λ to Vector{Dual}, complex λ to Vector{Complex{Dual}}, complex U to Matrix{Complex{Dual}}.
• M[j] - Diagonal(λ_parts[j]) can go. It only zeroes the diagonal, which #757's _lyap_div_zero_diag!! does itself, and the results are identical. That's an n² allocation per direction.
• The four renamed call sites, plus partials.(A, j).

It goes the other way in one respect: since the check sits in the shared _lyap_div kernel, #757's new Hermitian{<:Dual} and Hermitian{<:Complex{<:Dual}} methods inherit it once this lands, and _use_symmetric's shortcut picks up #757's single decomposition and BLAS reassociation for free. Nothing in #757 fixes anything above, since it is Hermitian-only and the phase convention never arises there.