[BUG] qml.exp with jax.jit causes TracerBoolConversionError

[justinpickering]

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

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

[justinpickering]

reported by @albi3ro.