Symmetric/Hermitian matrix function rules

src/rulesets/LinearAlgebra/structured.jl

@@ -47,6 +47,225 @@ end
 _symmetric_back(ΔΩ) = UpperTriangular(ΔΩ) + LowerTriangular(ΔΩ)' - Diagonal(ΔΩ)
 _symmetric_back(ΔΩ::Union{Diagonal,UpperTriangular}) = ΔΩ
 
+#####
+##### `Symmetric{<:Real}`/`Hermitian` eigendecomposition
+#####
+
+function frule((_, ΔA), ::typeof(eigen), A::LinearAlgebra.RealHermSymComplexHerm)
+    F = eigen(A)
+    ∂F = Thunk() do
+        λ, U = F
+        M = U' * ΔA * U
+        ∂λ = real(diag(M)) # if ΔA is Hermitian, so is U′ΔAU, so its diagonal is real
+        # K is skew-hermitian with zero diag
+        K = M ./ _nonzero.(λ' .- λ)
+        _setdiag!(K, Zero())
+        ∂U = U * K
+        return Composite{typeof(F)}(values = ∂λ, vectors = ∂U)
+    end
+    return F, ∂F
+end
+
+function rrule(::typeof(eigen), A::LinearAlgebra.RealHermSymComplexHerm)
+    F = eigen(A)
+    function eigen_pullback(ΔF)
+        ∂A = Thunk() do
+            λ, U = F
+            ∂λ, ∂U = ΔF.values, ΔF.vectors
+            # K is skew-hermitian
+            K = U' * ∂U
+            # unstable for degenerate matrices
+            U′∂AU = K ./ (_nonzero.(λ' .- λ))
+            setdiag!(U'∂AU, ∂λ)
+            return _symherm(A, U * U′∂AU * U')
+        end
+        return NO_FIELDS, ∂A
+    end
+    return F, eigen_pullback
+end
+
+function frule((_, ΔA), ::typeof(eigvals), A::LinearAlgebra.RealHermSymComplexHerm)
+    λ, U = eigen(A)
+    return λ, @thunk real(diag(U' * ΔA * U))
+end
+
+function rrule(::typeof(eigvals), A::LinearAlgebra.RealHermSymComplexHerm)
+    F, back = rrule(eigen, A)
+    λ, U = F
+    function eigvals_pullback(Δλ)
+        ∂A = Thunk() do
+            ∂F = Composite{typeof(F)}(values = Δλ)
+            return unthunk(back(∂F)[2])
+        end
+        return NO_FIELDS, ∂A
+    end
+    return λ, eigvals_pullback
+end
+
+_nonzero(x) = abs(x) > eps(eltype(x)) ? x : eps(eltype(x)) * sign(x)
+
+_setdiag!(A, d) = (A[diagind(A)] = d)
+_setdiag!(A, d::AbstractZero) = (A[diagind(A)] .= 0)
+_setdiag!(A::AbstractZero, d) = Diagonal(d)
+
+_pureimag(x) = x - real(x)
+
+function _realifydiag!(A)
+    for i in diagind(A)
+        @inbounds A[i] = real(A[i])
+    end
+    return A
+end
+_realifydiag!(A::LinearAlgebra.RealHermSym) = A
+
+_realifydiag(A) = A - _pureimag(Diagonal(A))
+_realifydiag(A::LinearAlgebra.RealHermSym) = A
+
+# constrain B to have same Symmetric/Hermitian type as A
+_symherm(A::Symmetric, B) = Symmetric(B, Symbol(A.uplo))
+_symherm(A::Hermitian, B) = Hermitian(B, Symbol(A.uplo))
+_symherm(f, B::AbstractMatrix{<:Real}, uplo) = Symmetric(B, uplo)
+_symherm(f, B::AbstractMatrix{<:Complex}, uplo) = Hermitian(B, uplo)
+
+# call _symherm but enforce diagonal constraint
+function _symherm!(args...)
+    S = _symherm(args...)
+    _realifydiag!(S)
+    return S
+end
+
+#####
+##### `Symmetric{<:Real}`/`Hermitian` power series functions
+#####
+
+# Currently only defined for series functions whose codomain is ℝ
+# These are type-stable and closed under `func`
+
+# The efficient way to do this is probably to AD Base.power_by_squaring
+function frule((_, ΔA, _), ::typeof(^), A::LinearAlgebra.RealHermSymComplexHerm, p::Integer)
+    λ, U = eigen(A)
+    λᵖ = λ .^ p
+    Y = _symherm!(^, U * Diagonal(λᵖ) * U', :U)
+    ∂Y = Thunk() do
+        dλᵖ_dλ = p .* λ .^ (p - 1)
+        ∂Λ = U' * ΔA * U
+        U′∂YU = _muldiffquotmat(^, λ, λᵖ, dλᵖ_dλ, ∂Λ)
+        return _symherm!(Y, U * U′∂YU * U')
+    end
+    return Y, ∂Y
+end
+
+function rrule(::typeof(^), A::LinearAlgebra.RealHermSymComplexHerm, p::Integer)
+    λ, U = eigen(A)
+    λᵖ = λ .^ p
+    Y = _symherm!(^, U * Diagonal(λᵖ) * U', :U)
+    function pow_pullback(ΔY)
+        ∂A = Thunk() do
+            dλᵖ_dλ = p .* λ .^ (p - 1)
+            ∂Λᵖ = U' * _realifydiag(ΔY) * U
+            # TODO: make sure that the `conj` is needed
+            U′∂AU = _muldiffquotmat(^, λ, λᵖ, dλᵖ_dλ, ∂Λᵖ)
+            return _symherm(A, U * U′∂AU * U')
+        end
+        return NO_FIELDS, ∂A, DoesNotExist()
+    end
+    return Y, pow_pullback
+end
+
+# TODO: support log, sqrt, acos, asin, and non-int pow, which are type-unstable
+for func in (:exp, :cos, :sin, :tan, :cosh, :sinh, :tanh, :atan, :asinh, :atanh)
+    @eval begin
+        function frule((_, ΔA), ::typeof($func), A::LinearAlgebra.RealHermSymComplexHerm)
+            df = λi -> frule((Zero(), One()), $func, λi)
+            λ, U = eigen(A)
+            fλ_df_dλ = (df).(λ)
+            fλ = first.(fλ_df_dλ)
+            Y = _symherm!($func, U * Diagonal(fλ) * U', :U)
+            ∂Y = Thunk() do
+                df_dλ = last.(unthunk.(fλ_df_dλ))
+                ∂Λ = U' * ΔA * U
+                U′∂YU = _muldiffquotmat($func, λ, fλ, df_dλ, ∂Λ)
+                return _symherm!(Y, U * U′∂YU * U')
+            end
+            return Y, ∂Y
+        end
+
+        function rrule(::typeof($func), A::LinearAlgebra.RealHermSymComplexHerm)
+            df = λi -> frule((Zero(), One()), $func, λi)
+            λ, U = eigen(A)
+            fλ_df_dλ = (df).(λ)
+            fλ = first.(fλ_df_dλ)
+            Y = _symherm!($func, U * Diagonal(fλ) * U', :U)
+            function $(Symbol("$(func)_pullback"))(ΔY)
+                ∂A = Thunk() do
+                    df_dλ = unthunk.(last.(fλ_df_dλ))
+                    ∂fΛ = U' * _realifydiag(ΔY) * U
+                    U′∂AU = _muldiffquotmat($func, λ, fλ, df_dλ, ∂fΛ)
+                    return _symherm(A, U * U′∂AU * U')
+                end
+                return NO_FIELDS, ∂A
+            end
+            return Y, $(Symbol("$(func)_pullback"))
+        end
+    end
+end
+
+function frule((_, ΔA), ::typeof(sincos), A::LinearAlgebra.RealHermSymComplexHerm)
+    λ, U = eigen(A)
+    sinλ, cosλ = sin.(λ), cos.(λ)
+    sinA = _symherm!(A, U * Diagonal(sinλ) * U', :U)
+    cosA = _symherm!(A, U * Diagonal(cosλ) * U', :U)
+    sincosA = (sinA, cosA)
+    ∂sincosA = Thunk() do
+        ∂Λ = U' * ΔA * U
+        U′∂sinAU = _muldiffquotmat(sin, λ, sinλ, cosλ, ∂Λ)
+        ∂sinA = _symherm!(sinA, U * U′∂sinAU * U')
+        U′∂cosAU = _muldiffquotmat(cos, λ, cosλ, -sinλ, ∂Λ)
+        ∂cosA = _symherm!(cosA, U * U′∂cosAU * U')
+        return Composite{typeof(sincosA)}(∂sinA, ∂cosA)
+    end
+    return sincosA, ∂sincosA
+end
+
+function rrule(::typeof(sincos), A::LinearAlgebra.RealHermSymComplexHerm)
+    λ, U = eigen(A)
+    sinλ, cosλ = sin.(λ), cos.(λ)
+    sinA = _symherm!(A, U * Diagonal(sinλ) * U', :U)
+    cosA = _symherm!(A, U * Diagonal(cosλ) * U', :U)
+    sincosA = (sinA, cosA)
+    function sincos_pullback(ΔsincosA)
+        ∂A = Thunk() do
+            ΔsinA, ΔcosA = ΔsincosA
+            ∂sinΛ, ∂cosΛ = U' * _realifydiag(ΔsinA) * U, U' * _realifydiag(ΔcosA) * U
+            inds = eachindex(λ)
+            U′∂AU = @inbounds begin
+                _diffquot.(sin, inds, inds', Ref(λ), Ref(sinλ), Ref(cosλ)) .* ∂sinΛ .+
+                _diffquot.(cos, inds, inds', Ref(λ), Ref(cosλ), Ref(-sinλ)) .* ∂cosΛ
+            end
+            return _symherm(A, U * U′∂AU * U')
+        end
+        return NO_FIELDS, ∂A
+    end
+    return sincosA, sincos_pullback
+end
+
+# difference quotient, i.e. Pᵢⱼ = (f(λᵢ) - f(λⱼ)) / (λᵢ - λⱼ), with f'(λᵢ) when i==j
+Base.@propagate_inbounds function _diffquot(f, i, j, λ, fλ, df_dλ)
+    i == j && return df_dλ[i]
+    Δλ = λ[i] - λ[j]
+    T = real(eltype(λ))
+    # Handle degenerate eigenvalues by taylor expanding Δfλ / Δλ as Δλ → 0
+    abs2(Δλ) < eps(T) && return (df_dλ[i] + df_dλ[j]) / 2
+    Δfλ = fλ[i] - fλ[j]
+    return Δfλ / Δλ
+end
+
+# multiply Δ by the matrix of difference quotients P
+function _muldiffquotmat(f, λ, fλ, df_dλ, Δ)
+    inds = eachindex(λ)
+    return @inbounds _diffquot.(f, inds, inds', Ref(λ), Ref(fλ), Ref(df_dλ)) .* Δ
+end
+
 #####
 ##### `Adjoint`
 #####