Update QuantumScriptSerializer.serialize_observables to better handle new operator arithmetic

.github/CHANGELOG.md

@@ -24,6 +24,9 @@
   [(#607)](https://github.com/PennyLaneAI/pennylane-lightning/pull/607)
   [(#628)](https://github.com/PennyLaneAI/pennylane-lightning/pull/628)
 
+* Add Vector-Jacobian Product calculation support to `lightning.qubit`.
+  [(#644)](https://github.com/PennyLaneAI/pennylane-lightning/pull/644)
+
 * Add support for using new operator arithmetic as the default.
   [(#649)](https://github.com/PennyLaneAI/pennylane-lightning/pull/649)
 
@@ -43,6 +46,9 @@
 * Add `isort` to `requirements-dev.txt` and run before `black` upon `make format` to sort Python imports.
   [(#623)](https://github.com/PennyLaneAI/pennylane-lightning/pull/623)
 
+* Improve support for new operator arithmetic with `QuantumScriptSerializer.serialize_observables`.
+  [(#)]()
+
 ### Documentation
 
 ### Bug fixes

.github/workflows/tests_gpu_cuda.yml

@@ -119,9 +119,7 @@ jobs:
 
       - name: Install required packages
         run: |
-          # Omitting the installation of cmake v3.29.0 due to
-          # https://github.com/scikit-build/cmake-python-distributions/pull/474
-          python -m pip install ninja "cmake!=3.29.0" custatevec-cu${{ matrix.cuda_version }}
+          python -m pip install ninja cmake custatevec-cu${{ matrix.cuda_version }}
           sudo apt-get -y -q install liblapack-dev
 
       - name: Build and run unit tests
@@ -243,9 +241,7 @@ jobs:
         run: |
           cd main
           python -m pip install -r requirements-dev.txt
-          # Omitting the installation of cmake v3.29.0 due to
-          # https://github.com/scikit-build/cmake-python-distributions/pull/474
-          python -m pip install "cmake!=3.29.0" custatevec-cu${{ matrix.cuda_version }} openfermionpyscf
+          python -m pip install cmake custatevec-cu${{ matrix.cuda_version }} openfermionpyscf
 
       - name: Checkout PennyLane for release build
         if: inputs.pennylane-version == 'release'

.github/workflows/tests_linux_x86_mpi_gpu.yml

@@ -94,9 +94,7 @@ jobs:
       - name: Install required packages
         run: |
           python -m pip install -r requirements-dev.txt
-          # Omitting the installation of cmake v3.29.0 due to
-          # https://github.com/scikit-build/cmake-python-distributions/pull/474
-          python -m pip install "cmake!=3.29.0" custatevec-cu12
+          python -m pip install cmake custatevec-cu12
           sudo apt-get -y -q install liblapack-dev
 
       - name: Validate GPU version and installed compiler and modules

.github/workflows/wheel_macos_x86_64.yml

@@ -68,9 +68,7 @@ jobs:
         run: |
           mkdir -p ${{ github.workspace}}/Kokkos_install/${{ matrix.exec_model }}
           cd kokkos
-          # Omitting the installation of cmake v3.29.0 due to
-          # https://github.com/scikit-build/cmake-python-distributions/pull/474
-          python -m pip install "cmake!=3.29.0" ninja
+          python -m pip install cmake ninja
 
           cmake -BBuild . -DCMAKE_INSTALL_PREFIX=${{ github.workspace}}/Kokkos_install/${{ matrix.exec_model }} \
                           -DKokkos_ENABLE_COMPLEX_ALIGN=OFF \

.github/workflows/wheel_noarch.yml

@@ -44,9 +44,7 @@ jobs:
 
       - name: Install CMake and ninja
         run: |
-          # Omitting the installation of cmake v3.29.0 due to
-          # https://github.com/scikit-build/cmake-python-distributions/pull/474
-          python -m pip install --upgrade "cmake!=3.29.0" ninja
+          python -m pip install --upgrade cmake ninja
 
       - name: Build wheels
         if: ${{ matrix.pl_backend == 'lightning_qubit'}}

.github/workflows/wheel_win_x86_64.yml

@@ -64,9 +64,7 @@ jobs:
       - name: Install dependencies
         if: steps.kokkos-cache.outputs.cache-hit != 'true'
         run: |
-          # Omitting the installation of cmake v3.29.0 due to
-          # https://github.com/scikit-build/cmake-python-distributions/pull/474
-          python -m pip install "cmake!=3.29.0" build
+          python -m pip install cmake build
 
       - name: Build Kokkos core library
         if: steps.kokkos-cache.outputs.cache-hit != 'true'

pennylane_lightning/core/_serialize.py

@@ -21,7 +21,6 @@
     BasisState,
     DeviceError,
     Hadamard,
-    Hamiltonian,
     Identity,
     PauliX,
     PauliY,
@@ -34,7 +33,7 @@
 )
 from pennylane.math import unwrap
 from pennylane.operation import Tensor
-from pennylane.ops import Prod, SProd, Sum
+from pennylane.ops import Prod, SProd, Sum, Hamiltonian, LinearCombination
 from pennylane.tape import QuantumTape
 
 pauli_name_map = {
@@ -258,9 +257,8 @@ def _pauli_sentence(self, observable, wires_map: dict = None):
         terms = [self._pauli_word(pw, wires_map) for pw in pwords]
         coeffs = np.array(coeffs).astype(self.rtype)
 
-        # TODO: Add this
-        # if len(terms) == 1 and coeffs[0] == 1.0:
-        #     return terms[0]
+        if len(terms) == 1 and coeffs[0] == 1.0:
+            return terms[0]
 
         if self.split_obs:
             return [self.hamiltonian_obs([c], [t]) for (c, t) in zip(coeffs, terms)]
@@ -269,22 +267,22 @@ def _pauli_sentence(self, observable, wires_map: dict = None):
     # pylint: disable=protected-access, too-many-return-statements
     def _ob(self, observable, wires_map: dict = None):
         """Serialize a :class:`pennylane.operation.Observable` into an Observable."""
-        if isinstance(observable, (Prod, Sum, SProd)) and observable.pauli_rep is not None:
+        if isinstance(observable, (PauliX, PauliY, PauliZ, Identity, Hadamard)):
+            return self._named_obs(observable, wires_map)
+        if isinstance(observable, Hamiltonian):
+            return self._hamiltonian(observable, wires_map)
+        if observable.pauli_rep is not None:
             return self._pauli_sentence(observable.pauli_rep, wires_map)
         if isinstance(observable, Tensor) or (
             isinstance(observable, Prod) and not observable.has_overlapping_wires
         ):
             return self._tensor_ob(observable, wires_map)
-        if observable.name in ("Hamiltonian", "LinearCombination"):
+        if isinstance(observable, (LinearCombination, Prod, SProd, Sum)):
+            # if isinstance(observable, (Sum, Prod, SProd)):
+            #     observable = observable.simplify()
             return self._hamiltonian(observable, wires_map)
         if observable.name == "SparseHamiltonian":
             return self._sparse_hamiltonian(observable, wires_map)
-        if isinstance(observable, (PauliX, PauliY, PauliZ, Identity, Hadamard)):
-            return self._named_obs(observable, wires_map)
-        if observable.pauli_rep is not None:
-            return self._pauli_sentence(observable.pauli_rep, wires_map)
-        # if isinstance(observable, (Prod, Sum)):
-        #     return self._hamiltonian(observable, wires_map)
         return self._hermitian_ob(observable, wires_map)
 
     def serialize_observables(self, tape: QuantumTape, wires_map: dict = None) -> List:
@@ -313,7 +311,9 @@ def serialize_observables(self, tape: QuantumTape, wires_map: dict = None) -> Li
                 offset_indices.append(offset_indices[-1] + 1)
         return serialized_obs, offset_indices
 
-    def serialize_ops(self, tape: QuantumTape, wires_map: dict = None) -> Tuple[
+    def serialize_ops(
+        self, tape: QuantumTape, wires_map: dict = None
+    ) -> Tuple[
         List[List[str]],
         List[np.ndarray],
         List[List[int]],

pennylane_lightning/core/_version.py

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

pennylane_lightning/lightning_qubit/_adjoint_jacobian.py

@@ -293,7 +293,16 @@ def calculate_vjp(self, tape: QuantumTape, grad_vec):
                 "Adjoint differentiation does not support State measurements."
             )
 
+        if any(m.return_type is not Expectation for m in tape.measurements):
+            raise QuantumFunctionError(
+                "Adjoint differentiation method does not support expectation return type "
+                "mixed with other return types"
+            )
+
         # Proceed, because tape_return_type is Expectation.
+        if qml.math.ndim(grad_vec) == 0:
+            grad_vec = (grad_vec,)
+
         if len(grad_vec) != len(measurements):
             raise ValueError(
                 "Number of observables in the tape must be the same as the "

pennylane_lightning/lightning_qubit/lightning_qubit.py

@@ -15,8 +15,9 @@
 This module contains the LightningQubit class that inherits from the new device interface.
 """
 from dataclasses import replace
+from numbers import Number
 from pathlib import Path
-from typing import Callable, Optional, Sequence, Union
+from typing import Callable, Optional, Sequence, Tuple, Union
 
 import numpy as np
 import pennylane as qml
@@ -129,6 +130,55 @@ def simulate_and_jacobian(circuit: QuantumTape, state: LightningStateVector, bat
     return res, jac
 
 
+def vjp(
+    circuit: QuantumTape, cotangents: Tuple[Number], state: LightningStateVector, batch_obs=False
+):
+    """Compute the Vector-Jacobian Product (VJP) for a single quantum script.
+    Args:
+        circuit (QuantumTape): The single circuit to simulate
+        cotangents (Tuple[Number, Tuple[Number]]): Gradient-output vector. Must
+            have shape matching the output shape of the corresponding circuit. If
+            the circuit has a single output, ``cotangents`` may be a single number,
+            not an iterable of numbers.
+        state (LightningStateVector): handle to Lightning state vector
+        batch_obs (bool): Determine whether we process observables in parallel when
+            computing the VJP. This value is only relevant when the lightning
+            qubit is built with OpenMP.
+    Returns:
+        TensorLike: The VJP of the quantum script
+    """
+    circuit = circuit.map_to_standard_wires()
+    state.reset_state()
+    final_state = state.get_final_state(circuit)
+    return LightningAdjointJacobian(final_state, batch_obs=batch_obs).calculate_vjp(
+        circuit, cotangents
+    )
+
+
+def simulate_and_vjp(
+    circuit: QuantumTape, cotangents: Tuple[Number], state: LightningStateVector, batch_obs=False
+):
+    """Simulate a single quantum script and compute its Vector-Jacobian Product (VJP).
+    Args:
+        circuit (QuantumTape): The single circuit to simulate
+        cotangents (Tuple[Number, Tuple[Number]]): Gradient-output vector. Must
+            have shape matching the output shape of the corresponding circuit. If
+            the circuit has a single output, ``cotangents`` may be a single number,
+            not an iterable of numbers.
+        state (LightningStateVector): handle to Lightning state vector
+        batch_obs (bool): Determine whether we process observables in parallel when
+            computing the jacobian. This value is only relevant when the lightning
+            qubit is built with OpenMP.
+    Returns:
+        Tuple[TensorLike]: The results of the simulation and the calculated VJP
+    Note that this function can return measurements for non-commuting observables simultaneously.
+    """
+    circuit = circuit.map_to_standard_wires()
+    res = simulate(circuit, state)
+    _vjp = LightningAdjointJacobian(state, batch_obs=batch_obs).calculate_vjp(circuit, cotangents)
+    return res, _vjp
+
+
 _operations = frozenset(
     {
         "Identity",
@@ -569,3 +619,79 @@ def execute_and_compute_derivatives(
             simulate_and_jacobian(c, self._statevector, batch_obs=batch_obs) for c in circuits
         )
         return tuple(zip(*results))
+
+    def supports_vjp(
+        self,
+        execution_config: Optional[ExecutionConfig] = None,
+        circuit: Optional[QuantumTape] = None,
+    ) -> bool:
+        """Whether or not this device defines a custom vector jacobian product.
+        ``LightningQubit`` supports adjoint differentiation with analytic results.
+        Args:
+            execution_config (ExecutionConfig): The configuration of the desired derivative calculation
+            circuit (QuantumTape): An optional circuit to check derivatives support for.
+        Returns:
+            Bool: Whether or not a derivative can be calculated provided the given information
+        """
+        return self.supports_derivatives(execution_config, circuit)
+
+    def compute_vjp(
+        self,
+        circuits: QuantumTape_or_Batch,
+        cotangents: Tuple[Number],
+        execution_config: ExecutionConfig = DefaultExecutionConfig,
+    ):
+        r"""The vector jacobian product used in reverse-mode differentiation. ``LightningQubit`` uses the
+        adjoint differentiation method to compute the VJP.
+        Args:
+            circuits (Union[QuantumTape, Sequence[QuantumTape]]): the circuit or batch of circuits
+            cotangents (Tuple[Number, Tuple[Number]]): Gradient-output vector. Must have shape matching the output shape of the
+                corresponding circuit. If the circuit has a single output, ``cotangents`` may be a single number, not an iterable
+                of numbers.
+            execution_config (ExecutionConfig): a datastructure with all additional information required for execution
+        Returns:
+            tensor-like: A numeric result of computing the vector jacobian product
+        **Definition of vjp:**
+        If we have a function with jacobian:
+        .. math::
+            \vec{y} = f(\vec{x}) \qquad J_{i,j} = \frac{\partial y_i}{\partial x_j}
+        The vector jacobian product is the inner product of the derivatives of the output ``y`` with the
+        Jacobian matrix. The derivatives of the output vector are sometimes called the **cotangents**.
+        .. math::
+            \text{d}x_i = \Sigma_{i} \text{d}y_i J_{i,j}
+        **Shape of cotangents:**
+        The value provided to ``cotangents`` should match the output of :meth:`~.execute`. For computing the full Jacobian,
+        the cotangents can be batched to vectorize the computation. In this case, the cotangents can have the following
+        shapes. ``batch_size`` below refers to the number of entries in the Jacobian:
+        * For a state measurement, the cotangents must have shape ``(batch_size, 2 ** n_wires)``
+        * For ``n`` expectation values, the cotangents must have shape ``(n, batch_size)``. If ``n = 1``,
+          then the shape must be ``(batch_size,)``.
+        """
+        batch_obs = execution_config.device_options.get("batch_obs", self._batch_obs)
+        return tuple(
+            vjp(circuit, cots, self._statevector, batch_obs=batch_obs)
+            for circuit, cots in zip(circuits, cotangents)
+        )
+
+    def execute_and_compute_vjp(
+        self,
+        circuits: QuantumTape_or_Batch,
+        cotangents: Tuple[Number],
+        execution_config: ExecutionConfig = DefaultExecutionConfig,
+    ):
+        """Calculate both the results and the vector jacobian product used in reverse-mode differentiation.
+        Args:
+            circuits (Union[QuantumTape, Sequence[QuantumTape]]): the circuit or batch of circuits to be executed
+            cotangents (Tuple[Number, Tuple[Number]]): Gradient-output vector. Must have shape matching the output shape of the
+                corresponding circuit. If the circuit has a single output, ``cotangents`` may be a single number, not an iterable
+                of numbers.
+            execution_config (ExecutionConfig): a datastructure with all additional information required for execution
+        Returns:
+            Tuple, Tuple: the result of executing the scripts and the numeric result of computing the vector jacobian product
+        """
+        batch_obs = execution_config.device_options.get("batch_obs", self._batch_obs)
+        results = tuple(
+            simulate_and_vjp(circuit, cots, self._statevector, batch_obs=batch_obs)
+            for circuit, cots in zip(circuits, cotangents)
+        )
+        return tuple(zip(*results))

requirements-dev.txt

@@ -14,6 +14,6 @@ clang-tidy~=16.0
 clang-format~=16.0
 isort==5.13.2
 click==8.0.4
-cmake==3.28.4
+cmake
 custatevec-cu12
 pylint