Fix issue 545 - autodiff with batched-arrayreg (#547)

Author: GiggleLiu

lib/YaoBlocks/src/autodiff/chainrules_patch.jl

@@ -79,18 +79,19 @@ end
 function rrule(::typeof(apply), reg::AbstractArrayReg, block::AbstractBlock)
     out = apply(reg, block)
     out, function (outδ)
-        (in, inδ), paramsδ = apply_back((copy(out), tangent_to_reg(typeof(out), outδ)), block)
+        (in, inδ), paramsδ = apply_back((copy(out), tangent_to_reg(out, outδ)), block)
         return (NoTangent(), inδ, create_circuit_tangent(block, paramsδ))
     end
 end
 function rrule(::typeof(apply), reg::AbstractArrayReg, block::AbstractAdd)
     out = apply(reg, block)
     out, function (outδ)
-        (in, inδ), paramsδ = apply_back((copy(out), tangent_to_reg(typeof(out), outδ)), block; in = reg)
+        (in, inδ), paramsδ = apply_back((copy(out), tangent_to_reg(out, outδ)), block; in = reg)
         return (NoTangent(), inδ, create_circuit_tangent(block, paramsδ))
     end
 end
-tangent_to_reg(::Type{T}, reg) where T<:AbstractArrayReg = reg isa Tangent ? T(reg.state) : reg
+tangent_to_reg(reg::T, grad) where {D, T<:ArrayReg{D}} = grad isa Tangent ? ArrayReg{D}(grad.state) : grad
+tangent_to_reg(reg::T, grad) where {D, T<:BatchedArrayReg{D}} = grad isa Tangent ? BatchedArrayReg{D}(grad.state, reg.nbatch) : grad
 
 
 function rrule(::typeof(dispatch), block::AbstractBlock, params)

lib/YaoBlocks/test/autodiff/chainrules_patch.jl

@@ -285,4 +285,38 @@ end
     loss3(c, rand(nparameters(c)))
     zygote_grad = Zygote.gradient(θs->loss3(c, θs), rand(nparameters(c)))[1]
     @test zygote_grad isa Vector
+end
+
+@testset "issue #545" begin
+    function variational_circuit(q, nlayer)
+        circuit = chain(q)
+        for i = 1:nlayer
+            push!(circuit, put(q, i=>chain(Rx(0.0), Ry(0.0), Rz(0.0))))
+        end
+        circuit
+    end
+    function bug_mwe()    
+        q = 4
+        vqc = variational_circuit(q, 2)
+        param = 4pi*rand(nparameters(vqc))
+
+        # bstate = rand_state(q) # works OK    
+        bstate = rand_state(q, nbatch=2) # Zygote errors
+        
+        function loss(param)
+            vqc_p = dispatch(vqc, param)
+            bstate_p = apply(bstate, vqc_p)
+            pr = probs(bstate_p)
+            pr_s = sum(pr, dims = 2)
+            r = -sum(pr_s[1:2^(4-1)])
+            return r
+        end   
+
+        l = loss(param)
+        grad = Zygote.gradient(loss, param)[1]
+
+        return grad
+    end
+
+    @test bug_mwe() isa Vector
 end