A bug about degenerated eigsolve ad

[XingyuZhang2018]

There is a bug in the current AD implementation for eigsolve when the eigenvalues are degenerate.

A simple example

using KrylovKit
using Zygote

@testset "degenerated eigsolve ad" begin
    function num_grad(f, K::Number; δ::Real=1e-5)
        if eltype(K) == ComplexF64
            (f(K + δ / 2) - f(K - δ / 2)) / δ + 
                (f(K + δ / 2 * 1.0im) - f(K - δ / 2 * 1.0im)) / δ * 1.0im
        else
            (f(K + δ / 2) - f(K - δ / 2)) / δ
        end
    end

    function num_grad(f, a::AbstractArray; δ::Real=1e-5)
        b = copy(a)
        df = map(CartesianIndices(b)) do i
            foo = x -> (ac = copy(b); ac[i] = x; f(ac))
            num_grad(foo, b[i], δ=δ)
        end
        return df
    end

    Random.seed!(42)
    M = [zeros(2,2) rand(ComplexF64, 2,2); rand(ComplexF64, 2,2) zeros(2,2)]
    N = rand(ComplexF64, 4, 4)

    function foo(M)
        λ, V = eigsolve(M, 1, :LM)
        p = real(λ[1]) > 0 ? 1 : 2
        V1 = V[p]
        return norm(V1' * N * V1)
    end
    # @show Zygote.gradient(foo, M)[1] num_grad(foo, M)
    @show norm(Zygote.gradient(foo, M)[1] - num_grad(foo, M))
end

Even fixing eigvector doesn't get the correct gradient.

[Jutho]

For now, you can fix it by using:

λ, V = eigsolve(M, 1, :LM; alg_rrule=GMRES())

where you might want to tune GMRES some more with
GMRES(krylovdim = ..., maxiter=...)

I will look into fixing this more generally.

[Jutho]

This should be fixed now also for Arnoldi() by #99

[Jutho]

I will tag a new version 0.8.2 now, so maybe you can check whether this now fully solves the problem. Also, I am interested in benchmark comparisons between Arnoldi and GMRES for your usecase.

[XingyuZhang2018]

Thank you! I will try and give feedback.