Repository: PennyLaneAI/pennylane

Issue:
[BUG] Uncompatibility between "qml.density_matrix" and "default.mixed" in jax.jit
### Expected behavior

Works without problems.

### Actual behavior

Error. More info below.

### Additional information

As you can see in the source code; tracing out qubits, the mixed device and jax can not go together. You can also see an example with the "default.qubit" device and it works without problems although I am tracing out qubits and running with jax as well. I'm guessing there's no reason for it to not work with the mixed device so I'm opening this bug.

### Source code

```shell
import pennylane as qml
import jax

nqubits = 2

dev = qml.device("default.qubit", wires=nqubits)
@jax.jit
@qml.qnode(dev, interface="jax")
def circuit_pure(state_ini):
    qml.QubitStateVector(state_ini, wires=[0])
    return qml.density_matrix([0])
    

dev = qml.device("default.mixed", wires=nqubits)
@jax.jit
@qml.qnode(dev, interface="jax")
def circuit_mixed(state_ini):
    qml.QubitDensityMatrix(state_ini, wires=[0])
    return qml.density_matrix([0])


pure_ini = np.array([1,0])
pure_out = circuit_pure(pure_ini)

mixed_ini = np.array([[1,0],[0,0]])
mixed_out = circuit_mixed(mixed_ini)
```


### Tracebacks

```shell
---------------------------------------------------------------------------
TracerBoolConversionError                 Traceback (most recent call last)
File d:\Github\qurriculum_learning\Phase Recognition\test.py:29
     26 pure_out = circuit_pure(pure_ini)
     28 mixed_ini = np.array([[1,0],[0,0]])
---> 29 mixed_out = circuit_mixed(mixed_ini)

    [... skipping hidden 12 frame]

File d:\Github\venv\Lib\site-packages\pennylane\workflow\qnode.py:1046, in QNode.__call__(self, *args, **kwargs)
   1041         full_transform_program._set_all_argnums(
   1042             self, args, kwargs, argnums
   1043         )  # pylint: disable=protected-access
   1045 # pylint: disable=unexpected-keyword-arg
-> 1046 res = qml.execute(
   1047     (self._tape,),
   1048     device=self.device,
   1049     gradient_fn=self.gradient_fn,
   1050     interface=self.interface,
   1051     transform_program=full_transform_program,
   1052     config=config,
   1053     gradient_kwargs=self.gradient_kwargs,
   1054     override_shots=override_shots,
   1055     **self.execute_kwargs,
   1056 )
   1058 res = res[0]
   1060 # convert result to the interface in case the qfunc has no parameters

File d:\Github\venv\Lib\site-packages\pennylane\workflow\execution.py:684, in execute(tapes, device, gradient_fn, interface, transform_program, config, grad_on_execution, gradient_kwargs, cache, cachesize, max_diff, override_shots, expand_fn, max_expansion, device_batch_transform, device_vjp)
    682 # Exiting early if we do not need to deal with an interface boundary
    683 if no_interface_boundary_required:
--> 684     results = inner_execute(tapes)
    685     return post_processing(results)
    687 _grad_on_execution = False

File d:\Github\venv\Lib\site-packages\pennylane\workflow\execution.py:283, in _make_inner_execute.<locals>.inner_execute(tapes, **_)
    281 if numpy_only:
    282     tapes = tuple(qml.transforms.convert_to_numpy_parameters(t) for t in tapes)
--> 283 return cached_device_execution(tapes)

File d:\Github\venv\Lib\site-packages\pennylane\workflow\execution.py:409, in cache_execute.<locals>.wrapper(tapes, **kwargs)
    404         return (res, []) if return_tuple else res
    406 else:
    407     # execute all unique tapes that do not exist in the cache
    408     # convert to list as new device interface returns a tuple
--> 409     res = list(fn(tuple(execution_tapes.values()), **kwargs))
    411 final_res = []
    413 for i, tape in enumerate(tapes):

File ~\AppData\Local\Programs\Python\Python312\Lib\contextlib.py:81, in ContextDecorator.__call__.<locals>.inner(*args, **kwds)
     78 @wraps(func)
     79 def inner(*args, **kwds):
     80     with self._recreate_cm():
---> 81         return func(*args, **kwds)

File d:\Github\venv\Lib\site-packages\pennylane\_qubit_device.py:459, in QubitDevice.batch_execute(self, circuits)
    454 for circuit in circuits:
    455     # we need to reset the device here, else it will
    456     # not start the next computation in the zero state
    457     self.reset()
--> 459     res = self.execute(circuit)
    460     results.append(res)
    462 if self.tracker.active:

File d:\Github\venv\Lib\site-packages\pennylane\devices\default_mixed.py:685, in DefaultMixed.execute(self, circuit, **kwargs)
    683         wires_list.append(m.wires)
    684     self.measured_wires = qml.wires.Wires.all_wires(wires_list)
--> 685 return super().execute(circuit, **kwargs)

File d:\Github\venv\Lib\site-packages\pennylane\_qubit_device.py:277, in QubitDevice.execute(self, circuit, **kwargs)
    274 self.check_validity(circuit.operations, circuit.observables)
    276 # apply all circuit operations
--> 277 self.apply(circuit.operations, rotations=self._get_diagonalizing_gates(circuit), **kwargs)
    279 # generate computational basis samples
    280 sample_type = (SampleMP, CountsMP, ClassicalShadowMP, ShadowExpvalMP)

File d:\Github\venv\Lib\site-packages\pennylane\devices\default_mixed.py:699, in DefaultMixed.apply(self, operations, rotations, **kwargs)
    693         raise DeviceError(
    694             f"Operation {operation.name} cannot be used after other Operations have already been applied "
    695             f"on a {self.short_name} device."
    696         )
    698 for operation in operations:
--> 699     self._apply_operation(operation)
    701 # store the pre-rotated state
    702 self._pre_rotated_state = self._state

File d:\Github\venv\Lib\site-packages\pennylane\devices\default_mixed.py:604, in DefaultMixed._apply_operation(self, operation)
    601     return
    603 if isinstance(operation, QubitDensityMatrix):
--> 604     self._apply_density_matrix(operation.parameters[0], wires)
    605     return
    607 if isinstance(operation, Snapshot):

File d:\Github\venv\Lib\site-packages\pennylane\devices\default_mixed.py:576, in DefaultMixed._apply_density_matrix(self, state, device_wires)
    574 transpose_axes = left_axes + right_axes
    575 rho = qnp.transpose(rho, axes=transpose_axes)
--> 576 assert qnp.allclose(
    577     qnp.trace(qnp.reshape(rho, (2**self.num_wires, 2**self.num_wires))),
    578     1.0,
    579     atol=tolerance,
    580 )
    582 self._state = qnp.asarray(rho, dtype=self.C_DTYPE)
    583 self._pre_rotated_state = self._state

    [... skipping hidden 1 frame]

File d:\Github\venv\Lib\site-packages\jax\_src\core.py:1510, in concretization_function_error.<locals>.error(self, arg)
   1509 def error(self, arg):
-> 1510   raise TracerBoolConversionError(arg)

TracerBoolConversionError: Attempted boolean conversion of traced array with shape bool[]..
The error occurred while tracing the function circuit_mixed at <ipython-input-130-a9b188255961>:18 for jit. This concrete value was not available in Python because it depends on the value of the argument state_ini.
See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerBoolConversionError
```


