Autodiff through chain of chain

[artix41]

When running the code below,

function U()
    return chain(Rz(0), Rx(0), Rz(0))
end
v_unitary = chain(n, put(1=>U()), put(2=>U()), 
                  control(2, 1=>X), 
                  put(1=>Rz(0)), put(2=>Ry(0)),
                  control(1, 2=>X),
                  put(2=>Ry(0)),
                  control(2, 1=>X),
                  put(1=>U()), put(2=>U()));
c = v_unitary |> autodiff(:BP)

I get the following output:

nqubits: 2
chain
├─ put on (1)
│  └─ chain
│     ├─ rot(Z gate, 0.0)
│     ├─ rot(X gate, 0.0)
│     └─ rot(Z gate, 0.0)
├─ put on (2)
│  └─ chain
│     ├─ rot(Z gate, 0.0)
│     ├─ rot(X gate, 0.0)
│     └─ rot(Z gate, 0.0)
├─ control(2)
│  └─ (1,) X gate
├─ [∂] put on (1)
│     └─ rot(Z gate, 0.0)
├─ [∂] put on (2)
│     └─ rot(Y gate, 0.0)
├─ control(1)
│  └─ (2,) X gate
├─ [∂] put on (2)
│     └─ rot(Y gate, 0.0)
├─ control(2)
│  └─ (1,) X gate
├─ put on (1)
│  └─ chain
│     ├─ rot(Z gate, 0.0)
│     ├─ rot(X gate, 0.0)
│     └─ rot(Z gate, 0.0)
└─ put on (2)
   └─ chain
      ├─ rot(Z gate, 0.0)
      ├─ rot(X gate, 0.0)
      └─ rot(Z gate, 0.0)

meaning that it didn't find chains defined with U() to be differentiable. Is there something wrong with my syntax or is it a bug?

[GiggleLiu]

autodiff(:BP) is a legacy feature that supports only limited features, it will be deprecated soon, please use Yao.AD insteads, that is a more powerful classical AD framework. In Yao.AD, you don't need to warp a node with Diff.

On the other side, autodiff(:QC) will capture all differentiable nodes. But it is only for quantum faithful simulation. We are now polishing these interfaces in #11 .

[artix41]

I see, thanks for your answer! Do you have any example of how to use Yao.AD?

[GiggleLiu]

For observable loss, a single prime would work.
https://github.com/QuantumBFS/YaoBlocks.jl/blob/3bb64e810e38ff8f02ac8975b7d04206eed8232f/test/autodiff/apply_back.jl#L81

For general backward, you need to call apply_back explicitly like in: https://github.com/QuantumBFS/YaoBlocks.jl/blob/3bb64e810e38ff8f02ac8975b7d04206eed8232f/src/autodiff/apply_back.jl#L167

Also, matrix representation can be back propagated with mat_back.

Note: need master branch.

[artix41]

So if I want to take the gradient of

L = f(expect(H, state => circuit))

with respect to the parameters of the the circuit (f being any differentiable function), what should I do? (the apply_back example is a bit cryptic)

[Roger-luo]

BTW, I'll hook backwards to ChainRules in the weekend. (I have 2 assignment due in the week) which should solve this issue once and for all. But @GiggleLiu 's solution is what you can do before we release this.

apply_back is the adjoint (or backward function of apply) which means you need to apply the chain rule yourself on the loss if there is no AD framework (which is actually straightforward tho)

[GiggleLiu]

@artix41

For your case

y = f(expect(H, state => circuit))
L = f(y)

You need to compute dL/dy first (with a classical AD framework, Flux, KNet, Zygote et. al.).
Then, with expect'(H, state => circuit) (notice the prime here), one can obtain dy/dstate and dy/dparams. The final adjoint is (dL/dy) * (dy/dparams).

I will write a working example ASAP.

[wangleiphy]

This repo contains an example of interfacing Yao's AD (expect'(H, state => circuit)) with Zygote
https://github.com/wangleiphy/BetaVQE.jl

[GiggleLiu]

Sorry for the late reply, also see this gate learning example

https://github.com/QuantumBFS/QuAlgorithmZoo.jl/blob/master/examples/PortZygote/gate_learning.jl

[yuguoshao]

Sorry to ask a question, Now I can get the gradient of a loss function (like observable), But is there any way if I want to get the Jacobian matrix of state vector (\frac{\partial |x>}{\partial \theta}) ?

[GiggleLiu]

When the output is larger than input, forwarddiff is much faster. I can give you a demo tomorrow (in bed now). You can also try to implement one if you need it urgently. It should be very easy.

https://github.com/JuliaDiff/ForwardDiff.jl

[yuguoshao]

Honestly, I faced some problem when I tired to do like this way (using Zygote). Maybe your demo will help me a lot. I'm not in a hurry, I wish you a good night's sleep. Thank you very much!

[GiggleLiu]

Once the above PR get merged and tagged, you should be able to update your package and use non-inplace API to compute jacobians.
If you are eager to use this API, you can also install that branch manually by typing

pkg> add YaoBlocks#non-inplace-API

[Roger-luo]

Regarding to this issue itself, I'm hoping to get rid of Zygote at some point entirely from YaoCompiler side... and replace with Enzyme or Diffractor, in near term it would be Enzyme, since ChainRules doesn't have an API definition of in-place operations yet. But not sure if this would happen this year, it will depend on the progress of ChainRules too.