Fix eigen/eigvals rules for Hermitian StridedMatrix

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"

src/rulesets/LinearAlgebra/factorization.jl

@@ -293,61 +293,61 @@ end
 #####
 
 # TODO:
-# - support correct differential of phase convention when A is hermitian
 # - simplify when A is diagonal
 # - support degenerate matrices (see #144)
 
-function frule((_, ΔA), ::typeof(eigen!), A::StridedMatrix{T}; kwargs...) where {T<:BlasFloat}
+# 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...
+) where {T<:BlasFloat}
+    ΔA = unthunk(Ȧ)
     Δ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 and `uplo`
+    _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)
+    # 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 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' \ ∂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 = mul!(∂K, isherm ? V * ∂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
 
@@ -408,23 +408,21 @@ 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
-    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 +438,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`
 #####

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())
@@ -360,36 +363,60 @@ 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
 
-                    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 = if eltype(F.values) == T
+                        eigvals
+                    else
+                        x -> eigvals(Matrix(hermitianpart(x)))
+                    end
+                    @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))
+
+                    # 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)
-                    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
 
                     C = _eigvecs_stabilize_mat(F.vectors)
                     CT = Tangent{typeof(F)}
 
+                    # eigenvalues are differentiable in all directions their type admits
+                    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]...)
                         ∂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,7 +425,14 @@ 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
+
+                    # 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
@@ -411,12 +445,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 +471,49 @@ 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))
+                    @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))
+
+                    # 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)
-                    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
+
+                    # 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