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" diff --git a/src/rulesets/LinearAlgebra/factorization.jl b/src/rulesets/LinearAlgebra/factorization.jl index e04bd0e21..669c9a717 100644 --- a/src/rulesets/LinearAlgebra/factorization.jl +++ b/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` ##### diff --git a/test/rulesets/LinearAlgebra/factorization.jl b/test/rulesets/LinearAlgebra/factorization.jl index f5cc8be7e..0b7fcc725 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()) @@ -360,22 +363,39 @@ 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 @@ -383,13 +403,20 @@ end 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