Error in Catalyst MitigationLoweringPass when shots argument is specified

[cosenal]

When the shots argument is passed to the device, the following snippet raises an error.

Code

import pennylane as qml
from catalyst import qjit, mitigate_with_zne
from pennylane.transforms import richardson_extrapolate


dev = qml.device("qrack.simulator", wires=2, shots=1)

@qml.qnode(dev)
def circuit():
    qml.Hadamard(wires=0)
    qml.CNOT(wires=[0, 1])
    return qml.expval(qml.PauliY(wires=0))

@qjit
def mitigated_circuit():
    return mitigate_with_zne(
        circuit, 
        scale_factors=[1, 3, 5], 
        extrapolate=richardson_extrapolate,
        folding="global"
    )()

mitigated_circuit()

Traceback

Traceback (most recent call last):
  File "/Users/cosenal/Coding/PennyLaneAI/catalyst/frontend/catalyst/compiler.py", line 559, in run_from_ir
    compiler_output = run_compiler_driver(
                      ^^^^^^^^^^^^^^^^^^^^
RuntimeError: Compilation failed:
mitigated_circuit:16:11: error: 'func.func' op entry block must have 1 arguments to match function signature
    %10 = "mitigation.zne"(%9) {callee = @circuit, folding = #mitigation<folding global>} : (tensor<3xi64>) -> tensor<3xf64>
          ^
mitigated_circuit:16:11: note: see current operation: 
"func.func"() <{function_type = (!quantum.reg) -> !quantum.reg, sym_name = "circuit.withoutMeasurements", sym_visibility = "private"}> ({
^bb0(%arg0: tensor<1xi64>, %arg1: !quantum.reg):
  %0 = "arith.constant"() <{value = 1.5707963267948966 : f64}> : () -> f64
  %1 = "quantum.extract"(%arg1) <{idx_attr = 0 : i64}> : (!quantum.reg) -> !quantum.bit
  %2 = "quantum.custom"(%1) <{gate_name = "Hadamard", operandSegmentSizes = array<i32: 0, 1, 0, 0>, resultSegmentSizes = array<i32: 1, 0>}> : (!quantum.bit) -> !quantum.bit
  %3 = "quantum.extract"(%arg1) <{idx_attr = 1 : i64}> : (!quantum.reg) -> !quantum.bit
  %4:2 = "quantum.custom"(%2, %3) <{gate_name = "CNOT", operandSegmentSizes = array<i32: 0, 2, 0, 0>, resultSegmentSizes = array<i32: 2, 0>}> : (!quantum.bit, !quantum.bit) -> (!quantum.bit, !quantum.bit)
  %5 = "quantum.custom"(%0, %4#0) <{gate_name = "RX", operandSegmentSizes = array<i32: 1, 1, 0, 0>, resultSegmentSizes = array<i32: 1, 0>}> : (f64, !quantum.bit) -> !quantum.bit
  %6 = "quantum.insert"(%arg1, %5) <{idx_attr = 0 : i64}> : (!quantum.reg, !quantum.bit) -> !quantum.reg
  %7 = "quantum.insert"(%6, %4#1) <{idx_attr = 1 : i64}> : (!quantum.reg, !quantum.bit) -> !quantum.reg
  "func.return"(%7) : (!quantum.reg) -> ()
}) : () -> ()
While processing 'MitigationLoweringPass' pass of the 'QuantumCompilationPass' pipeline
Failed to lower MLIR module


The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/Users/cosenal/Coding/PennyLaneAI/catalyst-playground/test_zne_qrack.py", line 14, in <module>
    @qjit
     ^^^^
  File "/opt/miniconda3/envs/catalyst/lib/python3.11/site-packages/pennylane/logging/decorators.py", line 61, in wrapper_entry
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/Users/cosenal/Coding/PennyLaneAI/catalyst/frontend/catalyst/jit.py", line 377, in qjit
    return QJIT(fn, CompileOptions(**kwargs))
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/miniconda3/envs/catalyst/lib/python3.11/site-packages/pennylane/logging/decorators.py", line 65, in wrapper_exit
    output = func(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^
  File "/Users/cosenal/Coding/PennyLaneAI/catalyst/frontend/catalyst/jit.py", line 445, in __init__
    self.aot_compile()
  File "/opt/miniconda3/envs/catalyst/lib/python3.11/site-packages/pennylane/logging/decorators.py", line 61, in wrapper_entry
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/Users/cosenal/Coding/PennyLaneAI/catalyst/frontend/catalyst/jit.py", line 491, in aot_compile
    self.compiled_function, self.qir = self.compile()
                                       ^^^^^^^^^^^^^^
  File "/Users/cosenal/Coding/PennyLaneAI/catalyst/frontend/catalyst/debug/instruments.py", line 143, in wrapper
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/opt/miniconda3/envs/catalyst/lib/python3.11/site-packages/pennylane/logging/decorators.py", line 61, in wrapper_entry
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/Users/cosenal/Coding/PennyLaneAI/catalyst/frontend/catalyst/jit.py", line 672, in compile
    shared_object, llvm_ir, _ = self.compiler.run(self.mlir_module, self.workspace)
                                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/miniconda3/envs/catalyst/lib/python3.11/site-packages/pennylane/logging/decorators.py", line 61, in wrapper_entry
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/Users/cosenal/Coding/PennyLaneAI/catalyst/frontend/catalyst/compiler.py", line 607, in run
    return self.run_from_ir(
           ^^^^^^^^^^^^^^^^^
  File "/opt/miniconda3/envs/catalyst/lib/python3.11/site-packages/pennylane/logging/decorators.py", line 61, in wrapper_entry
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/Users/cosenal/Coding/PennyLaneAI/catalyst/frontend/catalyst/compiler.py", line 571, in run_from_ir
    raise CompileError(*e.args) from e
catalyst.utils.exceptions.CompileError: Compilation failed:
mitigated_circuit:16:11: error: 'func.func' op entry block must have 1 arguments to match function signature
    %10 = "mitigation.zne"(%9) {callee = @circuit, folding = #mitigation<folding global>} : (tensor<3xi64>) -> tensor<3xf64>
          ^
mitigated_circuit:16:11: note: see current operation: 
"func.func"() <{function_type = (!quantum.reg) -> !quantum.reg, sym_name = "circuit.withoutMeasurements", sym_visibility = "private"}> ({
^bb0(%arg0: tensor<1xi64>, %arg1: !quantum.reg):
  %0 = "arith.constant"() <{value = 1.5707963267948966 : f64}> : () -> f64
  %1 = "quantum.extract"(%arg1) <{idx_attr = 0 : i64}> : (!quantum.reg) -> !quantum.bit
  %2 = "quantum.custom"(%1) <{gate_name = "Hadamard", operandSegmentSizes = array<i32: 0, 1, 0, 0>, resultSegmentSizes = array<i32: 1, 0>}> : (!quantum.bit) -> !quantum.bit
  %3 = "quantum.extract"(%arg1) <{idx_attr = 1 : i64}> : (!quantum.reg) -> !quantum.bit
  %4:2 = "quantum.custom"(%2, %3) <{gate_name = "CNOT", operandSegmentSizes = array<i32: 0, 2, 0, 0>, resultSegmentSizes = array<i32: 2, 0>}> : (!quantum.bit, !quantum.bit) -> (!quantum.bit, !quantum.bit)
  %5 = "quantum.custom"(%0, %4#0) <{gate_name = "RX", operandSegmentSizes = array<i32: 1, 1, 0, 0>, resultSegmentSizes = array<i32: 1, 0>}> : (f64, !quantum.bit) -> !quantum.bit
  %6 = "quantum.insert"(%arg1, %5) <{idx_attr = 0 : i64}> : (!quantum.reg, !quantum.bit) -> !quantum.reg
  %7 = "quantum.insert"(%6, %4#1) <{idx_attr = 1 : i64}> : (!quantum.reg, !quantum.bit) -> !quantum.reg
  "func.return"(%7) : (!quantum.reg) -> ()
}) : () -> ()
While processing 'MitigationLoweringPass' pass of the 'QuantumCompilationPass' pipeline
Failed to lower MLIR module

Notes

1. This doesn't happen with the lightning.qubit simulator
2. The error occurs no matter the value of shots, so I suspect it's about how the argument is lowered to the MLIR functions.
3. It only happens when lowering the ZNE mitigation function, and now, for example, when qjiting the original circuit.