Wrong Jacobian

[Chen Zhao]

using Yao
include("chainrules_patch.jl")
using Zygote
circ = put(1, 1=> Rx(0))
function loss(circ, θ)
    circ = dispatch!(circ, [θ])
    ψ = zero_state(1)
    ψ = apply!(ψ, circ)
    return abs2.(state(ψ))
end
function loss2(circ, θ)
    circ = dispatch!(circ, [θ]) # this works
    ψ = zero_state(1)
    ψ = apply!(ψ, circ)
    return abs2.(state(ψ))[2]
end

The gradient at pi should be zero.

julia> gradient(x -> loss2(circ, x), pi)
(6.123233995736766e-17,)

However, the Jacobian is wrong. Here, the second element should be approximate to zero.

julia> jacobian(x -> loss(circ, x), pi)
([-6.123233995736766e-17, 1.0],)

[Jinguo Liu (刘金国)]

julia> using Yao

julia> include("chainrules_patch.jl")


julia> using Zygote

julia> circ = put(1, 1=> Rx(0))
nqubits: 1
put on (1)
└─ rot(X, 0.0)

julia> function loss(circ, θ)
           circ = dispatch!(circ, [θ])
           ψ = zero_state(1)
           ψ = apply!(ψ, circ)
           return abs2.(state(ψ))
       end
loss (generic function with 1 method)

julia> function loss2(circ, θ)
           circ = dispatch!(circ, [θ]) # this works
           ψ = zero_state(1)
           ψ = apply!(ψ, circ)
           return abs2.(state(ψ))[2]
       end
loss2 (generic function with 1 method)

julia> gradient(x -> loss2(circ, x), pi)
(6.123233995736766e-17,)

julia> jacobian(x -> loss2(circ, x), pi)
([6.123233995736766e-17],)

Hmm, please definitely check your zygote version!

[Chen Zhao]

I checked my Zygote ver. which is 0.6.12.
The error is in Jacobian of loss not loss2.

julia> jacobian(x -> loss(circ, x), pi)
([-6.123233995736766e-17, 1.0],)

It should be approximate to ([0, 0], ).

[Chen Zhao]

I updated Zygote to 0.6.19. And this issue is still there.

[Jinguo Liu (刘金国)]

I just confirmed your issue. But I think it is probably an issue of Zygote's withjacobian function.

function withjacobian(f, args...)
  y, back = pullback(_jvec∘f, args...)
  out = map(args) do x
    T = promote_type(eltype(x), eltype(y))
    dx = x isa AbstractArray ? similar(x, T, length(y), length(x)) :
      x isa Number ? similar(y, T, length(y)) :
      nothing
  end
  delta = _eyelike(y)
   @show delta
    @show back(delta[:,2])
    @show back(delta[:,1])
    @show back(delta[:,2])
  for k in LinearIndices(y)
    grads = back(delta[:,k])
    for (dx, grad) in zip(out, grads)
      dx isa AbstractArray || continue
      _gradcopy!(view(dx,k,:), grad)
    end
  end
  (val=y, grad=out)
end

I tried the following order of computing gradients,

julia> jacobian(x -> loss(circ, x), pi)
delta = [1.0 0.0; 0.0 1.0]
back(delta[:, 2]) = (6.123233995736766e-17,)
back(delta[:, 1]) = (-0.0,)
([6.123233995736766e-17, -1.0],)

julia> jacobian(x -> loss(circ, x), pi)
delta = [1.0 0.0; 0.0 1.0]
back(delta[:, 2]) = (6.123233995736766e-17,)
back(delta[:, 1]) = (-0.0,)
back(delta[:, 2]) = (6.123233995736766e-17,)
([7.498798913309288e-33, -1.8369701987210297e-16],)

julia> jacobian(x -> loss(circ, x), pi)
delta = [1.0 0.0; 0.0 1.0]
back(delta[:, 1]) = (-6.123233995736766e-17,)
back(delta[:, 2]) = (1.0,)
([6.123233995736766e-17, -1.0],)

My theory is, during computing the first gradient, the cached data is destroyed. The quick fix for you is just compute column by column, each time re-compute cached states. We really need an expert to look into this.

[Jinguo Liu (刘金国)]

BTW, do you have such issue when using the previous patch for zygote?

[Jinguo Liu (刘金国)]

You can also try the following approach to fix your issue

1. implement non-inplace version dispatch and apply,
2. rewrite rules for these functions,
3. submit a PR to QuAlgorithmZoo if you have time. :P

then it should works

[Jinguo Liu (刘金国)]

Just pushed a fix to this patch. Now the Jacobian should work. You need to update YaoBlocks to v0.11.5. to make it work. Now, the patch is implemented on apply and dispatch functions.

[Chen Zhao]

@GiggleLiu Thanks! It has been fixed.