### System information

```shell
Name: PennyLane
Version: 0.35.0.dev0
Summary: PennyLane is a Python quantum machine learning library by Xanadu Inc.
Home-page: https://github.com/PennyLaneAI/pennylane
Author: 
Author-email: 
License: Apache License 2.0
Location: d:\Github\venv\Lib\site-packages
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, pennylane-lightning, requests, rustworkx, scipy, semantic-version, toml, typing-extensions
Required-by: PennyLane-Lightning, PennyLane-qiskit

Platform info:           Windows-10-10.0.19045-SP0
Python version:          3.12.1
Numpy version:           1.26.2
Scipy version:           1.11.4
Installed devices:
- default.clifford (PennyLane-0.35.0.dev0)
- default.gaussian (PennyLane-0.35.0.dev0)
- default.mixed (PennyLane-0.35.0.dev0)
- default.qubit (PennyLane-0.35.0.dev0)
- default.qubit.autograd (PennyLane-0.35.0.dev0)
- default.qubit.jax (PennyLane-0.35.0.dev0)
- default.qubit.legacy (PennyLane-0.35.0.dev0)
- default.qubit.tf (PennyLane-0.35.0.dev0)
- default.qubit.torch (PennyLane-0.35.0.dev0)
- default.qutrit (PennyLane-0.35.0.dev0)
- null.qubit (PennyLane-0.35.0.dev0)
- lightning.qubit (PennyLane-Lightning-0.34.0)
- qiskit.aer (PennyLane-qiskit-0.34.0)
- qiskit.basicaer (PennyLane-qiskit-0.34.0)
- qiskit.ibmq (PennyLane-qiskit-0.34.0)
- qiskit.ibmq.circuit_runner (PennyLane-qiskit-0.34.0)
- qiskit.ibmq.sampler (PennyLane-qiskit-0.34.0)
- qiskit.remote (PennyLane-qiskit-0.34.0)
```


