Update the LQ new device API to work with Catalyst

.github/CHANGELOG.md

@@ -86,6 +86,9 @@
 * Fix the failed observable serialization unit tests.
   [(#683)](https://github.com/PennyLaneAI/pennylane-lightning/pull/683)
 
+* Update the `LightningQubit` new device API to work with Catalyst.
+  [(#665)](https://github.com/PennyLaneAI/pennylane-lightning/pull/665)
+
 ### Contributors
 
 This release contains contributions from (in alphabetical order):

pennylane_lightning/core/_version.py

@@ -16,4 +16,4 @@
    Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.36.0-dev29"
+__version__ = "0.36.0-dev30"

pennylane_lightning/core/src/utils/config.h

@@ -14,10 +14,10 @@
 
 /**
  * @file
- * Record the path to scipy.libs at compile time.
+ * Config file for the path to scipy.libs at compile time.
  */
 
 #ifndef CONFIG_H
 #define CONFIG_H
 #define SCIPY_LIBS_PATH ""
-#endif
+#endif

pennylane_lightning/lightning_kokkos/lightning_kokkos.toml

@@ -18,6 +18,7 @@ observables = [
         "SProd",
         "Prod",
         "Exp",
+        "LinearCombination",
 ]
 
 # The union of all gate types listed in this section must match what

pennylane_lightning/lightning_qubit/lightning_qubit.py

@@ -34,7 +34,7 @@
     validate_observables,
 )
 from pennylane.measurements import MidMeasureMP
-from pennylane.operation import Tensor
+from pennylane.operation import DecompositionUndefinedError, Operator, Tensor
 from pennylane.ops import Prod, SProd, Sum
 from pennylane.tape import QuantumScript, QuantumTape
 from pennylane.transforms.core import TransformProgram
@@ -264,8 +264,6 @@
         "QFT",
         "ECR",
         "BlockEncode",
-        "MidMeasureMP",
-        "Conditional",
     }
 )
 # The set of supported operations.
@@ -292,17 +290,23 @@
 # The set of supported observables.
 
 
-def stopping_condition(op: qml.operation.Operator) -> bool:
+def stopping_condition(op: Operator) -> bool:
     """A function that determines whether or not an operation is supported by ``lightning.qubit``."""
     return op.name in _operations
 
 
-def accepted_observables(obs: qml.operation.Operator) -> bool:
+def stopping_condition_shots(op: Operator) -> bool:
+    """A function that determines whether or not an operation is supported by ``lightning.qubit``
+    with finite shots."""
+    return op.name in _operations or isinstance(op, (MidMeasureMP, qml.ops.op_math.Conditional))
+
+
+def accepted_observables(obs: Operator) -> bool:
     """A function that determines whether or not an observable is supported by ``lightning.qubit``."""
     return obs.name in _observables
 
 
-def adjoint_observables(obs: qml.operation.Operator) -> bool:
+def adjoint_observables(obs: Operator) -> bool:
     """A function that determines whether or not an observable is supported by ``lightning.qubit``
     when using the adjoint differentiation method."""
     if isinstance(obs, qml.Projector):
@@ -336,7 +340,7 @@
 
     try:
         prog((circuit,))
-    except (qml.operation.DecompositionUndefinedError, qml.DeviceError, AttributeError):
+    except (DecompositionUndefinedError, qml.DeviceError, AttributeError):
         return False
     return True
 
@@ -356,7 +360,11 @@
     name = "adjoint + lightning.qubit"
     program.add_transform(no_sampling, name=name)
     program.add_transform(
-        decompose, stopping_condition=adjoint_ops, name=name, skip_initial_state_prep=False
+        decompose,
+        stopping_condition=adjoint_ops,
+        stopping_condition_shots=stopping_condition_shots,
+        name=name,
+        skip_initial_state_prep=False,
     )
     program.add_transform(validate_observables, accepted_observables, name=name)
     program.add_transform(
@@ -408,7 +416,9 @@
     _CPP_BINARY_AVAILABLE = LQ_CPP_BINARY_AVAILABLE
     _new_API = True
     _backend_info = backend_info if LQ_CPP_BINARY_AVAILABLE else None
-    _config = Path(__file__).parent / "lightning_qubit.toml"
+
+    # This `config` is used in Catalyst-Frontend
+    config = Path(__file__).parent / "lightning_qubit.toml"
 
     # TODO: Move supported ops/obs to TOML file
     operations = _operations
@@ -463,7 +473,7 @@
             self._num_burnin = num_burnin
         else:
             self._kernel_name = None
-            self._num_burnin = None
+            self._num_burnin = 0
 
     @property
     def name(self):
@@ -515,19 +525,24 @@
         * Currently does not intrinsically support parameter broadcasting
 
         """
-        config = self._setup_execution_config(execution_config)
+        exec_config = self._setup_execution_config(execution_config)
         program = TransformProgram()
 
         program.add_transform(validate_measurements, name=self.name)
         program.add_transform(validate_observables, accepted_observables, name=self.name)
         program.add_transform(validate_device_wires, self.wires, name=self.name)
         program.add_transform(mid_circuit_measurements, device=self)
-        program.add_transform(decompose, stopping_condition=stopping_condition, name=self.name)
+        program.add_transform(
+            decompose,
+            stopping_condition=stopping_condition,
+            stopping_condition_shots=stopping_condition_shots,
+            name=self.name,
+        )
         program.add_transform(qml.transforms.broadcast_expand)
 
-        if config.gradient_method == "adjoint":
+        if exec_config.gradient_method == "adjoint":
             _add_adjoint_transforms(program)
-        return program, config
+        return program, exec_config
 
     # pylint: disable=unused-argument
     def execute(

pennylane_lightning/lightning_qubit/lightning_qubit.toml

@@ -19,6 +19,7 @@ observables = [
         "SProd",
         "Prod",
         "Exp",
+        "LinearCombination",
 ]
 
 # The union of all gate types listed in this section must match what

tests/new_api/test_device.py

@@ -33,6 +33,7 @@
     decompose,
     no_sampling,
     stopping_condition,
+    stopping_condition_shots,
     validate_adjoint_trainable_params,
     validate_device_wires,
     validate_measurements,
@@ -107,6 +108,7 @@ def test_add_adjoint_transforms(self):
         expected_program.add_transform(
             decompose,
             stopping_condition=adjoint_ops,
+            stopping_condition_shots=stopping_condition_shots,
             name=name,
             skip_initial_state_prep=False,
         )
@@ -184,7 +186,7 @@ def process_and_execute(device, tape):
         "batch_obs": False,
         "mcmc": False,
         "kernel_name": None,
-        "num_burnin": None,
+        "num_burnin": 0,
     }
 
     @pytest.mark.parametrize(
@@ -222,7 +224,7 @@ def process_and_execute(device, tape):
                         "batch_obs": False,
                         "mcmc": True,
                         "kernel_name": None,
-                        "num_burnin": None,
+                        "num_burnin": 0,
                     },
                 ),
             ),
@@ -260,7 +262,10 @@ def test_preprocess(self, adjoint):
             qml.devices.preprocess.mid_circuit_measurements, device=device
         )
         expected_program.add_transform(
-            decompose, stopping_condition=stopping_condition, name=device.name
+            decompose,
+            stopping_condition=stopping_condition,
+            stopping_condition_shots=stopping_condition_shots,
+            name=device.name,
         )
         expected_program.add_transform(qml.transforms.broadcast_expand)
 
@@ -270,6 +275,7 @@ def test_preprocess(self, adjoint):
             expected_program.add_transform(
                 decompose,
                 stopping_condition=adjoint_ops,
+                stopping_condition_shots=stopping_condition_shots,
                 name=name,
                 skip_initial_state_prep=False,
             )