update Zygote example

Author: GiggleLiu

Project.toml

@@ -3,21 +3,21 @@ uuid = "65c24e16-9b0a-11e8-1353-efc5bc5f6586"
 version = "0.1.0"
 
 [deps]
+BitBasis = "50ba71b6-fa0f-514d-ae9a-0916efc90dcf"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
-BitBasis = "50ba71b6-fa0f-514d-ae9a-0916efc90dcf"
 Yao = "5872b779-8223-5990-8dd0-5abbb0748c8c"
 YaoArrayRegister = "e600142f-9330-5003-8abb-0ebd767abc51"
 YaoBlocks = "418bc28f-b43b-5e0b-a6e7-61bbc1a2c1df"
 YaoExtensions = "7a06699c-c960-11e9-3c98-9f78548b5f0f"
 
 [compat]
-julia = "1"
-Yao = "0.6.0"
 BitBasis = "0.6"
+Yao = "0.6.0"
 YaoArrayRegister = "0.6"
 YaoBlocks = "0.8, 0.9, 0.10"
 YaoExtensions = "0.2"
+julia = "1"
 
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

examples/PortZygote/shared_parameters.jl

@@ -0,0 +1,25 @@
+include("zygote_patch.jl")
+
+import YaoExtensions, Random
+
+c = YaoExtensions.variational_circuit(5)
+h = YaoExtensions.heisenberg(5)
+
+function loss(h, c, θ) where N
+    # the assign is nessesary!
+    c = dispatch!(c, fill(θ, nparameters(c)))
+    reg = apply!(zero_state(nqubits(c)), c)
+    real(expect(h, reg))
+end
+
+reg0 = zero_state(5)
+zygote_grad = Zygote.gradient(θ->loss(h, c, θ), 0.5)[1]
+
+
+# check gradients
+using Test
+dispatch!(c, fill(0.5, nparameters(c)))
+greg, gparams = expect'(h, zero_state(5)=>c)
+true_grad = sum(gparams)
+
+@test true_grad ≈ true_grad

examples/PortZygote/zygote_patch.jl

@@ -10,6 +10,13 @@ using Yao, Yao.AD
     end
 end
 
+@adjoint function Yao.dispatch!(block::AbstractBlock, params)
+    out = dispatch!(block, params)
+    out, function (outδ)
+        (nothing, outδ)
+    end
+end
+
 @adjoint function Matrix(block::AbstractBlock)
     out = Matrix(block)
     out, function (outδ)
@@ -40,3 +47,8 @@ end
 @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),)
+Zygote.@nograd Yao.nparameters
+Zygote.@nograd Yao.zero_state
+Zygote.@nograd Yao.rand_state
+Zygote.@nograd Yao.uniform_state
+Zygote.@nograd Yao.product_state