Repository: PennyLaneAI/pennylane

Issue:
[BUG] Incompatibility of jax with device default.mixed and qml.QubitDensityMatrix()
### Expected behavior

jax.jit can be wrapped around a default.mixed device with qml.QubitDensityMatrix() function.

### Actual behavior

Error: TracerBoolConversionError: Attempted boolean conversion of traced array with shape bool[]..
The error occurred while tracing the function variational_circuit_mixed at <ipython-input-13-66020cf653fc>:9 for jit. This concrete value was not available in Python because it depends on the value of the argument state_ini.

### Additional information

As you can see in the source code, the variational_circuit_qubit() works correctly, while variational_circuit_mixed() stumps upon an error due to the jit compilation. It's the same scenario, but with the mixed device it's not working and, as far as I understand, there should be compatibility between jax and the default.mixed device.

### Source code

```shell
import pennylane as qml
import numpy as np
import jax

nqubits = 1

dev = qml.device("default.qubit", wires=nqubits)
@jax.jit
@qml.qnode(dev, interface="jax")
def variational_circuit_qubit(weights, state_ini):
    qml.QubitStateVector(state_ini, wires=range(nqubits))
    return qml.state()

dev = qml.device("default.mixed", wires=nqubits)
@jax.jit
@qml.qnode(dev, interface="jax")
def variational_circuit_mixed(weights, state_ini):
    qml.QubitDensityMatrix(state_ini, wires=range(nqubits))
    return qml.state()


state_ini = np.array([1,0])
rho_ini = np.tensordot(state_ini, state_ini, axes=0)

state_out = variational_circuit_qubit(0, state_ini)
rho_out = variational_circuit_mixed(0, rho_ini)
```


### Tracebacks