### Existing GitHub issues

- [X] I have searched existing GitHub issues to make sure the issue does not already exist.


Discussion:
Hi @erikrecio, sorry about that! My fix the last time around missed the case when the operator is not applied to all device wires, but rather just a subset. I've opened a follow-up that fixes the case for applying the operation to a subsystem as well.
This is great thanks!

Resolving patch:
diff --git a/pennylane/_qubit_device.py b/pennylane/_qubit_device.py
--- a/pennylane/_qubit_device.py
+++ b/pennylane/_qubit_device.py
@@ -41,6 +41,7 @@
     MeasurementProcess,
     MeasurementTransform,
     MeasurementValue,
+    MidMeasureMP,
     MutualInfoMP,
     ProbabilityMP,
     SampleMeasurement,
@@ -273,8 +274,23 @@ def execute(self, circuit, **kwargs):
 
         self.check_validity(circuit.operations, circuit.observables)
 
+        has_mcm = any(isinstance(op, MidMeasureMP) for op in circuit.operations)
+        if has_mcm:
+            kwargs["mid_measurements"] = {}
         # apply all circuit operations
-        self.apply(circuit.operations, rotations=self._get_diagonalizing_gates(circuit), **kwargs)
+        self.apply(
+            circuit.operations,
+            rotations=self._get_diagonalizing_gates(circuit),
+            **kwargs,
+        )
+        if has_mcm:
+            mid_measurements = kwargs["mid_measurements"]
+            mid_values = np.array(tuple(mid_measurements.values()))
+            if np.any(mid_values == -1):
+                for k, v in tuple(mid_measurements.items()):
+                    if v == -1:
+                        mid_measurements.pop(k)
+                return None, mid_measurements
 
         # generate computational basis samples
         sample_type = (SampleMP, CountsMP, ClassicalShadowMP, ShadowExpvalMP)
@@ -323,7 +339,7 @@ def execute(self, circuit, **kwargs):
             )
             self.tracker.record()
 
-        return results
+        return (results, mid_measurements) if has_mcm else results
 
     def shot_vec_statistics(self, circuit: QuantumTape):
         """Process measurement results from circuit execution using a device
diff --git a/pennylane/transforms/dynamic_one_shot.py b/pennylane/transforms/dynamic_one_shot.py
--- a/pennylane/transforms/dynamic_one_shot.py
+++ b/pennylane/transforms/dynamic_one_shot.py
@@ -84,6 +84,12 @@ def func(x, y):
     if not any(isinstance(o, MidMeasureMP) for o in tape.operations):
         return (tape,), null_postprocessing
 
+    for m in tape.measurements:
+        if not isinstance(m, (CountsMP, ExpectationMP, ProbabilityMP, SampleMP, VarianceMP)):
+            raise TypeError(
+                f"Native mid-circuit measurement mode does not support {type(m).__name__} measurements."
+            )
+
     aux_tape = init_auxiliary_tape(tape)
     output_tapes = [aux_tape] * tape.shots.total_shots
 
@@ -113,6 +119,29 @@ def processing_fn(results, has_partitioned_shots=None):
     return output_tapes, processing_fn
 
 
+@dynamic_one_shot.custom_qnode_transform
+def _dynamic_one_shot_qnode(self, qnode, targs, tkwargs):
+    """Custom qnode transform for ``dynamic_one_shot``."""
+    if tkwargs.get("device", None):
+        raise ValueError(
+            "Cannot provide a 'device' value directly to the dynamic_one_shot decorator "
+            "when transforming a QNode."
+        )
+    if qnode.device is not None:
+        support_mcms = hasattr(qnode.device, "capabilities") and qnode.device.capabilities().get(
+            "supports_mid_measure", False
+        )
+        support_mcms = support_mcms or isinstance(
+            qnode.device, qml.devices.default_qubit.DefaultQubit
+        )
+        if not support_mcms:
+            raise TypeError(
+                f"Device {qnode.device.name} does not support mid-circuit measurements natively, and hence it does not support the dynamic_one_shot transform. `default.qubit` and `lightning.qubit` currently support mid-circuit measurements and the dynamic_one_shot transform."
+            )
+    tkwargs.setdefault("device", qnode.device)
+    return self.default_qnode_transform(qnode, targs, tkwargs)
+
+
 def init_auxiliary_tape(circuit: qml.tape.QuantumScript):
     """Creates an auxiliary circuit to perform one-shot mid-circuit measurement calculations.
 
