Detecting mismatching operations with Catalyst

[erick-xanadu]

Describe the bug
Hi,

I am working on Catalyst which is JIT compiler for PennyLane. We have an integration with Amazon Braket that allows users to write quantum circuits expressed as PennyLane's tapes and generate OpenQASM which is then sent to the Amazon Braket for execution.

Catalyst differs a little bit from PennyLane. In particular, Catalyst uses a TOML file to specify which gates and observables are available on the target device. For example, this is a TOML file we use for testing a test device.. PennyLane, on the other hand uses the operations and observables property on the instance of the device being executed. This sometimes leads to a conflict that may result in errors.

In particular amazon-braket-pennylane-plugins 1.27.1 release contained a modification to BraketQubitDevice's observables property (patch inlined below):

diff --git a/src/braket/pennylane_plugin/braket_device.py b/src/braket/pennylane_plugin/braket_device.py
index a45de88..f3866c3 100644
--- a/src/braket/pennylane_plugin/braket_device.py
+++ b/src/braket/pennylane_plugin/braket_device.py
@@ -162,7 +162,7 @@ class BraketQubitDevice(QubitDevice):
 
     @property
     def observables(self) -> frozenset[str]:
-        base_observables = frozenset(super().observables)
+        base_observables = frozenset(super().observables - {"SProd", "Sum"})
         # This needs to be here bc expectation(ax+by)== a*expectation(x)+b*expectation(y)
         # is only true when shots=0
         if not self.shots:

The commit removes SProd and Sum observables from the list of supported observables by the device. However, SProd and Sum were not removed Catalyst's TOML file which corresponds to the local braket simulator. And this lead to errors.

The error went away with the release of amazon-braket-pennylane-plugin version 1.27.2 which added back the Sum and SProd observables to the list of supported observables. However, it would be nice if we could guarantee that amazon-braket-pennylane-plugin's list of supported observables and operations doesn't go out of sync with Catalyst's.

To reproduce
pip install pennylane-catalyst amazon-braket-pennylane-plugin==1.27.1

import pennylane as qml

@qml.qjit
@qml.qnode(qml.device("braket.local.qubit", wires=2))
def all_measurements(x):
    qml.RY(x, wires=0)
    return (
        qml.expval(qml.PauliZ(0) @ qml.PauliZ(1)),
        qml.var(qml.PauliZ(1) @ qml.PauliZ(0)),
    )

all_measurements(3.14)

will produce the following error:

catalyst.utils.exceptions.CompileError: Observables in qml.device.observables and specification file do not match.
Observables that present only in spec: {'Sum', 'SProd'}

Expected behavior
It should not throw the CompileError and return instead the following values:

(array(-0.99999873), array(2.53654331e-06))

System information
A description of your system. Please provide:

• Amazon Braket Python PennyLane Plugin version: 1.27.1
• Amazon Braket Python SDK version: amazon-braket-sdk==1.78.0
• Amazon Braket Python Schemas version: amazon-braket-sdk==1.78.0
• Amazon Braket Python Default Simulator version:
• Python version: 3.10

[virajvchaudhari]

Hello @erick-xanadu, thank you for bringing this issue to our attention, we will look into this. Do you have any recommendations that you want to prescribe to help us keep these dependencies in sync?

[erick-xanadu]

Hi @virajvchaudhari, we are still working on it. We have an idea of how to fix it, but we want to wait until after some design changes happen on our repository. Will keep you updated. Thanks!