new, port zygote

Author: GiggleLiu

examples/PortZygote/gate_learning.jl

@@ -0,0 +1,31 @@
+using YaoExtensions, Yao
+using Test, Random
+using QuAlgorithmZoo: Adam, update!
+
+include("zygote_patch.jl")
+
+function loss(u, ansatz)
+    m = Matrix(ansatz)
+    sum(abs.(u .- m))
+end
+
+function learn_su4(u::AbstractMatrix; optimizer=Adam(lr=0.1), niter=100)
+    ansatz = general_U4() * put(2, 1=>phase(0.0))  # initial values are 0, here, we attach a global phase.
+    params = parameters(ansatz)
+    for i=1:1000
+        println("Step = $i, loss = $(loss(u,ansatz))")
+        grad = gradient(ansatz->loss(u, ansatz), ansatz)[1]
+        update!(params, grad, optimizer)
+        dispatch!(ansatz, params)
+    end
+    return ansatz
+end
+
+using Random
+Random.seed!(2)
+u = rand_unitary(4)
+using LinearAlgebra
+#u[:,1] .*= -conj(det(u))
+#@show det(u)
+c = learn_su4(u; optimizer=Adam(lr=0.005))
+det(mat(c))

examples/PortZygote/simple_example.jl

@@ -0,0 +1,17 @@
+include("zygote_patch.jl")
+
+import YaoExtensions, Random
+
+c = YaoExtensions.variational_circuit(5)
+dispatch!(c, :random)
+
+function loss(reg::AbstractRegister, circuit::AbstractBlock{N}) where N
+    #copy(reg) |> circuit
+    reg = apply!(copy(reg), circuit)
+    st = state(reg)
+    sum(real(st.*st))
+end
+
+reg0 = zero_state(5)
+paramsδ = gradient(c->loss(reg0, c), c)[1]
+regδ = gradient(reg->loss(reg, c), reg0)[1]

examples/PortZygote/zygote_patch.jl

@@ -0,0 +1,42 @@
+using Zygote
+using Zygote: @adjoint
+using Yao, Yao.AD
+
+@adjoint function apply!(reg::ArrayReg, block::AbstractBlock)
+    out = apply!(reg, block)
+    out, function (outδ)
+        (in, inδ), paramsδ = apply_back((out, outδ), block)
+        return (inδ, paramsδ)
+    end
+end
+
+@adjoint function Matrix(block::AbstractBlock)
+    out = Matrix(block)
+    out, function (outδ)
+        paramsδ = mat_back(block, outδ)
+        return (paramsδ,)
+    end
+end
+
+@adjoint function ArrayReg{B}(raw::AbstractArray) where B
+    ArrayReg{B}(raw), adjy->(reshape(adjy.state, size(raw)),)
+end
+
+@adjoint function ArrayReg{B}(raw::ArrayReg) where B
+    ArrayReg{B}(raw), adjy->(adjy,)
+end
+
+@adjoint function ArrayReg(raw::AbstractArray)
+    ArrayReg(raw), adjy->(reshape(adjy.state, size(raw)),)
+end
+
+@adjoint function copy(reg::ArrayReg) where B
+    copy(reg), adjy->(adjy,)
+end
+
+@adjoint state(reg::ArrayReg) = state(reg), adjy->(ArrayReg(adjy),)
+@adjoint statevec(reg::ArrayReg) = statevec(reg), adjy->(ArrayReg(adjy),)
+@adjoint state(reg::AdjointArrayReg) = state(reg), adjy->(ArrayReg(adjy')',)
+@adjoint statevec(reg::AdjointArrayReg) = statevec(reg), adjy->(ArrayReg(adjy')',)
+@adjoint parent(reg::AdjointArrayReg) = parent(reg), adjy->(adjy',)
+@adjoint Base.adjoint(reg::ArrayReg) = Base.adjoint(reg), adjy->(parent(adjy),)