[BUG] Uncompatibility between "qml.density_matrix" and "default.mixed" in jax.jit

[erikrecio]

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

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

---------------------------------------------------------------------------
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

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.

[timmysilv]

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.

[erikrecio]

This is great thanks!

[erikrecio]

Hello, is this already fixed? I see a fix has been merged, but upon upgrading my pennylane version with the one on github I keep getting the same error.

[timmysilv]

Hi, this fix was merged into a branch used to prepare for the upcoming 0.35 release, so it's not yet on master. #5274 (should merge sometime today) will bring those changes into master if you want to install from Github after that. It will also be in the 0.35 release next week. Will close this as the fix has been accepted