Non unitary custom gate within qml.ctrl...?

Hi everyone, I am back with another question today.
is it possible to include a mid-circuit-measurement, which is not unitary, inside the qml.ctrl() operation?

I show you the code here:

import pennylane as qml
import numpy as np
import matplotlib.pyplot as plt

qubits = list(range(3))
dev = qml.device("default.qubit", wires=len(qubits))

@qml.qnode(dev, interface="numpy")
def circuit_mid_measurement(weights: np.array) -> float:
    """
    Circuit with the mid-measurement operation inside.
    :param weights: (np.array) The weights to feed the circuit with;
    :return: measurement (float) The measurement.
    """
    qml.RX(weights[0], wires=qubits[0])
    m_outcome = qml.measure(wires=qubits[1])
    qml.cond(m_outcome, qml.U3)(weights[1], weights[2], weights[3], wires=qubits[2])
    measurement = qml.expval(qml.PauliZ(wires=qubits[0]))
    return measurement
    
trainable_weights = np.random.uniform(-1., 1., 4)
qml.draw_mpl(circuit_mid_measurement)(trainable_weights)
plt.show()
# This cell is done only for plotting the mid-circuit custom unitary I would like to add inside
# a larger quantum circuit.

def mid_circuit_measure(weights: np.array) -> None:
    """
    Circuit with the mid-measurement operation inside.
    :param weights: (np.array) The weights to feed the circuit with;
    :return: None.
    """
    qml.RX(weights[0], wires=qubits[0])
    m_outcome = qml.measure(wires=qubits[1])
    qml.cond(m_outcome, qml.U3)(weights[1], weights[2], weights[3], wires=qubits[2])


dev2 = qml.device("default.qubit", wires=len(qubits)+1)
@qml.qnode(device=dev2, interface="numpy")
def circuit_total(weights: np.array) -> float:
    """
    Total circuit which applies a customized operation with a mid-circuit measurement inside.
    :param weights: (np.array) The weights to feed the circuit with;
    :return: measurement: (float) The measured value of the circuit.
    """
    qml.ctrl(op=mid_circuit_measure, control=3, control_values=0)(
        weights[:4]
                                                                  )
    measurement = qml.expval(qml.PauliZ(wires=qubits[0]))
    return measurement
                      
trainable_weights = np.random.uniform(-1., 1., 4)
qml.draw_mpl(circuit_total)(trainable_weights)
plt.show()

# In this cell I got the error:
ValueError: The object measure(wires=[1]) of type
<class 'pennylane.measurements.mid_measure.MidMeasureMP'> is not an Operator
or callable.
This error might occur if you apply ctrl to a list of operations instead of a
function or Operator.

Hey @checcopo!

We had a similar question here: 'NoneType' object is not iterable error when working with mid-circuit measurements or qml.ctrl

I think you’d be able to accomplish what you want to do by following my response to that question. Let me know!