fixed catalyst to correctly extract number of shots for legacyDeviceFacade (#1035)

**Context:**
This PR is to fix an issue that will be caused after merging the
following PR:
PennyLaneAI/pennylane#6046

The problem is that LegacyDeviceFacade will cause the Legacy Devices to
use the Shots class instead of an integer which causes a crash in
runtime.

**Description of the Change:**
This PR ensures that the any device that is an instance of the new
device API uses the shots class to extract the total number of shots.

**Related GitHub Issues:**
[sc-70792]

---------

Co-authored-by: Ahmed Darwish <[email protected]>

Author: mehrdad2m

.dep-versions

@@ -8,11 +8,11 @@ enzyme=v0.0.130
 
 # For a custom PL version, update the package version here and at
 # 'doc/requirements.txt
-pennylane=0.38.0.dev11
+pennylane=0.38.0.dev21
 
 # For a custom LQ/LK version, update the package version here and at
 # 'doc/requirements.txt'. Also, update the 'LIGHTNING_GIT_TAG' at
 # 'runtime/Makefile' and at all GitHub workflows, using the exact
 # commit hash corresponding to the merged PR that implements the
 # desired feature.
-lightning=0.38.0-dev34
+lightning=0.38.0-dev50

doc/changelog.md

@@ -456,6 +456,9 @@
 
 <h3>Bug fixes</h3>
 
+* Fix a bug where LegacyDevice number of shots is not correctly extracted when using the legacyDeviceFacade.
+  [(#1035)](https://github.com/PennyLaneAI/catalyst/pull/1035)
+
 * Catalyst no longer generates a `QubitUnitary` operation during decomposition if a device doesn't
   support it. Instead, the operation that would lead to a `QubitUnitary` is either decomposed or
   raises an error.

doc/requirements.txt

@@ -30,6 +30,6 @@ lxml_html_clean
 
 # Pre-install PL development wheels
 --extra-index-url https://test.pypi.org/simple/
-pennylane-lightning-kokkos==0.38.0-dev34
-pennylane-lightning==0.38.0-dev34
-pennylane==0.38.0.dev11
+pennylane-lightning-kokkos==0.38.0-dev50
+pennylane-lightning==0.38.0-dev50
+pennylane==0.38.0.dev21

frontend/catalyst/device/qjit_device.py

@@ -44,6 +44,7 @@
     verify_operations,
 )
 from catalyst.logging import debug_logger, debug_logger_init
+from catalyst.third_party.cuda import SoftwareQQPP
 from catalyst.utils.exceptions import CompileError
 from catalyst.utils.patching import Patcher
 from catalyst.utils.runtime_environment import get_lib_path
@@ -133,7 +134,7 @@
 
 def get_device_shots(dev):
     """Helper function to get device shots."""
-    return dev.shots if isinstance(dev, qml.devices.LegacyDevice) else dev.shots.total_shots
+    return dev.shots.total_shots if isinstance(dev, qml.devices.Device) else dev.shots
 
 
 @dataclass
@@ -153,8 +154,8 @@ def extract_backend_info(device: qml.QubitDevice, capabilities: DeviceCapabiliti
     to a valid TOML config file."""
 
     dname = device.name
-    if isinstance(device, qml.devices.LegacyDevice):
-        dname = device.short_name
+    if isinstance(device, qml.devices.LegacyDeviceFacade):
+        dname = device.target_device.short_name
 
     device_name = ""
     device_lpath = ""
@@ -189,18 +190,18 @@ def extract_backend_info(device: qml.QubitDevice, capabilities: DeviceCapabiliti
         device_kwargs["device_type"] = dname
         device_kwargs["backend"] = (
             # pylint: disable=protected-access
-            device._device._delegate.DEVICE_ID
+            device.target_device._device._delegate.DEVICE_ID
         )
     elif dname == "braket.aws.qubit":  # pragma: no cover
         device_kwargs["device_type"] = dname
         device_kwargs["device_arn"] = device._device._arn  # pylint: disable=protected-access
-        if device._s3_folder:  # pylint: disable=protected-access
+        if device.target_device._s3_folder:  # pylint: disable=protected-access
             device_kwargs["s3_destination_folder"] = str(
-                device._s3_folder  # pylint: disable=protected-access
+                device.target_device._s3_folder  # pylint: disable=protected-access
             )
 
     for k, v in capabilities.options.items():
-        if hasattr(device, v):
+        if hasattr(device, v) and not k in device_kwargs:
             device_kwargs[k] = getattr(device, v)
 
     return BackendInfo(dname, device_name, device_lpath, device_kwargs)
@@ -506,6 +507,8 @@ def preprocess(
     def _measurement_transform_program(self):
 
         measurement_program = TransformProgram()
+        if isinstance(self.original_device, SoftwareQQPP):
+            return measurement_program
 
         supports_sum_observables = any(
             obs in self.qjit_capabilities.native_obs for obs in ("Sum", "Hamiltonian")

frontend/catalyst/qfunc.py

@@ -253,8 +253,8 @@ def processing_fn(results):
     total_shots = get_device_shots(dev)
 
     new_dev = copy(dev)
-    if isinstance(new_dev, qml.devices.LegacyDevice):
-        new_dev.shots = 1  # pragma: no cover
+    if isinstance(new_dev, qml.devices.LegacyDeviceFacade):
+        new_dev.target_device.shots = 1  # pragma: no cover
     else:
         new_dev._shots = qml.measurements.Shots(1)
     single_shot_qnode.device = new_dev

frontend/catalyst/third_party/cuda/__init__.py

@@ -127,7 +127,10 @@ class BaseCudaInstructionSet(qml.QubitDevice):
         "SWAP",
         "CSWAP",
     ]
-    observables = []
+    observables = [
+        "PauliX",
+        "PauliZ",
+    ]
     config = Path(__file__).parent / "cuda_quantum.toml"
 
     def __init__(self, shots=None, wires=None):

frontend/catalyst/third_party/cuda/catalyst_to_cuda_interpreter.py

@@ -825,9 +825,16 @@ def cudaq_backend_info(device, _capabilities) -> BackendInfo:
             catalyst-specific. We need to make this API a bit nicer for third-party compilers.
             """
             device_name = (
-                device.short_name if isinstance(device, qml.devices.LegacyDevice) else device.name
+                device.target_device.short_name
+                if isinstance(device, qml.devices.LegacyDeviceFacade)
+                else device.name
             )
-            return BackendInfo(device_name, device.name, "", {})
+            interface_name = (
+                device.target_device.name
+                if isinstance(device, qml.devices.LegacyDeviceFacade)
+                else device.name
+            )
+            return BackendInfo(device_name, interface_name, "", {})
 
         with Patcher(
             (QFunc, "extract_backend_info", cudaq_backend_info),

frontend/catalyst/third_party/cuda/cuda_quantum.toml

@@ -34,9 +34,17 @@ CSWAP = { properties = [ "invertible" ] }
 
 # Observables supported natively by the device
 [operators.observables]
+PauliX = {}
+PauliZ = {}
+Sum = {}
 
 [measurement_processes]
 
+Expval                 = {}
+State                  = { condition = [ "analytic" ] }
+Sample                 = { condition = [ "finiteshots" ] }
+Counts                  = { condition = [ "finiteshots" ] }
+
 [compilation]
 # If the device is compatible with qjit
 qjit_compatible = true

frontend/test/pytest/test_cuda_integration.py

@@ -68,7 +68,7 @@ def test_measurement_side_effect(self):
         from catalyst.third_party.cuda import cudaqjit as cjit
 
         @cjit
-        @qml.qnode(qml.device("softwareq.qpp", wires=1, shots=30))
+        @qml.qnode(qml.device("softwareq.qpp", wires=1, shots=None))
         def circuit():
             qml.RX(jnp.pi / 4, wires=[0])
             measure(0)

frontend/test/pytest/test_measurements_shots_results.py

@@ -428,7 +428,13 @@ def circuit(theta, phi, varphi):
             qml.CNOT(wires=[1, 2])
             return qml.var(0.2 * qml.PauliZ(wires=0) + 0.5 * qml.Hadamard(wires=1))
 
-        if isinstance(dev, qml.devices.Device):
+        if isinstance(dev, qml.devices.LegacyDeviceFacade):
+            with pytest.raises(
+                RuntimeError,
+                match=r"Cannot split up terms in sums for MeasurementProcess <class 'pennylane.measurements.var.VarianceMP'>",
+            ):
+                circuit(0.432, 0.123, -0.543)
+        else:
             # TODO: only raises with the new API, Kokkos should also raise an error.
             with pytest.raises(
                 TypeError,
@@ -488,10 +494,7 @@ def test_missing_shots_value(self, backend, meas_fun):
         def circuit():
             return meas_fun(wires=0)
 
-        # ValueError is legacy behaviour with the old device API
-        error_type = ValueError if isinstance(dev, qml.devices.LegacyDevice) else CompileError
-
-        with pytest.raises(error_type, match="cannot work with shots=None"):
+        with pytest.raises(CompileError, match="cannot work with shots=None"):
             qjit(circuit)
 
     def test_multiple_return_values(self, backend, tol_stochastic):