qml.capture.enable()

@subroutine
def flip(w1, w2, w3, theta):
    @qml.for_loop(0, 2, 1)
    def loop_rx(i, _theta):
        qml.X(w1)
        qml.Y(w2)
        qml.Z(w3)
        qml.ctrl(qml.RX, (w1, w2))(_theta, wires=0)
        return jnp.sin(_theta)
    _ = loop_rx(theta)

@qjit
@qml.qnode(qml.device("lightning.qubit", wires=1))
def circuit():
    with qml.allocate(2) as q1:
        with qml.allocate(3) as q2:
            flip(q1[0], q1[1], q2[2], 1.23)

    return qml.expval(qml.Z(0))

print(circuit())