From d2a624e28f919f8b8eeb5a242dc20813f9f081b6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Fri, 23 Oct 2026 14:41:00 +0200 Subject: [PATCH 1/8] Fix eigen/eigvals rules for Hermitian `StridedMatrix` MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit At Hermitian `A`, the `eigen`/`eigvals` rules for `StridedMatrix` delegated to the rules for `Hermitian(A)`. The resulting (co)tangents were those of `x -> eigen(Hermitian(x))`, which only reads the upper triangle, so they were wrong for unstructured matrices (FluxML/Zygote.jl#1369). Now the eigenvalue (co)tangents are `uᵢ' ΔA uᵢ` and `U Diagonal(Δλ) U'`, with the type of the primal eigenvalues. The eigenvectors are only differentiable along Hermitian directions (phase convention of the Hermitian solver), so their (co)tangents use the Hermitian part. This also fixes the failures on Julia 1.14, where LinearAlgebra.jl#1717 made eigenvalues of complex matrices complex also for Hermitian matrices, and removes a `Core.Box` in the `eigen` pullback that broke inference. Co-Authored-By: Claude Opus 5.5 --- Project.toml | 2 +- src/rulesets/LinearAlgebra/factorization.jl | 80 ++++++++------------ test/rulesets/LinearAlgebra/factorization.jl | 64 +++++++++------- 3 files changed, 69 insertions(+), 77 deletions(-) diff --git a/Project.toml b/Project.toml index 116735a0c..aeb76428a 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "ChainRules" uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2" -version = "1.74.2" +version = "1.74.3" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" diff --git a/src/rulesets/LinearAlgebra/factorization.jl b/src/rulesets/LinearAlgebra/factorization.jl index e04bd0e21..fcab1dcbf 100644 --- a/src/rulesets/LinearAlgebra/factorization.jl +++ b/src/rulesets/LinearAlgebra/factorization.jl @@ -293,61 +293,57 @@ end ##### # TODO: -# - support correct differential of phase convention when A is hermitian # - simplify when A is diagonal # - support degenerate matrices (see #144) +# NOTE: if `A` is Hermitian, `eigen` uses the solver for Hermitian matrices. Its eigenvectors +# are unitary and follow a different phase convention than the eigenvectors of non-Hermitian +# matrices, so they are only differentiable along Hermitian directions: the rules use the +# Hermitian part of the (co)tangents of the eigenvectors. The eigenvalues are differentiable +# in all directions their type admits. + function frule((_, ΔA), ::typeof(eigen!), A::StridedMatrix{T}; kwargs...) where {T<:BlasFloat} ΔA isa AbstractZero && return (eigen!(A; kwargs...), ΔA) - if ishermitian(A) - sortby = get(kwargs, :sortby, LinearAlgebra.eigsortby) - return if sortby === nothing - frule((ZeroTangent(), Hermitian(ΔA)), eigen!, Hermitian(A)) - else - frule((ZeroTangent(), Hermitian(ΔA)), eigen!, Hermitian(A); sortby=sortby) - end - end + isherm = ishermitian(A) F = eigen!(A; kwargs...) λ, V = F.values, F.vectors - tmp = V \ ΔA + tmp = isherm ? V' * ΔA : V \ ΔA ∂K = tmp * V ∂Kdiag = @view ∂K[diagind(∂K)] ∂λ = eltype(λ) <: Real ? real.(∂Kdiag) : copy(∂Kdiag) + isherm && hermitianpart!(∂K) ∂K ./= transpose(λ) .- λ fill!(∂Kdiag, 0) ∂V = mul!(tmp, V, ∂K) - _eigen_norm_phase_fwd!(∂V, A, V) + # `A` is overwritten but the phase convention depends only on its type + _eigen_norm_phase_fwd!(∂V, isherm ? Hermitian(A) : A, V) ∂F = Tangent{typeof(F)}(values = ∂λ, vectors = ∂V) return F, ∂F end function rrule(::typeof(eigen), A::StridedMatrix{T}; kwargs...) where {T<:Union{Real,Complex}} + isherm = ishermitian(A) F = eigen(A; kwargs...) - function eigen_pullback(ΔF::Tangent) + function eigen_pullback(ΔF::ChainRulesCore.AbstractTangent) + _ΔF = unthunk(ΔF) + _ΔF isa AbstractZero && return (NoTangent(), _ΔF) λ, V = F.values, F.vectors - Δλ, ΔV = ΔF.values, ΔF.vectors + Δλ, ΔV = _ΔF.values, _ΔF.vectors ΔV isa AbstractZero && Δλ isa AbstractZero && return (NoTangent(), Δλ + ΔV) - if eltype(λ) <: Real && ishermitian(A) - hermA = Hermitian(A) - ∂V = ΔV isa AbstractZero ? ΔV : copyto!(similar(ΔV), ΔV) - ∂hermA = eigen_rev!(hermA, λ, V, Δλ, ∂V) - ∂Atriu = _symherm_back(typeof(hermA), ∂hermA, Symbol(hermA.uplo)) - ∂A = ∂Atriu isa AbstractTriangular ? triu!(∂Atriu.data) : ∂Atriu - elseif ΔV isa AbstractZero + if ΔV isa AbstractZero ∂K = Diagonal(Δλ) - ∂A = V' \ ∂K * V' + ∂A = isherm ? V * ∂K * V' : V' \ ∂K * V' else ∂V = copyto!(similar(ΔV), ΔV) - _eigen_norm_phase_rev!(∂V, A, V) + _eigen_norm_phase_rev!(∂V, isherm ? Hermitian(A) : A, V) ∂K = V' * ∂V ∂K ./= λ' .- conj.(λ) + isherm && hermitianpart!(∂K) ∂K[diagind(∂K)] .= Δλ - ∂A = mul!(∂K, V' \ ∂K, V') + ∂A = isherm ? V * ∂K * V' : mul!(∂K, V' \ ∂K, V') end return NoTangent(), T <: Real ? real(∂A) : ∂A end - eigen_pullback(Ȳ::AbstractThunk) = eigen_pullback(unthunk(Ȳ)) - eigen_pullback(ΔF::AbstractZero) = (NoTangent(), ΔF) return F, eigen_pullback end @@ -410,21 +406,16 @@ end function frule((_, ΔA), ::typeof(eigvals!), A::StridedMatrix{T}; kwargs...) where {T<:BlasFloat} ΔA isa AbstractZero && return eigvals!(A; kwargs...), ΔA - if ishermitian(A) - λ, ∂λ = frule((ZeroTangent(), Hermitian(ΔA)), eigvals!, Hermitian(A)) - sortby = get(kwargs, :sortby, LinearAlgebra.eigsortby) - _sorteig!_fwd(∂λ, λ, sortby) + isherm = ishermitian(A) + F = eigen!(A; kwargs...) + λ, V = F.values, F.vectors + tmp = isherm ? V' * ΔA : V \ ΔA + ∂λ = similar(λ) + # diag(tmp * V) without computing full matrix product + if eltype(∂λ) <: Real + broadcast!((a, b) -> sum(real ∘ prod, zip(a, b)), ∂λ, eachrow(tmp), eachcol(V)) else - F = eigen!(A; kwargs...) - λ, V = F.values, F.vectors - tmp = V \ ΔA - ∂λ = similar(λ) - # diag(tmp * V) without computing full matrix product - if eltype(∂λ) <: Real - broadcast!((a, b) -> sum(real ∘ prod, zip(a, b)), ∂λ, eachrow(tmp), eachcol(V)) - else - broadcast!((a, b) -> sum(prod, zip(a, b)), ∂λ, eachrow(tmp), eachcol(V)) - end + broadcast!((a, b) -> sum(prod, zip(a, b)), ∂λ, eachrow(tmp), eachcol(V)) end return λ, ∂λ end @@ -440,17 +431,6 @@ function rrule(::typeof(eigvals), A::StridedMatrix{T}; kwargs...) where {T<:Unio return λ, eigvals_pullback end -# adapted from LinearAlgebra.sorteig! -function _sorteig!_fwd(Δλ, λ, sortby) - Δλ isa AbstractZero && return (sort!(λ; by=sortby), Δλ) - if sortby !== nothing - p = sortperm(λ; alg=QuickSort, by=sortby) - permute!(λ, p) - permute!(Δλ, p) - end - return (λ, Δλ) -end - ##### ##### `cholesky` ##### diff --git a/test/rulesets/LinearAlgebra/factorization.jl b/test/rulesets/LinearAlgebra/factorization.jl index f5cc8be7e..46fcc449f 100644 --- a/test/rulesets/LinearAlgebra/factorization.jl +++ b/test/rulesets/LinearAlgebra/factorization.jl @@ -266,13 +266,13 @@ end @test Ḟ_ad.vectors ≈ Ḟ_fd.vectors @test frule((ZeroTangent(), ZeroTangent()), eigen!, copy(X)) == (F, ZeroTangent()) - @testset "tangents are real when outputs are" begin - # hermitian matrices have real eigenvalues and, when real, real eigenvectors + @testset "tangent types match outputs" begin + # hermitian matrices have real eigenvalues (of complex type for complex + # matrices on Julia >= 1.14) and, when real, real eigenvectors X = Matrix(Hermitian(randn(T, n, n))) - Ẋ = Matrix(Hermitian(rand_tangent(X))) - _, Ḟ = frule((ZeroTangent(), Ẋ), eigen!, X) - @test eltype(Ḟ.values) <: Real - T <: Real && @test eltype(Ḟ.vectors) <: Real + F, Ḟ = frule((ZeroTangent(), rand_tangent(X)), eigen!, X) + @test eltype(Ḟ.values) == eltype(F.values) + @test eltype(Ḟ.vectors) == eltype(F.vectors) end end @@ -349,9 +349,12 @@ end return F end + # eigenvectors are only differentiable along Hermitian directions + _eigvecs_stable(x) = _eigen_stable(Matrix(hermitianpart(x))).vectors + n = 10 @testset "eigen!(::Matrix{$T})" for T in (Float64, ComplexF64) - A, ΔA = Matrix(Hermitian(randn(T, n, n))), Matrix(Hermitian(randn(T, n, n))) + A, ΔA = Matrix(Hermitian(randn(T, n, n))), randn(T, n, n) F = eigen!(copy(A)) @test frule((ZeroTangent(), ZeroTangent()), eigen!, copy(A)) == (F, ZeroTangent()) @@ -361,21 +364,18 @@ end @test ∂F_ad.values isa typeof(F.values) @test ∂F_ad.vectors isa typeof(F.vectors) - f = x -> asnt(eigen(Matrix(Hermitian(x)))) - ∂F_fd = jvp(_fdm, f, (A, ΔA)) - @test ∂F_ad.values ≈ ∂F_fd.values + # eigenvalues are differentiable in all directions their type admits + f_values = eltype(F.values) == T ? eigvals : x -> eigvals(Matrix(hermitianpart(x))) + @test ∂F_ad.values ≈ jvp(_fdm, f_values, (A, ΔA)) - f_stable = x -> asnt(_eigen_stable(Matrix(Hermitian(x)))) - F_stable = f_stable(A) - ∂F_stable_fd = jvp(_fdm, f_stable, (A, ΔA)) C = _eigvecs_stabilize_mat(F.vectors) - @test ∂F_ad.vectors * C ≈ ∂F_stable_fd.vectors + @test ∂F_ad.vectors * C ≈ jvp(_fdm, _eigvecs_stable, (A, ΔA)) end @testset "eigen(::Matrix{$T})" for T in (Float64, ComplexF64) - A, ΔU, Δλ = Matrix(Hermitian(randn(T, n, n))), randn(T, n, n), randn(n) - + A = Matrix(Hermitian(randn(T, n, n))) F = eigen(A) + ΔU, Δλ = randn(T, n, n), randn(eltype(F.values), n) ΔF = Tangent{typeof(F)}(; values=Δλ, vectors=ΔU) F_ad, back = rrule(eigen, A) @test F_ad == F @@ -383,13 +383,16 @@ end C = _eigvecs_stabilize_mat(F.vectors) CT = Tangent{typeof(F)} + # eigenvalues are differentiable in all directions their type admits + f_values = eltype(F.values) == T ? eigvals : x -> eigvals(Matrix(hermitianpart(x))) + @testset for nzprops in ([:values], [:vectors], [:values, :vectors]) ∂F = CT(; [s => getproperty(ΔF, s) for s in nzprops]...) ∂F_stable = (; [s => copy(getproperty(ΔF, s)) for s in nzprops]...) :vectors in nzprops && rmul!(∂F_stable.vectors, C) f_stable = function(x) - F_ = _eigen_stable(Matrix(Hermitian(x))) + F_ = (values=f_values(x), vectors=_eigvecs_stable(x)) return (; (s => getproperty(F_, s) for s in nzprops)...) end @@ -398,6 +401,7 @@ end @test ∂A isa typeof(A) ∂A_fd = j′vp(_fdm, f_stable, ∂F_stable, A)[1] @test ∂A ≈ ∂A_fd + @test @maybe_inferred(back(@thunk(∂F))) == (∂self, ∂A) end end end @@ -411,12 +415,12 @@ end test_frule(eigvals!, X) @test frule((ZeroTangent(), ZeroTangent()), eigvals!, copy(X))[2] == ZeroTangent() - @testset "tangents are real when outputs are" begin - # hermitian matrices have real eigenvalues + @testset "tangent types match outputs" begin + # hermitian matrices have real eigenvalues (of complex type for complex + # matrices on Julia >= 1.14) X = Matrix(Hermitian(randn(T, n, n))) - Ẋ = Matrix(Hermitian(rand_tangent(X))) - _, λ̇ = frule((ZeroTangent(), Ẋ), eigvals!, X) - @test eltype(λ̇) <: Real + λ, λ̇ = frule((ZeroTangent(), rand_tangent(X)), eigvals!, X) + @test eltype(λ̇) == eltype(λ) end end @@ -437,24 +441,32 @@ end @testset "hermitian matrices" begin n = 10 @testset "eigvals!(::Matrix{$T})" for T in (Float64, ComplexF64) - A, ΔA = Matrix(Hermitian(randn(T, n, n))), Matrix(Hermitian(randn(T, n, n))) + A, ΔA = Matrix(Hermitian(randn(T, n, n))), randn(T, n, n) λ = eigvals!(copy(A)) λ_ad, ∂λ_ad = frule((ZeroTangent(), copy(ΔA)), eigvals!, copy(A)) @test λ_ad ≈ λ # inexact because frule uses eigen not eigvals @test ∂λ_ad isa typeof(λ) - @test ∂λ_ad ≈ jvp(_fdm, A -> eigvals(Matrix(Hermitian(A))), (A, ΔA)) + # eigenvalues are differentiable in all directions their type admits + f = eltype(λ) == T ? eigvals : x -> eigvals(Matrix(hermitianpart(x))) + @test ∂λ_ad ≈ jvp(_fdm, f, (A, ΔA)) end @testset "eigvals(::Matrix{$T})" for T in (Float64, ComplexF64) - A, Δλ = Matrix(Hermitian(randn(T, n, n))), randn(n) + A = Matrix(Hermitian(randn(T, n, n))) λ = eigvals(A) + Δλ = randn(eltype(λ), n) λ_ad, back = rrule(eigvals, A) @test λ_ad ≈ λ # inexact because rrule uses eigen not eigvals ∂self, ∂A = @maybe_inferred back(Δλ) @test ∂self === NoTangent() @test ∂A isa typeof(A) - @test ∂A ≈ j′vp(_fdm, A -> eigvals(Matrix(Hermitian(A))), Δλ, A)[1] + # eigenvalues are differentiable in all directions their type admits + f = eltype(λ) == T ? eigvals : x -> eigvals(Matrix(hermitianpart(x))) + @test ∂A ≈ j′vp(_fdm, f, Δλ, A)[1] @test @maybe_inferred(back(ZeroTangent())) == (NoTangent(), ZeroTangent()) + + # gradient of `sum ∘ eigvals = tr` (FluxML/Zygote.jl#1369) + @test back(ones(eltype(λ), n))[2] ≈ I end end end From 5f23afecd8d781d37aeb0ffeb6fd8aa65c978238 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Fri, 23 Oct 2026 15:43:56 +0200 Subject: [PATCH 2/8] Unthunk tangents in eigen!/eigvals! frules and reduce allocations Reuse the cotangent buffer for the product with the eigenvectors in the `eigen` pullback for Hermitian matrices. Co-Authored-By: Claude Opus 5.5 --- src/rulesets/LinearAlgebra/factorization.jl | 10 ++++++---- test/rulesets/LinearAlgebra/factorization.jl | 4 ++++ 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/src/rulesets/LinearAlgebra/factorization.jl b/src/rulesets/LinearAlgebra/factorization.jl index fcab1dcbf..9ac1bea50 100644 --- a/src/rulesets/LinearAlgebra/factorization.jl +++ b/src/rulesets/LinearAlgebra/factorization.jl @@ -302,7 +302,8 @@ end # Hermitian part of the (co)tangents of the eigenvectors. The eigenvalues are differentiable # in all directions their type admits. -function frule((_, ΔA), ::typeof(eigen!), A::StridedMatrix{T}; kwargs...) where {T<:BlasFloat} +function frule((_, Ȧ), ::typeof(eigen!), A::StridedMatrix{T}; kwargs...) where {T<:BlasFloat} + ΔA = unthunk(Ȧ) ΔA isa AbstractZero && return (eigen!(A; kwargs...), ΔA) isherm = ishermitian(A) F = eigen!(A; kwargs...) @@ -332,7 +333,7 @@ function rrule(::typeof(eigen), A::StridedMatrix{T}; kwargs...) where {T<:Union{ ΔV isa AbstractZero && Δλ isa AbstractZero && return (NoTangent(), Δλ + ΔV) if ΔV isa AbstractZero ∂K = Diagonal(Δλ) - ∂A = isherm ? V * ∂K * V' : V' \ ∂K * V' + ∂A = (isherm ? V * ∂K : V' \ ∂K) * V' else ∂V = copyto!(similar(ΔV), ΔV) _eigen_norm_phase_rev!(∂V, isherm ? Hermitian(A) : A, V) @@ -340,7 +341,7 @@ function rrule(::typeof(eigen), A::StridedMatrix{T}; kwargs...) where {T<:Union{ ∂K ./= λ' .- conj.(λ) isherm && hermitianpart!(∂K) ∂K[diagind(∂K)] .= Δλ - ∂A = isherm ? V * ∂K * V' : mul!(∂K, V' \ ∂K, V') + ∂A = mul!(∂K, isherm ? V * ∂K : V' \ ∂K, V') end return NoTangent(), T <: Real ? real(∂A) : ∂A end @@ -404,7 +405,8 @@ end ##### `eigvals` ##### -function frule((_, ΔA), ::typeof(eigvals!), A::StridedMatrix{T}; kwargs...) where {T<:BlasFloat} +function frule((_, Ȧ), ::typeof(eigvals!), A::StridedMatrix{T}; kwargs...) where {T<:BlasFloat} + ΔA = unthunk(Ȧ) ΔA isa AbstractZero && return eigvals!(A; kwargs...), ΔA isherm = ishermitian(A) F = eigen!(A; kwargs...) diff --git a/test/rulesets/LinearAlgebra/factorization.jl b/test/rulesets/LinearAlgebra/factorization.jl index 46fcc449f..78455eadf 100644 --- a/test/rulesets/LinearAlgebra/factorization.jl +++ b/test/rulesets/LinearAlgebra/factorization.jl @@ -363,6 +363,9 @@ end @test ∂F_ad isa Tangent{typeof(F)} @test ∂F_ad.values isa typeof(F.values) @test ∂F_ad.vectors isa typeof(F.vectors) + _, ∂F_thunk = frule((ZeroTangent(), @thunk(copy(ΔA))), eigen!, copy(A)) + @test ∂F_thunk.values == ∂F_ad.values + @test ∂F_thunk.vectors == ∂F_ad.vectors # eigenvalues are differentiable in all directions their type admits f_values = eltype(F.values) == T ? eigvals : x -> eigvals(Matrix(hermitianpart(x))) @@ -446,6 +449,7 @@ end λ_ad, ∂λ_ad = frule((ZeroTangent(), copy(ΔA)), eigvals!, copy(A)) @test λ_ad ≈ λ # inexact because frule uses eigen not eigvals @test ∂λ_ad isa typeof(λ) + @test frule((ZeroTangent(), @thunk(copy(ΔA))), eigvals!, copy(A))[2] == ∂λ_ad # eigenvalues are differentiable in all directions their type admits f = eltype(λ) == T ? eigvals : x -> eigvals(Matrix(hermitianpart(x))) @test ∂λ_ad ≈ jvp(_fdm, f, (A, ΔA)) From d2005d5ea2096d05a05463b4b4b0478e064bc741 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Fri, 23 Oct 2026 23:27:46 +0200 Subject: [PATCH 3/8] Apply batched suggestions from code review Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- src/rulesets/LinearAlgebra/factorization.jl | 8 ++++++-- test/rulesets/LinearAlgebra/factorization.jl | 15 ++++++++++++--- 2 files changed, 18 insertions(+), 5 deletions(-) diff --git a/src/rulesets/LinearAlgebra/factorization.jl b/src/rulesets/LinearAlgebra/factorization.jl index 9ac1bea50..ff9545a5c 100644 --- a/src/rulesets/LinearAlgebra/factorization.jl +++ b/src/rulesets/LinearAlgebra/factorization.jl @@ -302,7 +302,9 @@ end # Hermitian part of the (co)tangents of the eigenvectors. The eigenvalues are differentiable # in all directions their type admits. -function frule((_, Ȧ), ::typeof(eigen!), A::StridedMatrix{T}; kwargs...) where {T<:BlasFloat} +function frule( + (_, Ȧ), ::typeof(eigen!), A::StridedMatrix{T}; kwargs... +) where {T<:BlasFloat} ΔA = unthunk(Ȧ) ΔA isa AbstractZero && return (eigen!(A; kwargs...), ΔA) isherm = ishermitian(A) @@ -405,7 +407,9 @@ end ##### `eigvals` ##### -function frule((_, Ȧ), ::typeof(eigvals!), A::StridedMatrix{T}; kwargs...) where {T<:BlasFloat} +function frule( + (_, Ȧ), ::typeof(eigvals!), A::StridedMatrix{T}; kwargs... +) where {T<:BlasFloat} ΔA = unthunk(Ȧ) ΔA isa AbstractZero && return eigvals!(A; kwargs...), ΔA isherm = ishermitian(A) diff --git a/test/rulesets/LinearAlgebra/factorization.jl b/test/rulesets/LinearAlgebra/factorization.jl index 78455eadf..188597224 100644 --- a/test/rulesets/LinearAlgebra/factorization.jl +++ b/test/rulesets/LinearAlgebra/factorization.jl @@ -368,7 +368,11 @@ end @test ∂F_thunk.vectors == ∂F_ad.vectors # eigenvalues are differentiable in all directions their type admits - f_values = eltype(F.values) == T ? eigvals : x -> eigvals(Matrix(hermitianpart(x))) + f_values = if eltype(F.values) == T + eigvals + else + x -> eigvals(Matrix(hermitianpart(x))) + end @test ∂F_ad.values ≈ jvp(_fdm, f_values, (A, ΔA)) C = _eigvecs_stabilize_mat(F.vectors) @@ -387,7 +391,11 @@ end CT = Tangent{typeof(F)} # eigenvalues are differentiable in all directions their type admits - f_values = eltype(F.values) == T ? eigvals : x -> eigvals(Matrix(hermitianpart(x))) + f_values = if eltype(F.values) == T + eigvals + else + x -> eigvals(Matrix(hermitianpart(x))) + end @testset for nzprops in ([:values], [:vectors], [:values, :vectors]) ∂F = CT(; [s => getproperty(ΔF, s) for s in nzprops]...) @@ -449,7 +457,8 @@ end λ_ad, ∂λ_ad = frule((ZeroTangent(), copy(ΔA)), eigvals!, copy(A)) @test λ_ad ≈ λ # inexact because frule uses eigen not eigvals @test ∂λ_ad isa typeof(λ) - @test frule((ZeroTangent(), @thunk(copy(ΔA))), eigvals!, copy(A))[2] == ∂λ_ad + @test frule((ZeroTangent(), @thunk(copy(ΔA))), eigvals!, copy(A))[2] == + ∂λ_ad # eigenvalues are differentiable in all directions their type admits f = eltype(λ) == T ? eigvals : x -> eigvals(Matrix(hermitianpart(x))) @test ∂λ_ad ≈ jvp(_fdm, f, (A, ΔA)) From 322181420623a4b5819ebbc8e7a986e855819e43 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Sat, 24 Oct 2026 01:08:59 +0200 Subject: [PATCH 4/8] Bump version from 1.74.3 to 1.74.4 --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index aeb76428a..9cf6ab670 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "ChainRules" uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2" -version = "1.74.3" +version = "1.74.4" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" From 32a50cd26b04ac17b02d61f13fa2e831cddf207c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Sat, 24 Oct 2026 01:10:29 +0200 Subject: [PATCH 5/8] Apply suggestion from @devmotion --- src/rulesets/LinearAlgebra/factorization.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/rulesets/LinearAlgebra/factorization.jl b/src/rulesets/LinearAlgebra/factorization.jl index ff9545a5c..aafad5bb2 100644 --- a/src/rulesets/LinearAlgebra/factorization.jl +++ b/src/rulesets/LinearAlgebra/factorization.jl @@ -327,7 +327,7 @@ end function rrule(::typeof(eigen), A::StridedMatrix{T}; kwargs...) where {T<:Union{Real,Complex}} isherm = ishermitian(A) F = eigen(A; kwargs...) - function eigen_pullback(ΔF::ChainRulesCore.AbstractTangent) + function eigen_pullback(ΔF) _ΔF = unthunk(ΔF) _ΔF isa AbstractZero && return (NoTangent(), _ΔF) λ, V = F.values, F.vectors From 8836bf6610106f8173f6134b72b3952d7fedf3a0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Sun, 25 Oct 2026 22:07:49 +0100 Subject: [PATCH 6/8] Project eigenvalue cotangents and test sortby for Hermitian eigen rules Cotangents of real-typed eigenvalues are projected onto the reals again, as in the previous Hermitian code path. Adds tests that (co)tangents follow the order of the eigenpairs for custom `sortby`. Co-Authored-By: Claude Opus 5.5 --- src/rulesets/LinearAlgebra/factorization.jl | 15 +++++---- test/rulesets/LinearAlgebra/factorization.jl | 34 ++++++++++++++++++++ 2 files changed, 42 insertions(+), 7 deletions(-) diff --git a/src/rulesets/LinearAlgebra/factorization.jl b/src/rulesets/LinearAlgebra/factorization.jl index aafad5bb2..669c9a717 100644 --- a/src/rulesets/LinearAlgebra/factorization.jl +++ b/src/rulesets/LinearAlgebra/factorization.jl @@ -296,11 +296,10 @@ end # - simplify when A is diagonal # - support degenerate matrices (see #144) -# NOTE: if `A` is Hermitian, `eigen` uses the solver for Hermitian matrices. Its eigenvectors -# are unitary and follow a different phase convention than the eigenvectors of non-Hermitian -# matrices, so they are only differentiable along Hermitian directions: the rules use the -# Hermitian part of the (co)tangents of the eigenvectors. The eigenvalues are differentiable -# in all directions their type admits. +# NOTE: for Hermitian `A`, `eigen` uses the Hermitian solver, whose eigenvectors are unitary +# but follow a different phase convention. They are only differentiable along Hermitian +# directions, so their (co)tangents use the Hermitian part of `V' * ΔA * V`. The eigenvalues +# are differentiable in all directions their type admits. function frule( (_, Ȧ), ::typeof(eigen!), A::StridedMatrix{T}; kwargs... @@ -318,7 +317,7 @@ function frule( ∂K ./= transpose(λ) .- λ fill!(∂Kdiag, 0) ∂V = mul!(tmp, V, ∂K) - # `A` is overwritten but the phase convention depends only on its type + # `A` is overwritten but the phase convention depends only on its type and `uplo` _eigen_norm_phase_fwd!(∂V, isherm ? Hermitian(A) : A, V) ∂F = Tangent{typeof(F)}(values = ∂λ, vectors = ∂V) return F, ∂F @@ -327,11 +326,13 @@ end function rrule(::typeof(eigen), A::StridedMatrix{T}; kwargs...) where {T<:Union{Real,Complex}} isherm = ishermitian(A) F = eigen(A; kwargs...) + # cotangents of real-typed eigenvalues are real + project_λ = ProjectTo(F.values) function eigen_pullback(ΔF) _ΔF = unthunk(ΔF) _ΔF isa AbstractZero && return (NoTangent(), _ΔF) λ, V = F.values, F.vectors - Δλ, ΔV = _ΔF.values, _ΔF.vectors + Δλ, ΔV = project_λ(_ΔF.values), _ΔF.vectors ΔV isa AbstractZero && Δλ isa AbstractZero && return (NoTangent(), Δλ + ΔV) if ΔV isa AbstractZero ∂K = Diagonal(Δλ) diff --git a/test/rulesets/LinearAlgebra/factorization.jl b/test/rulesets/LinearAlgebra/factorization.jl index 188597224..0b7fcc725 100644 --- a/test/rulesets/LinearAlgebra/factorization.jl +++ b/test/rulesets/LinearAlgebra/factorization.jl @@ -377,6 +377,19 @@ end C = _eigvecs_stabilize_mat(F.vectors) @test ∂F_ad.vectors * C ≈ jvp(_fdm, _eigvecs_stable, (A, ΔA)) + + # tangents follow the order of the eigenpairs + _, ∂F_nothing = frule( + (ZeroTangent(), copy(ΔA)), eigen!, copy(A); sortby=nothing + ) + @test ∂F_nothing.values ≈ ∂F_ad.values + @test ∂F_nothing.vectors ≈ ∂F_ad.vectors + F_rev, ∂F_rev = frule( + (ZeroTangent(), copy(ΔA)), eigen!, copy(A); sortby=λ -> -real(λ) + ) + @test F_rev.values == reverse(F.values) + @test ∂F_rev.values ≈ reverse(∂F_ad.values) + @test ∂F_rev.vectors ≈ reverse(∂F_ad.vectors; dims=2) end @testset "eigen(::Matrix{$T})" for T in (Float64, ComplexF64) @@ -414,6 +427,12 @@ end @test ∂A ≈ ∂A_fd @test @maybe_inferred(back(@thunk(∂F))) == (∂self, ∂A) end + + # cotangents follow the order of the eigenpairs + F_rev, back_rev = rrule(eigen, A; sortby=λ -> -real(λ)) + @test F_rev.values == reverse(F.values) + ΔF_rev = CT(; values=reverse(Δλ), vectors=reverse(ΔU; dims=2)) + @test back_rev(ΔF_rev)[2] ≈ back(ΔF)[2] end end end @@ -462,6 +481,12 @@ end # eigenvalues are differentiable in all directions their type admits f = eltype(λ) == T ? eigvals : x -> eigvals(Matrix(hermitianpart(x))) @test ∂λ_ad ≈ jvp(_fdm, f, (A, ΔA)) + + # tangents follow the order of the eigenvalues + _, ∂λ_rev = frule( + (ZeroTangent(), copy(ΔA)), eigvals!, copy(A); sortby=λ -> -real(λ) + ) + @test ∂λ_rev ≈ reverse(∂λ_ad) end @testset "eigvals(::Matrix{$T})" for T in (Float64, ComplexF64) @@ -480,6 +505,15 @@ end # gradient of `sum ∘ eigvals = tr` (FluxML/Zygote.jl#1369) @test back(ones(eltype(λ), n))[2] ≈ I + + # imaginary parts of cotangents of real-typed eigenvalues are ignored + if eltype(λ) <: Real + @test back(complex.(Δλ, randn(n)))[2] == ∂A + end + + # cotangents follow the order of the eigenvalues + _, back_rev = rrule(eigvals, A; sortby=λ -> -real(λ)) + @test back_rev(reverse(Δλ))[2] ≈ ∂A end end end From 5afd369f3deba33a29d42721e0c6f74ea9998bfb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Sun, 25 Oct 2026 22:13:22 +0100 Subject: [PATCH 7/8] Bump version from 1.74.4 to 1.74.5 --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 9cf6ab670..5ef530989 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "ChainRules" uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2" -version = "1.74.4" +version = "1.74.5" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" From 01fc267652c00b1102cbf272a9806e7a5125b15f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Mon, 26 Oct 2026 13:19:49 +0100 Subject: [PATCH 8/8] Bump version from 1.74.5 to 1.74.6 --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 5ef530989..4bf00b042 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "ChainRules" uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2" -version = "1.74.5" +version = "1.74.6" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"