Calculating Jacobian using Zygote

[Radoica Draškić]

I'm trying to calculate the Jacobian for an example that is very similar to simple_example.jl from QuAlgorithmZoo.jl.
If I modify simple_example.jl such that the return value is a single element array then calculating Jacobian like below
returns (nothing, ) while calculating the gradient on a scalar output works.
I suspect this has something to do with the way jacobian is implemented (this is written in the Zygote doc "For arguments of any type except Number & AbstractArray, the result is nothing."), however, I don't know what to do about it.

using Zygote
include("chainrules_patch.jl")

import YaoExtensions, Random

c = YaoExtensions.variational_circuit(5)
dispatch!(c, :random)

function loss(reg::AbstractRegister, circuit::AbstractBlock{N}) where N
    reg = apply(copy(reg), circuit)
    st = state(reg)
    [sum(real(st.*st))]
end

reg0 = zero_state(5)
print(jacobian(c->loss(reg0, c), c))

[Jinguo Liu (刘金国)]

Thanks for reporting the issue. You do not need chainrules_patch.jl anymore in the latest Yao. Now you should define the loss function in a different (more natural) way

julia> function loss(reg::AbstractRegister, circuit, params)
           reg = apply(copy(reg), dispatch(circuit, params))
           st = state(reg)
           [sum(real(st.*st))]
       end
loss (generic function with 3 methods)

julia> Zygote.jacobian(p->loss(reg0, c, p), parameters(c))
([-0.07283422404949928 -0.12449066973901673 … 0.25057711019332585 -0.30987561522442586],)

If you have an output with many outputs, you should use the ForwardDiff, it works good.

julia> ForwardDiff.jacobian(p->loss(ArrayReg(Complex{eltype(p)}.(reg0.state)), c, p), parameters(c))
1×50 Matrix{Float64}:
 -0.0728342  -0.124491  0.17673  -0.160745  0.0651408  -0.19335  …  0.08197  -0.162243  -0.017159  -0.171337  0.250577  -0.309876

In practise, you just compare their performance and pick a better one. The larger the output dimension, the worse the Zygote works.