Dynamic one shot does not work with ZNE

[dime10]

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'

[dime10]

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.

[dime10]

@dime10 Check if fixed

[dime10]

Can confirm this is now fixed, test here:

catalyst/frontend/test/lit/test_mitigation.py

Line 28 in f34646e

def mcm_method_with_zne():