diff --git a/pennylane/workflow/qnode.py b/pennylane/workflow/qnode.py
--- a/pennylane/workflow/qnode.py
+++ b/pennylane/workflow/qnode.py
@@ -1008,6 +1008,15 @@ def __call__(self, *args, **kwargs) -> qml.typing.Result:
         else:
             config = None
             full_transform_program = qml.transforms.core.TransformProgram(self.transform_program)
+        has_mcm_support = (
+            any(isinstance(op, MidMeasureMP) for op in self._tape)
+            and hasattr(self.device, "capabilities")
+            and self.device.capabilities().get("supports_mid_measure", False)
+        )
+        if has_mcm_support:
+            full_transform_program.add_transform(qml.dynamic_one_shot)
+            override_shots = 1
+
         # Add the gradient expand to the program if necessary
         if getattr(self.gradient_fn, "expand_transform", False):
             full_transform_program.insert_front_transform(


Test patch:
diff --git a/tests/devices/test_default_qubit_legacy.py b/tests/devices/test_default_qubit_legacy.py
--- a/tests/devices/test_default_qubit_legacy.py
+++ b/tests/devices/test_default_qubit_legacy.py
@@ -17,6 +17,7 @@
 # pylint: disable=too-many-arguments,too-few-public-methods
 # pylint: disable=protected-access,cell-var-from-loop
 import cmath
+import copy
 
 import math
 
@@ -89,6 +90,33 @@
 VARPHI = np.linspace(0.02, 1, 3)
 
 
+def test_qnode_native_mcm(mocker):
+    """Tests that the legacy devices may support native MCM execution via the dynamic_one_shot transform."""
+
+    class MCMDevice(DefaultQubitLegacy):
+        def apply(self, *args, **kwargs):
+            pass
+
+        @classmethod
+        def capabilities(cls):
+            default_capabilities = copy.copy(DefaultQubitLegacy.capabilities())
+            default_capabilities["supports_mid_measure"] = True
+            return default_capabilities
+
+    dev = MCMDevice(wires=1, shots=100)
+    dev.operations.add("MidMeasureMP")
+    spy = mocker.spy(qml.dynamic_one_shot, "_transform")
+
+    @qml.qnode(dev, interface=None, diff_method=None)
+    def func():
+        _ = qml.measure(0)
+        return qml.expval(op=qml.PauliZ(0))
+
+    res = func()
+    assert spy.call_count == 1
+    assert isinstance(res, float)
+
+
 def test_analytic_deprecation():
     """Tests if the kwarg `analytic` is used and displays error message."""
     msg = "The analytic argument has been replaced by shots=None. "
diff --git a/tests/test_qnode.py b/tests/test_qnode.py
--- a/tests/test_qnode.py
+++ b/tests/test_qnode.py
@@ -927,6 +927,22 @@ def conditional_ry_qnode(x, y):
         assert np.allclose(r2[1], mv_res(first_par))
         assert spy.call_count == call_count  # once for each preprocessing
 
+    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
+    def test_dynamic_one_shot_if_mcm_unsupported(self, dev_name):
+        dev = qml.device(dev_name, wires=2, shots=100)
+
+        with pytest.raises(
+            TypeError,
+            match="does not support mid-circuit measurements natively, and hence it does not support the dynamic_one_shot transform.",
+        ):
+
+            @qml.transforms.dynamic_one_shot
+            @qml.qnode(dev)
+            def _():
+                qml.RX(1.23, 0)
+                ms = [qml.measure(0) for _ in range(10)]
+                return qml.probs(op=ms)
+
     @pytest.mark.parametrize("basis_state", [[1, 0], [0, 1]])
     def test_sampling_with_mcm(self, basis_state, mocker):
         """Tests that a QNode with qml.sample and mid-circuit measurements