How to return a cudaq.qvector after measuring ancillas

[babreu-ncsa]

I'm building quantum states for distributed algorithms. One of the nodes in my network is an entanglement generator (G) that can send quantum states (qubits) to all other nodes. In this particular case, G creates a 4-qubit entangled state $\lvert \Psi \rangle$ and passes two to node A, one to node B, and another one to node C. This is just for context.

It happens that $\lvert \Psi \rangle$ is a superposition of states $\lvert \Psi \rangle = \lvert \Phi_1 \rangle + \lvert \Phi_2 \rangle$ and I know the circuits $C_i$ that generate $\lvert \Phi_i \rangle$, so I'm using controlled boxes and a rotated ancilla to tune the superposition. However, the state I want only is generated when the ancilla is measured at a certain state. So I want to build a cudaq.kernel that measures the ancilla and, if I get the right value, it returns me the quantum state with the remaining qubits. I'm not sure how to do this.

Here's what I have:

@cudaq.kernel
def c1(reg: cudaq.qvector):
    h(reg[0])
    for i in range(reg.size()-1):
        x.ctrl(reg[i], reg[i+1])
    x(reg[0])
    x(reg[1])

@cudaq.kernel
def controlled_c1(control: cudaq.qubit, target: cudaq.qvector):
    cudaq.control(c1, control, target)

@cudaq.kernel
def c2(reg: cudaq.qvector):
    x(reg[1])
    h(reg[0])
    x.ctrl(reg[0], reg[1])

    x(reg[3])
    h(reg[2])
    x.ctrl(reg[2], reg[3])

@cudaq.kernel
def controlled_c2(control: cudaq.qubit, target: cudaq.qvector):
    cudaq.control(c2, control, target)

@cudaq.kernel
def superpose():
    register = cudaq.qvector(4)
    ancilla = cudaq.qubit()

    h(ancilla)
    controlled_c1(ancilla, register)
    x(ancilla)
    controlled_c2(ancilla, register)
    h(ancilla)

    mz(ancilla)

I think I need some logic at the end of superpose() to return register given the ancilla state when $Z$-measured but I can't really find in the documentation how to handle kernel returns (does a kernel always need to be measured in the end?) and if I write def superpose() -> cudaq.qvector: to return register in the end (even without the postselection) I get errors like

RuntimeError: Unsupported CUDA Quantum kernel return type - !quake.veq<?>.

Any comments and suggestions are welcome! Thanks!!!

[babreu-ncsa]

I know I get the state I want because if I add

    anc = mz(ancilla)
    if(anc):
        reg = mz(register)

I get something like

{ 
  reg : { 1010:244 1001:254 0101:282 0011:515 1100:493 0110:262 }
   __global__ : { 10101:244 10010:270 01100:238 10011:254 11001:493 01011:282 01010:247 00111:515 01101:262 10100:258 11000:507 00110:526 }
   anc : { 1:2050 0:2046 }
}

when I sample the kernel so register is giving me the right states but I want the kernel to return the qvector object without measuring. That state will obviously be passed to subsequent kernels....