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:

Resolving patch:
diff --git a/pennylane/devices/default_mixed.py b/pennylane/devices/default_mixed.py
--- a/pennylane/devices/default_mixed.py
+++ b/pennylane/devices/default_mixed.py
@@ -573,7 +573,7 @@ def _apply_density_matrix(self, state, device_wires):
                     right_axes.append(index + self.num_wires)
             transpose_axes = left_axes + right_axes
             rho = qnp.transpose(rho, axes=transpose_axes)
-            assert qnp.allclose(
+            assert qml.math.is_abstract(rho) or qnp.allclose(
                 qnp.trace(qnp.reshape(rho, (2**self.num_wires, 2**self.num_wires))),
                 1.0,
                 atol=tolerance,


Test patch:
diff --git a/tests/devices/test_default_mixed_jax.py b/tests/devices/test_default_mixed_jax.py
--- a/tests/devices/test_default_mixed_jax.py
+++ b/tests/devices/test_default_mixed_jax.py
@@ -89,10 +89,11 @@ def circuit(a):
 
         assert np.allclose(state, expected, atol=tol, rtol=0)
 
-    def test_qubit_density_matrix_jit_compatible(self, mocker):
+    @pytest.mark.parametrize("n_qubits", [1, 2])
+    def test_qubit_density_matrix_jit_compatible(self, n_qubits, mocker):
         """Test that _apply_density_matrix works with jax-jit"""
 
-        dev = qml.device("default.mixed", wires=1)
+        dev = qml.device("default.mixed", wires=n_qubits)
         spy = mocker.spy(dev, "_apply_density_matrix")
 
         @jax.jit
@@ -106,7 +107,11 @@ def circuit(state_ini):
         rho_out = circuit(rho_ini)
         spy.assert_called_once()
         assert qml.math.get_interface(rho_out) == "jax"
-        assert np.array_equal(rho_out, [[1, 0], [0, 0]])
+
+        dim = 2**n_qubits
+        expected = np.zeros((dim, dim))
+        expected[0, 0] = 1.0
+        assert np.array_equal(rho_out, expected)
 
 
 class TestDtypePreserved: