How to define a trainable Unitary in Controlled-Unitary gate by ourselves?

Hi, I want to implement the Quantum circuit like the image below.

I know ControlledQubitUnitary might be possible to achieve my goal, but I want the value inside the Unitary to become trainable instead of the fixed Unitary matrix.
Can anyone show me how to do this?
Thanks!

Hey @mini! Great question. You can use qml.ctrl like this:

import pennylane as qml
from pennylane import numpy as np

dev = qml.device("default.qubit", wires=range(6))

def subcircuit(thetas):
    qml.RY(thetas[0], wires=0)
    qml.RY(thetas[1], wires=1)
    qml.RY(thetas[2], wires=2)

    qml.CNOT([0, 1])
    qml.CNOT([1, 2])

    qml.RY(thetas[3], wires=0)
    qml.RY(thetas[4], wires=1)
    qml.RY(thetas[5], wires=2)

    qml.CNOT([0, 1])
    qml.CNOT([1, 2])

@qml.qnode(dev)
def circuit(thetas):
    qml.ctrl(subcircuit, (3, 4, 5), control_values=(1,1,1))(thetas)
    return [qml.expval(qml.PauliY(i)) for i in dev.wires]

thetas = np.random.uniform(0, 2*np.pi, size=(6,))
print(qml.draw(circuit)(thetas))

0: ─╭RY(0.79)─────────────────────╭●────╭RY(0.02)─────────────────────╭●────┤  <Y>
1: ─│─────────╭RY(6.17)───────────├X─╭●─│─────────╭RY(3.36)───────────├X─╭●─┤  <Y>
2: ─│─────────│─────────╭RY(4.05)─│──├X─│─────────│─────────╭RY(2.21)─│──├X─┤  <Y>
3: ─├●────────├●────────├●────────├●─├●─├●────────├●────────├●────────├●─├●─┤  <Y>
4: ─├●────────├●────────├●────────├●─├●─├●────────├●────────├●────────├●─├●─┤  <Y>
5: ─╰●────────╰●────────╰●────────╰●─╰●─╰●────────╰●────────╰●────────╰●─╰●─┤  <Y>

Basically, ctrl can take any PennyLane Operator or a quantum function (the example above uses a quantum function called subcircuit). Hope this helps!

Thank you @isaacdevlugt ! I also wanted to change the layer number by myself, so I did this:

import pennylane as qml
from pennylane import numpy as np

dev = qml.device("default.qubit", wires=range(6))

def subcircuit(thetas):
    for ii in range(layer_num):
        for jj in range(3):
            qml.RY(thetas[ii][jj], wires=jj)
        for jj in range(2):
            qml.CNOT(wires=[jj,jj+1])
            

@qml.qnode(dev)
def circuit(thetas):
    qml.ctrl(subcircuit, (3, 4, 5), control_values=(1,1,1))(thetas)
    return [qml.expval(qml.PauliY(i)) for i in dev.wires]


layer_num = 3
thetas = np.random.uniform(0, 2*np.pi, size=(layer_num,3))
qml.draw_mpl(circuit)(thetas)

Thanks again for your guidance! :saluting_face: :raised_hands:

Nice! I think your subcircuit might need layer_num as an argument as well :slight_smile: