PennyLaneAI / PennyLaneAI/catalyst

[BUG] parameterizing coefficients of `qml.dot` fails with qjit

Open
#1,464 2 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

enhancement
Dominant language
Python
Stars
234
Forks
84
Avg merge
2d 15h
Merged PRs (30d)
66

Description

Experienced TracerBoolConversionError when trying to parameterize coeffients of qml.dot with qjit. Minimal example code below.

Expected behavior: (without qjit)

import pennylane as qml
import jax.numpy as jnp
from catalyst import qjit

@qml.qnode(qml.device('lightning.qubit', wires=1))
def circuit(x):
    coeffs = [x]
    ops = [qml.PauliX(0)]
    H = qml.dot(coeffs, ops) # H = qml.Hamiltonian(coeffs, ops)
    qml.evolve(H)
    return qml.expval(qml.Z(0))

print(circuit(jnp.array([0.5])))

Output:

>>>  [0.54030231]

Actual Behaviour: (with qjit)

import pennylane as qml
import jax.numpy as jnp
from catalyst import qjit

@qjit(autograph=True)
@qml.qnode(qml.device('lightning.qubit', wires=1))
def circuit(x):
    coeffs = [x]
    ops = [qml.PauliX(0)]
    H = qml.dot(coeffs, ops) # H = qml.Hamiltonian(coeffs, ops)
    qml.evolve(H)
    return qml.expval(qml.Z(0))

print(circuit(jnp.array([0.5])))

Output:

---------------------------------------------------------------------------
TracerBoolConversionError                 Traceback (most recent call last)
Cell In[7], line 14
     11     qml.evolve(H)
     12     return qml.expval(qml.Z(0))
---> 14 print(circuit(jnp.array([0.5])))

File ~/anaconda3/envs/pennylane40/lib/python3.12/site-packages/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 ~/anaconda3/envs/pennylane40/lib/python3.12/site-packages/catalyst/jit.py:521, in QJIT.__call__(self, *args, **kwargs)
    517         kwargs = {"static_argnums": self.compile_options.static_argnums, **kwargs}
    519     return self.user_function(*args, **kwargs)
--> 521 requires_promotion = self.jit_compile(args, **kwargs)
    523 # If we receive tracers as input, dispatch to the JAX integration.
    524 if any(isinstance(arg, jax.core.Tracer) for arg in tree_flatten(args)[0]):
...
File ~/anaconda3/envs/pennylane40/lib/python3.12/site-packages/pennylane/ops/functions/dot.py:171, in dot(coeffs, ops, pauli, grouping_type, method)
    168 # Convert possible PauliWord and PauliSentence instances to operation
    169 ops = [op.operation() if isinstance(op, (PauliWord, PauliSentence)) else op for op in ops]
--> 171 operands = [op if coeff == 1 else qml.s_prod(coeff, op) for coeff, op in zip(coeffs, ops)]
    172 return (
    173     operands[0]
    174     if len(operands) == 1
    175     else qml.sum(*operands, grouping_type=grouping_type, method=method)
    176 )

File ~/anaconda3/envs/pennylane40/lib/python3.12/site-packages/jax/_src/core.py:712, in Tracer.__bool__(self)
    710 def __bool__(self):
    711   check_bool_conversion(self)
--> 712   return self.aval._bool(self)

File ~/anaconda3/envs/pennylane40/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[1]..
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: /home/nishanth/anaconda3/envs/pennylane40/lib/python3.12/site-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-4.4.0-19041-Microsoft-x86_64-with-glibc2.35
Python version:          3.12.8
Numpy version:           2.0.2
Scipy version:           1.15.1
Installed devices:
- 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)
- 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)
- lightning.qubit (PennyLane_Lightning-0.40.0)

Additional Information:
Replacing qml.dot with qml.Hamiltonian gives the same error.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with the minimal qjit example in the issue and inspect the implementation around line 171 of pennylane/ops/functions/dot.py, where the traceback shows the coefficient comparison fails. Compare this path with qml.Hamiltonian, which shows the same error. Done means parameterized coefficients work under qjit without TracerBoolConversionError and preserve the expected result.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
compilers
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.