I have a relatively concise routine
cheby that evaluates
Φ = exp(-i H dt) Ψ for a Hermitian matrix
H and a complex vector
Ψ by expanding the exponential into Cheybchev polynomials. A minimal working example adapted from the
cheby routine in QuantumPropagators.jl is pasted below.
I’m finding that I get different results if I run
cheby directly or inside
Zygote.withgradient (where the Zygote result is wrong). I’ve managed to pinpoint it to one specific line that does a cyclic permutation of three vectors,
v0, v1, v2 = v1, v2, v0. It works if I rewrite it as
aux = 1 * v0; v0 = 1 * v1; v1 = 1 * v2; v2 = 1 * aux (and only if I have the factor
1 in there).
Does anybody have any insight into what is going on here? Should I simply report this as a bug in Zygote?
The full MWE is this:
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 * v0 v1 = c * (H * v0 - β * v0) ϕ = ϕ + a * 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