Building custom gates for variational purposes

[FabianBallar7]

Hey everyone!

I've been playing around with defining my own custom Primitive blocks. In particular, I'm trying to implement a gate which behaves as the fSim(θ,ϕ) at ϕ = 0. I followed the steps laid out in the docs and also got inspired from the fSim definition itself and wrote the following definition:

mutable struct XY{T<:Real} <: PrimitiveBlock{2}
    θ::T
end

YaoBlocks.print_block(io::IO, block::XY) = print(io, "XY(θ=$(block.θ)))")

function Base.:(==)(xy1::XY, xy2::XY)
    return xy1.θ == xy2.θ 
end

# Implement the matrix representation of the X+Y rotation
function Yao.mat(::Type{T}, g::XY) where T
    θ = g.θ
    T[1  0       0       0
      0  cos(θ/2) -1im*sin(θ/2)  0
      0  -1im*sin(θ/2)  cos(θ/2)  0
      0  0       0       1]
end


YaoBlocks.nqudits(g::XY) = 2

#Define them as parameters
Yao.niparams(::Type{<:XY}) = 1
Yao.iparams_eltype(::XY{T}) where T = T
Yao.getiparams(g::XY{T}) where T = g.θ
function Yao.setiparams!(g::XY{T}, θ::Number) where T
    g.θ = θ
    return g
end
#A version for vectors for when I want to do dispatch on a large circuit
function Yao.setiparams!(g::XY{T}, θ::AbstractArray) where T
    g.θ = θ[1] 
    return g
end
function Yao.setiparams(g::XY{T}, θ::AbstractArray) where T
    return XY(g.θ)
end

cache_key(gate::XY) = gate.θ

YaoBlocks.@dumpload_fallback XY XY

YaoPlots.get_brush_texts(c, b::XY) = (c.gatestyles.g, "XY($(YaoPlots.pretty_angle(b.θ)))")


function Yao.AD.mat_back!(::Type{T}, gate::XY{RT}, adjy::AbstractMatrix, collector::AbstractVector) where {T<:Complex,RT<:Real}
    θ = gate.θ
    
    # Derivative of the matrix with respect to θ
    dM = T[0  0           0           0
           0  -sin(θ/2)/2 -1im*cos(θ/2)/2  0
           0  -1im*cos(θ/2)/2 -sin(θ/2)/2  0
           0  0           0           0]
    
    # Calculate the gradient
    grad = real(tr(adjy' * dM))
    
    # Use projection function to accumulate the gradient
    pushfirst!(collector, projection(θ, grad))
    
    return collector
end

In terms of building circuits, dispatching parameters, and so on everything seems to be working. However, when trying to compute gradients I get errors. I've managed to solve a couple of them however now I'm getting the following error:

`Tangent` for type XY{Float64} is not defined!

I assume this is related to the tangent in ChainRulesCore, I wonder if I have to also define this on my side. I thought before proceeding, I'd ask if I took a wrong turn or if I was doing something flat-out wrong. I've tried looking for examples of building custom gates from scratch but only found very general instructions in the documentation.

I'd appreciate any help with this. Also, if this is solved, I'd love to make a more detailed tutorial on how to make your own gates, given that it becomes so essential at a certain point.

[GiggleLiu]

Nice try, please also check the sim defined here: https://github.com/QuantumBFS/Yao.jl/blob/master/src/EasyBuild/block_extension/FSimGate.jl

Regarding the error, have you checked these lines:

Yao.jl/lib/YaoBlocks/src/autodiff/chainrules_patch.jl

Line 47 in 3c6a767

@eval function recursive_create_tangent(c::$GT)

[FabianBallar7]

Thanks for your response!

Thank you for pointing me towards the chainrules_patch.jl, this did indeed help, although I still haven't managed to solve the issue just yet.

I figured that if I try to compute gradients with just the _,grad = expect'(...) it does work and I get correct gradients. However, the issue arises only when I compute gradients through Zygote Zygote.gradient() for some other more complicated loss functions I'm trying out.

In fact, the same is true for the FSimGate currently implemented in Yao. I tried computing the gradients through Zygote and since mat_back! is not defined for it, it cannot compute gradients in this way. This function is straightforward to implement, however still when I implement it, the Tangent error shows up.

I guess at this point its about writing the tangent function for the ::XY and ::FSimGate separately but I'm unsure if there's a better way. I'd appreciate any further help or ideas to try.

Thanks a lot again for all the help!

[GiggleLiu]

Glad it almost works. I think you have to implement the tangent function for ::FSimGate to obtain the gradients. This is the design of ChainRulesCore. It does not provide a default gradient implementation on user defined constructors.

[FabianBallar7]

Thanks again for the response Jinguo. I really appreciate it.

I've been thinking a bit more about this problem throughout the weekend. I agree with you that the problem is in defining a tangent for the FSimGate or any other user-defined gate for that matter when trying to use Zygote or other AD platforms through ChainRulesCore.

However, I'd like to ask for a bit of direction with respect to the definition of the tangent. In particular, I'm interested in understanding how exactly the tangents are defined for the primitive gates already built in Yao. My question comes mainly from the fact that my error which is precisely triggered by line 12 in the chainrules_patch.jl file which I link below.

Yao.jl/lib/YaoBlocks/src/autodiff/chainrules_patch.jl

Line 12 in 3c6a767

error("`Tangent` for type $(typeof(c)) is not defined!")

Here it's clear to me that if the recursive_create_tangent() function is called on the types of the custom gates it recognizes that it has not been defined and throws the error. I assume that the important part to define the tangent function for the new gates is to write the corresponding rrule functions for our new type such as the ones defined in the same file linked above starting around line 79.

However when I define this function for my new gate I get the same exact errors I'm guessing because, still, recursive_create_tangent() has not been defined for the custom type even if the tangent function exists.

Does this logic sound alright? I'm really wondering if defining these tangents requires actually adding more function and that's why I'm seeking advice potentially to see If I'm missing some particular function that remains to be defined. I'd appreciate any help, even if it's just pointing me in the right direction as to which functions I should write in order to have my new tangents recognized by ChainRulesCore.

I think if we pin this down it would be quite beneficial overall. Making it more accessible to define custom gates and use them in more complicated variational setups.

Again, I appreciate responses quite a lot!