[BUG] qml.exp with jax.jit causes TracerBoolConversionError #6932
Description
Expected behavior
The function should execute without errors, returning an expectation value.
Actual behavior
The error TracerBoolConversionError: Attempted boolean conversion of traced array occurs during execution.
Additional information
- The issue appears to be related to
qml.exp's decomposition process. - It may be attempting a boolean conversion of a JAX-traced value when checking
self.base.is_hermitian. - Using non-JAX NumPy values does not trigger this issue.
Source code
@qml.qnode(qml.device('lightning.qubit', wires=2))
def circuit(coeffs):
qml.exp(1j * coeffs[0] * qml.X(0) @ qml.X(1) + 1j * coeffs[1] * qml.Y(0) @ qml.Y(1))
return qml.expval(qml.Z(0))
jax.jit(circuit)(jax.numpy.array([1.0, 2.0]))Tracebacks
---------------------------------------------------------------------------
TracerBoolConversionError Traceback (most recent call last)
Cell In[8], line 6
3 qml.exp(1j * coeffs[0] * qml.X(0) @ qml.X(1) + 1j * coeffs[1] * qml.Y(0) @ qml.Y(1))
4 return qml.expval(qml.Z(0))
----> 6 jax.jit(circuit)(jax.numpy.array([1.0, 2.0]))
[... skipping hidden 11 frame]
File ~/Prog/pennylane/pennylane/workflow/qnode.py:843, in QNode.__call__(self, *args, **kwargs)
840 from ._capture_qnode import capture_qnode # pylint: disable=import-outside-toplevel
842 return capture_qnode(self, *args, **kwargs)
--> 843 return self._impl_call(*args, **kwargs)
File ~/Prog/pennylane/pennylane/workflow/qnode.py:817, in QNode._impl_call(self, *args, **kwargs)
814 # Calculate the classical jacobians if necessary
815 self._transform_program.set_classical_component(self, args, kwargs)
--> 817 res = qml.execute(
818 (tape,),
819 device=self.device,
820 diff_method=self.diff_method,
821 interface=interface,
822 transform_program=self._transform_program,
823 gradient_kwargs=self.gradient_kwargs,
824 **self.execute_kwargs,
825 )
826 res = res[0]
828 # convert result to the interface in case the qfunc has no parameters
File ~/Prog/pennylane/pennylane/workflow/execution.py:216, in execute(tapes, device, diff_method, interface, grad_on_execution, cache, cachesize, max_diff, device_vjp, gradient_kwargs, transform_program, mcm_config, config, inner_transform)
205 # Only need to calculate derivatives with jax when we know it will be executed later.
207 config = qml.devices.ExecutionConfig(
208 interface=interface,
209 gradient_method=diff_method,
(...)
214 derivative_order=max_diff,
215 )
--> 216 config = _resolve_execution_config(config, device, tapes, transform_program=transform_program)
218 config = replace(
219 config,
220 interface=interface,
221 derivative_order=max_diff,
222 )
224 transform_program = transform_program or qml.transforms.core.TransformProgram()
File ~/Prog/pennylane/pennylane/workflow/resolution.py:257, in _resolve_execution_config(execution_config, device, tapes, transform_program)
250 if (
251 "lightning" in device.name
252 and transform_program
253 and qml.metric_tensor in transform_program
254 and execution_config.gradient_method == "best"
255 ):
256 execution_config = replace(execution_config, gradient_method=qml.gradients.param_shift)
--> 257 execution_config = _resolve_diff_method(execution_config, device, tape=tapes[0])
259 if execution_config.use_device_jacobian_product and not device.supports_vjp(
260 execution_config, tapes[0]
261 ):
262 raise qml.QuantumFunctionError(
263 f"device_vjp=True is not supported for device {device},"
264 f" diff_method {execution_config.gradient_method},"
265 " and the provided circuit."
266 )
File ~/Prog/pennylane/pennylane/logging/decorators.py:61, in log_string_debug_func.<locals>.wrapper_entry(*args, **kwargs)
54 s_caller = "::L".join(
55 [str(i) for i in inspect.getouterframes(inspect.currentframe(), 2)[1][1:3]]
56 )
57 lgr.debug(
58 f"Calling {f_string} from {s_caller}",
59 **_debug_log_kwargs,
60 )
---> 61 return func(*args, **kwargs)
File ~/Prog/pennylane/pennylane/workflow/resolution.py:190, in _resolve_diff_method(initial_config, device, tape)
187 if diff_method is None:
188 return initial_config
--> 190 if device.supports_derivatives(initial_config, circuit=tape):
191 new_config = device.setup_execution_config(initial_config)
192 return new_config
File ~/Prog/pl/lib/python3.12/site-packages/pennylane_lightning/lightning_qubit/lightning_qubit.py:427, in LightningQubit.supports_derivatives(self, execution_config, circuit)
425 if circuit is None:
426 return True
--> 427 return _supports_adjoint(circuit=circuit)
File ~/Prog/pl/lib/python3.12/site-packages/pennylane_lightning/lightning_qubit/lightning_qubit.py:133, in _supports_adjoint(circuit)
130 _add_adjoint_transforms(prog)
132 try:
--> 133 prog((circuit,))
134 except (DecompositionUndefinedError, qml.DeviceError, AttributeError):
135 return False
File ~/Prog/pennylane/pennylane/transforms/core/transform_program.py:580, in TransformProgram.__call__(self, tapes)
578 if argnums is not None:
579 tape.trainable_params = argnums[j]
--> 580 new_tapes, fn = transform(tape, *targs, **tkwargs)
581 execution_tapes.extend(new_tapes)
583 fns.append(fn)
File ~/Prog/pennylane/pennylane/devices/preprocess.py:404, in decompose(tape, stopping_condition, stopping_condition_shots, skip_initial_state_prep, decomposer, name, error)
398 return (tape,), null_postprocessing
399 try:
401 new_ops = [
402 final_op
403 for op in tape.operations[len(prep_op) :]
--> 404 for final_op in _operator_decomposition_gen(
405 op,
406 stopping_condition,
407 decomposer=decomposer,
408 name=name,
409 error=error,
410 )
411 ]
412 except RecursionError as e:
413 raise error(
414 "Reached recursion limit trying to decompose operations. "
415 "Operator decomposition may have entered an infinite loop."
416 ) from e
File ~/Prog/pennylane/pennylane/devices/preprocess.py:64, in _operator_decomposition_gen(op, acceptance_function, decomposer, max_expansion, current_depth, name, error)
62 else:
63 try:
---> 64 decomp = decomposer(op)
65 current_depth += 1
66 except qml.operation.DecompositionUndefinedError as e:
File ~/Prog/pennylane/pennylane/devices/preprocess.py:387, in decompose.<locals>.decomposer(op)
386 def decomposer(op):
--> 387 return op.decomposition()
File ~/Prog/pennylane/pennylane/ops/op_math/exp.py:241, in Exp.decomposition(self)
229 r"""Representation of the operator as a product of other operators. Decomposes into
230 :class:`~.PauliRot` if the coefficient is imaginary and the base is a Pauli Word.
231
(...)
238 list[PauliRot]: decomposition of the operator
239 """
240 with qml.QueuingManager.stop_recording():
--> 241 d = self._recursive_decomposition(self.base, self.coeff)
243 if qml.QueuingManager.recording():
244 for op in d:
File ~/Prog/pennylane/pennylane/ops/op_math/exp.py:282, in Exp._recursive_decomposition(self, base, coeff)
276 if not self.num_steps: # if num_steps was not set
277 error_msg += (
278 " Please set a value to ``num_steps`` when instantiating the ``Exp`` operator "
279 "if a Suzuki-Trotter decomposition is required."
280 )
--> 282 if self.base.is_hermitian:
283 error_msg += (
284 " Decomposition is not defined for real coefficients of hermitian operators."
285 )
287 raise DecompositionUndefinedError(error_msg)
File ~/Prog/pennylane/pennylane/ops/op_math/composite.py:37, in handle_recursion_error.<locals>.wrapper(*args, **kwargs)
34 @wraps(func)
35 def wrapper(*args, **kwargs):
36 try:
---> 37 return func(*args, **kwargs)
38 except RecursionError as e:
39 raise RuntimeError(
40 "Maximum recursion depth reached! This is likely due to nesting too many levels "
41 "of composite operators. Try setting lazy=False when calling qml.sum, qml.prod, "
42 "and qml.s_prod, or use the +, @, and * operators instead. Alternatively, you "
43 "can periodically call qml.simplify on your operators."
44 ) from e
File ~/Prog/pennylane/pennylane/ops/op_math/sum.py:309, in Sum.is_hermitian(self)
306 if not math.is_abstract(coeffs_list[0]):
307 return not any(math.iscomplex(c) for c in coeffs_list)
--> 309 return all(s.is_hermitian for s in self)
File ~/Prog/pennylane/pennylane/ops/op_math/sum.py:309, in <genexpr>(.0)
306 if not math.is_abstract(coeffs_list[0]):
307 return not any(math.iscomplex(c) for c in coeffs_list)
--> 309 return all(s.is_hermitian for s in self)
File ~/Prog/pennylane/pennylane/ops/op_math/prod.py:246, in Prod.is_hermitian(self)
244 if qml.wires.Wires.shared_wires([o1.wires, o2.wires]):
245 return False
--> 246 return all(op.is_hermitian for op in self)
File ~/Prog/pennylane/pennylane/ops/op_math/prod.py:246, in <genexpr>(.0)
244 if qml.wires.Wires.shared_wires([o1.wires, o2.wires]):
245 return False
--> 246 return all(op.is_hermitian for op in self)
File ~/Prog/pennylane/pennylane/ops/op_math/composite.py:37, in handle_recursion_error.<locals>.wrapper(*args, **kwargs)
34 @wraps(func)
35 def wrapper(*args, **kwargs):
36 try:
---> 37 return func(*args, **kwargs)
38 except RecursionError as e:
39 raise RuntimeError(
40 "Maximum recursion depth reached! This is likely due to nesting too many levels "
41 "of composite operators. Try setting lazy=False when calling qml.sum, qml.prod, "
42 "and qml.s_prod, or use the +, @, and * operators instead. Alternatively, you "
43 "can periodically call qml.simplify on your operators."
44 ) from e
File ~/Prog/pennylane/pennylane/ops/op_math/sprod.py:209, in SProd.is_hermitian(self)
204 @property
205 @handle_recursion_error
206 def is_hermitian(self):
207 """If the base operator is hermitian and the scalar is real,
208 then the scalar product operator is hermitian."""
--> 209 return self.base.is_hermitian and not qml.math.iscomplex(self.scalar)
[... skipping hidden 1 frame]
File ~/Prog/pl/lib/python3.12/site-packages/jax/_src/core.py:1475, in concretization_function_error.<locals>.error(self, arg)
1474 def error(self, arg):
-> 1475 raise TracerBoolConversionError(arg)
TracerBoolConversionError: Attempted boolean conversion of traced array with shape bool[]..
The error occurred while tracing the function circuit at /var/folders/k1/0v_kvphn55lgf_45kntf1hqm0000gq/T/ipykernel_14463/3587553121.py:1 for jit. This concrete value was not available in Python because it depends on the value of the argument coeffs.
See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerBoolConversionErrorSystem information
Name: PennyLane
Version: 0.40.0
Summary: PennyLane is a cross-platform Python library for quantum computing, quantum machine learning, and quantum chemistry. Train a quantum computer the same way as a neural network.
Home-page: https://github.com/PennyLaneAI/pennylane
Author:
Author-email:
License: Apache License 2.0
Location: /usr/local/lib/python3.11/dist-packages
Requires: appdirs, autograd, autoray, cachetools, diastatic-malt, networkx, numpy, packaging, pennylane-lightning, requests, rustworkx, scipy, tomlkit, typing-extensions
Required-by: PennyLane-Catalyst, PennyLane_Lightning
Platform info: Linux-6.1.85+-x86_64-with-glibc2.35
Python version: 3.11.11
Numpy version: 1.26.4
Scipy version: 1.13.1
Installed devices:
- lightning.qubit (PennyLane_Lightning-0.40.0)
- nvidia.custatevec (PennyLane-Catalyst-0.10.0)
- nvidia.cutensornet (PennyLane-Catalyst-0.10.0)
- oqc.cloud (PennyLane-Catalyst-0.10.0)
- softwareq.qpp (PennyLane-Catalyst-0.10.0)
- default.clifford (PennyLane-0.40.0)
- default.gaussian (PennyLane-0.40.0)
- default.mixed (PennyLane-0.40.0)
- default.qubit (PennyLane-0.40.0)
- default.qutrit (PennyLane-0.40.0)
- default.qutrit.mixed (PennyLane-0.40.0)
- default.tensor (PennyLane-0.40.0)
- null.qubit (PennyLane-0.40.0)
- reference.qubit (PennyLane-0.40.0)Existing GitHub issues
- I have searched existing GitHub issues to make sure the issue does not already exist.