Repository navigation
Dynamic one shot does not work with ZNE #929
Description
import jax.numpy as jnp
import pennylane as qml
from catalyst import *
@qml.qnode(dev, mcm_method="one-shot")
def circuit():
qml.Hadamard(wires=0)
qml.CNOT(wires=[0, 1])
qml.Hadamard(wires=0)
qml.CNOT(wires=[0, 1])
qml.Hadamard(wires=0)
return qml.expval(qml.PauliY(0))
@qjit
def mitigated_circuit():
s = jax.numpy.array([1, 2])
return mitigate_with_zne(circuit, scale_factors=s)() File "catalyst/frontend/catalyst/jax_primitives.py", line 739, in _zne_lowering
_func_lowering(ctx, *args, call_jaxpr=jaxpr.eqns[0].params["call_jaxpr"], fn=fn, call=False)
~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^
KeyError: 'call_jaxpr'
Activity
dime10 commentedon Jul 15, 2024
Briefly looked into this issue, and the root cause is that the ZNE wrapper only works directly on QNodes, but the dynamic_one_shot transform replaces a QNode with a classical function that invokes a QNode, thus breaking the ZNE assumption.
I think a solution to this problem might be to rewrite ZNE to be part of the QNode transform program, that way it would be properly propagated by the dynamic_one_shot method. Haven't verified if this would actually work though.
Catalyst support for `qml.CosineWindow` (#1166)
dime10 commentedon Apr 1, 2025
@dime10 Check if fixed
dime10 commentedon Apr 1, 2025
Can confirm this is now fixed, test here:
catalyst/frontend/test/lit/test_mitigation.py
Line 28 in f34646e