```shell
---------------------------------------------------------------------------
TracerBoolConversionError                 Traceback (most recent call last)
File d:\Documents\1. GitHub\qurriculum_learning\Phase Recognition\test.py:27
     24 rho_ini = np.tensordot(state_ini, state_ini, axes=0)
     26 state_out = variational_circuit_qubit(0, state_ini)
---> 27 rho_out = variational_circuit_mixed(0, rho_ini)

    [... skipping hidden 12 frame]

File d:\Documents\1. GitHub\venv\lib\site-packages\pennylane\qnode.py:1027, in QNode.__call__(self, *args, **kwargs)
   1022         full_transform_program._set_all_argnums(
   1023             self, args, kwargs, argnums
   1024         )  # pylint: disable=protected-access
   1026 # pylint: disable=unexpected-keyword-arg
-> 1027 res = qml.execute(
   1028     (self._tape,),
   1029     device=self.device,
   1030     gradient_fn=self.gradient_fn,
   1031     interface=self.interface,
   1032     transform_program=full_transform_program,
   1033     config=config,
   1034     gradient_kwargs=self.gradient_kwargs,
   1035     override_shots=override_shots,
   1036     **self.execute_kwargs,
   1037 )
   1039 res = res[0]
   1041 # convert result to the interface in case the qfunc has no parameters

File d:\Documents\1. GitHub\venv\lib\site-packages\pennylane\interfaces\execution.py:616, 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)
    614 # Exiting early if we do not need to deal with an interface boundary
    615 if no_interface_boundary_required:
--> 616     results = inner_execute(tapes)
    617     return post_processing(results)
    619 _grad_on_execution = False

File d:\Documents\1. GitHub\venv\lib\site-packages\pennylane\interfaces\execution.py:249, in _make_inner_execute.<locals>.inner_execute(tapes, **_)
    247 if numpy_only:
    248     tapes = tuple(qml.transforms.convert_to_numpy_parameters(t) for t in tapes)
--> 249 return cached_device_execution(tapes)

File d:\Documents\1. GitHub\venv\lib\site-packages\pennylane\interfaces\execution.py:371, in cache_execute.<locals>.wrapper(tapes, **kwargs)
    366         return (res, []) if return_tuple else res
    368 else:
    369     # execute all unique tapes that do not exist in the cache
    370     # convert to list as new device interface returns a tuple
--> 371     res = list(fn(tuple(execution_tapes.values()), **kwargs))
    373 final_res = []
    375 for i, tape in enumerate(tapes):

File C:\Program Files\WindowsApps\PythonSoftwareFoundation.Python.3.10_3.10.3056.0_x64__qbz5n2kfra8p0\lib\contextlib.py:79, in ContextDecorator.__call__.<locals>.inner(*args, **kwds)
     76 @wraps(func)
     77 def inner(*args, **kwds):
     78     with self._recreate_cm():
---> 79         return func(*args, **kwds)

File d:\Documents\1. GitHub\venv\lib\site-packages\pennylane\_qubit_device.py:460, in QubitDevice.batch_execute(self, circuits)
    455 for circuit in circuits:
    456     # we need to reset the device here, else it will
    457     # not start the next computation in the zero state
    458     self.reset()
--> 460     res = self.execute(circuit)
    461     results.append(res)
    463 if self.tracker.active:

File d:\Documents\1. 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:\Documents\1. GitHub\venv\lib\site-packages\pennylane\_qubit_device.py:279, in QubitDevice.execute(self, circuit, **kwargs)
    276 self.check_validity(circuit.operations, circuit.observables)
    278 # apply all circuit operations
--> 279 self.apply(circuit.operations, rotations=self._get_diagonalizing_gates(circuit), **kwargs)
    281 # generate computational basis samples
    282 if self.shots is not None or circuit.is_sampled:

File d:\Documents\1. 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:\Documents\1. 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:\Documents\1. GitHub\venv\lib\site-packages\pennylane\devices\default_mixed.py:540, in DefaultMixed._apply_density_matrix(self, state, device_wires)
    537 if dm_dim != state.shape[0]:
    538     raise ValueError("Density matrix must be of length (2**wires, 2**wires)")
--> 540 if not qnp.allclose(
    541     qnp.trace(qnp.reshape(state, (state_dim, state_dim))), 1.0, atol=tolerance
    542 ):
    543     raise ValueError("Trace of density matrix is not equal one.")
    545 if len(device_wires) == self.num_wires and sorted(device_wires.labels) == list(
    546     device_wires.labels
    547 ):
    548     # Initialize the entire wires with the state

    [... skipping hidden 1 frame]

File d:\Documents\1. 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 variational_circuit_mixed at <ipython-input-15-745eb18162e5>:16 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.33.1
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:\documents\1. 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

Platform info:           Windows-10-10.0.19045-SP0
Python version:          3.10.11
Numpy version:           1.26.2
Scipy version:           1.11.4
Installed devices:
- default.gaussian (PennyLane-0.33.1)
- default.mixed (PennyLane-0.33.1)
- default.qubit (PennyLane-0.33.1)
- default.qubit.autograd (PennyLane-0.33.1)
- default.qubit.jax (PennyLane-0.33.1)
- default.qubit.legacy (PennyLane-0.33.1)
- default.qubit.tf (PennyLane-0.33.1)
- default.qubit.torch (PennyLane-0.33.1)
- default.qutrit (PennyLane-0.33.1)
- null.qubit (PennyLane-0.33.1)
- lightning.qubit (PennyLane-Lightning-0.33.1)
```


### 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
@@ -537,7 +537,7 @@ def _apply_density_matrix(self, state, device_wires):
         if dm_dim != state.shape[0]:
             raise ValueError("Density matrix must be of length (2**wires, 2**wires)")
 
-        if not qnp.allclose(
+        if not qml.math.is_abstract(state) and not qnp.allclose(
             qnp.trace(qnp.reshape(state, (state_dim, state_dim))), 1.0, atol=tolerance
         ):
             raise ValueError("Trace of density matrix is not equal one.")


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,6 +89,25 @@ def circuit(a):
 
         assert np.allclose(state, expected, atol=tol, rtol=0)
 
+    def test_qubit_density_matrix_jit_compatible(self, mocker):
+        """Test that _apply_density_matrix works with jax-jit"""
+
+        dev = qml.device("default.mixed", wires=1)
+        spy = mocker.spy(dev, "_apply_density_matrix")
+
+        @jax.jit
+        @qml.qnode(dev, interface="jax")
+        def circuit(state_ini):
+            qml.QubitDensityMatrix(state_ini, wires=[0])
+            return qml.state()
+
+        state_ini = jnp.array([1, 0])
+        rho_ini = jnp.tensordot(state_ini, state_ini, axes=0)
+        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]])
+
 
 class TestDtypePreserved:
     """Test that the user-defined dtype of the device is preserved for QNode