Incorrect function evaluation when swapping variables

[goerz]

The following example fails with Zygote 0.6.37, see also Discourse:

using Zygote
using LinearAlgebra
using Distributions
using Test

function random_hermitian_matrix(N, specrad=1.0)
    σ = 1 / √N
    d = Normal(0.0, σ)
    X = (rand(d, (N, N)) + rand(d, (N, N)) * 1im) / √2
    H = specrad * (X + X') / (2 * √2)
end


function random_state_vector(N)
    Ψ = rand(N) .* exp.((2π * im) .* rand(N))
    Ψ ./= norm(Ψ)
    return Ψ
end


function cheby(Ψ, H, dt)
    a = [0.9915910021578431, 0.18282201929219635, 0.008403088661031203, 0.000257307553262815]
    Δ = 6.0; E_min = -3.0
    β = (Δ / 2) + E_min
    c = -2im / Δ

    v0 = Ψ
    ϕ = a[1] * v0
    v1 = c * (H * v0 - β * v0)
    ϕ = ϕ + a[2] * v1

    c *= 2
    for i = 3:length(a)
        v2 = c * (H * v1 - β * v1) + v0
        ϕ = ϕ + a[i] * v2
        v0, v1, v2 = v1, v2, v0                                # doesn't work
        #aux = v0; v0 = v1; v1 = v2; v2 = aux                   # doesn't work
        #aux = 1 * v0; v0 = 1 * v1; v1 = 1 * v2; v2 = 1 * aux   # works
    end

    return exp(-1im * β * dt) * ϕ
end

N = 2
dt = 0.1

Ψ0 = random_state_vector(N)
Ψ1 = random_state_vector(N)
H0 = random_hermitian_matrix(N)
H1 = random_hermitian_matrix(N)

ϵ = 1.0
res1 = abs2(Ψ1 ⋅ cheby(Ψ0, H0 + ϵ * H1, dt))
res2, grad = Zygote.withgradient(ϵ -> abs2(Ψ1 ⋅ cheby(Ψ0, H0 + ϵ * H1, dt)), ϵ)

@test abs(res1 - res2) < 1e-12