-
Notifications
You must be signed in to change notification settings - Fork 40
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Unify Lightning Kokkos device and Lightning Qubit device under a Ligh…
…tning Base device (#876) ### Before submitting Please complete the following checklist when submitting a PR: - [ ] All new features must include a unit test. If you've fixed a bug or added code that should be tested, add a test to the [`tests`](../tests) directory! - [X] All new functions and code must be clearly commented and documented. If you do make documentation changes, make sure that the docs build and render correctly by running `make docs`. - [X] Ensure that the test suite passes, by running `make test`. - [X] Add a new entry to the `.github/CHANGELOG.md` file, summarizing the change, and including a link back to the PR. - [X] Ensure that code is properly formatted by running `make format`. When all the above are checked, delete everything above the dashed line and fill in the pull request template. ------------------------------------------------------------------------------------------------------------ **Context:** Following the design of **Lightning Qubit* uses the New Device API end to end, we can now migrate lightning.kokkos device to the new device API. **Description of the Change:** Increase the reusability of the code by creating base classes for Lightning StateVector, Measurements, and AjointJacobian. Create a new device base class for LightningQubit and LightningKokkos. **Benefits:** Improve the maintenance and readability of the Python front-end devices for LQ and LK. **Possible Drawbacks:** **Related GitHub Issues:** [sc-71783] --------- Co-authored-by: ringo-but-quantum <[email protected]> Co-authored-by: Amintor Dusko <[email protected]> Co-authored-by: Vincent Michaud-Rioux <[email protected]> Co-authored-by: Ali Asadi <[email protected]>
- Loading branch information
1 parent
db98cbc
commit 89a1bed
Showing
18 changed files
with
1,550 additions
and
2,048 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,288 @@ | ||
# Copyright 2018-2024 Xanadu Quantum Technologies Inc. | ||
|
||
# Licensed under the Apache License, Version 2.0 (the "License"); | ||
# you may not use this file except in compliance with the License. | ||
# You may obtain a copy of the License at | ||
|
||
# http://www.apache.org/licenses/LICENSE-2.0 | ||
|
||
# Unless required by applicable law or agreed to in writing, software | ||
# distributed under the License is distributed on an "AS IS" BASIS, | ||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
# See the License for the specific language governing permissions and | ||
# limitations under the License. | ||
r""" | ||
Internal methods for adjoint Jacobian differentiation method. | ||
""" | ||
|
||
from abc import ABC, abstractmethod | ||
from typing import Any, Callable, List | ||
|
||
import numpy as np | ||
import pennylane as qml | ||
from pennylane import BasisState, QuantumFunctionError, StatePrep | ||
from pennylane.measurements import Expectation, MeasurementProcess, State | ||
from pennylane.operation import Operation | ||
from pennylane.tape import QuantumTape | ||
|
||
from pennylane_lightning.core._serialize import QuantumScriptSerializer | ||
|
||
|
||
class LightningBaseAdjointJacobian(ABC): | ||
"""Lightning [Device] Adjoint Jacobian class | ||
A class that serves as a base class for Lightning state-vector simulators. | ||
Check and execute the adjoint Jacobian differentiation method. | ||
Args: | ||
qubit_state(Lightning [Device] StateVector): State Vector to calculate the adjoint Jacobian with. | ||
batch_obs(bool): If serialized tape is to be batched or not. | ||
""" | ||
|
||
def __init__(self, qubit_state: Any, batch_obs: bool) -> None: | ||
self._qubit_state = qubit_state | ||
self._batch_obs = batch_obs | ||
|
||
# Dummy for the C++ bindings | ||
self._jacobian_lightning: Callable = None | ||
self._create_ops_list_lightning: Callable = None | ||
|
||
@property | ||
def qubit_state(self): | ||
"""Returns a handle to the Lightning [Device] StateVector object.""" | ||
return self._qubit_state | ||
|
||
@property | ||
def state(self): | ||
"""Returns a handle to the Lightning internal data object.""" | ||
return self._qubit_state.state_vector | ||
|
||
@property | ||
def dtype(self): | ||
"""Returns the simulation data type.""" | ||
return self._qubit_state.dtype | ||
|
||
@abstractmethod | ||
def _adjoint_jacobian_dtype(self): | ||
"""Binding to Lightning [Device] Adjoint Jacobian C++ class. | ||
Returns: the AdjointJacobian class | ||
""" | ||
|
||
@staticmethod | ||
def _get_return_type( | ||
measurements: List[MeasurementProcess], | ||
): | ||
"""Get the measurement return type. | ||
Args: | ||
measurements (List[MeasurementProcess]): a list of measurement processes to check. | ||
Returns: | ||
None, Expectation or State: a common return type of measurements. | ||
""" | ||
if not measurements: | ||
return None | ||
|
||
if len(measurements) == 1 and measurements[0].return_type is State: | ||
return State | ||
|
||
return Expectation | ||
|
||
def _process_jacobian_tape(self, tape: QuantumTape, split_obs: bool = False): | ||
"""Process a tape, serializing and building a dictionary proper for | ||
the adjoint Jacobian calculation in the C++ layer. | ||
Args: | ||
tape (QuantumTape): Operations and measurements that represent instructions for execution on Lightning. | ||
split_obs (bool, optional): If splitting the observables in a list. Defaults to False. | ||
Returns: | ||
dictionary: dictionary providing serialized data for Jacobian calculation. | ||
""" | ||
use_csingle = self._qubit_state.dtype == np.complex64 | ||
|
||
use_mpi = False | ||
obs_serialized, obs_indices = QuantumScriptSerializer( | ||
self._qubit_state.device_name, use_csingle, use_mpi, split_obs | ||
).serialize_observables(tape) | ||
|
||
ops_serialized, use_sp = QuantumScriptSerializer( | ||
self._qubit_state.device_name, use_csingle, use_mpi, split_obs | ||
).serialize_ops(tape) | ||
|
||
ops_serialized = self._create_ops_list_lightning(*ops_serialized) | ||
|
||
# We need to filter out indices in trainable_params which do not | ||
# correspond to operators. | ||
trainable_params = sorted(tape.trainable_params) | ||
if len(trainable_params) == 0: | ||
return None | ||
|
||
tp_shift = [] | ||
record_tp_rows = [] | ||
all_params = 0 | ||
|
||
for op_idx, trainable_param in enumerate(trainable_params): | ||
# get op_idx-th operator among differentiable operators | ||
operation, _, _ = tape.get_operation(op_idx) | ||
if isinstance(operation, Operation) and not isinstance( | ||
operation, (BasisState, StatePrep) | ||
): | ||
# We now just ignore non-op or state preps | ||
tp_shift.append(trainable_param) | ||
record_tp_rows.append(all_params) | ||
all_params += 1 | ||
|
||
if use_sp: | ||
# When the first element of the tape is state preparation. Still, I am not sure | ||
# whether there must be only one state preparation... | ||
tp_shift = [i - 1 for i in tp_shift] | ||
|
||
return { | ||
"state_vector": self.state, | ||
"obs_serialized": obs_serialized, | ||
"ops_serialized": ops_serialized, | ||
"tp_shift": tp_shift, | ||
"record_tp_rows": record_tp_rows, | ||
"all_params": all_params, | ||
"obs_indices": obs_indices, | ||
} | ||
|
||
@staticmethod | ||
def _adjoint_jacobian_processing(jac): | ||
""" | ||
Post-process the Jacobian matrix returned by ``adjoint_jacobian`` for | ||
the new return type system. | ||
""" | ||
jac = np.squeeze(jac) | ||
|
||
if jac.ndim == 0: | ||
return np.array(jac) | ||
|
||
if jac.ndim == 1: | ||
return tuple(np.array(j) for j in jac) | ||
|
||
# must be 2-dimensional | ||
return tuple(tuple(np.array(j_) for j_ in j) for j in jac) | ||
|
||
def _handle_raises(self, tape: QuantumTape, is_jacobian: bool, grad_vec=None): | ||
"""Handle the raises related with the tape for computing the Jacobian with the adjoint method or the vector-Jacobian products.""" | ||
|
||
if tape.shots: | ||
raise QuantumFunctionError( | ||
"Requested adjoint differentiation to be computed with finite shots. " | ||
"The derivative is always exact when using the adjoint " | ||
"differentiation method." | ||
) | ||
|
||
tape_return_type = self._get_return_type(tape.measurements) | ||
|
||
if is_jacobian: | ||
if not tape_return_type: | ||
# the tape does not have measurements | ||
return True | ||
|
||
if tape_return_type is State: | ||
raise QuantumFunctionError( | ||
"Adjoint differentiation method does not support measurement StateMP." | ||
) | ||
|
||
if not is_jacobian: | ||
if qml.math.allclose(grad_vec, 0.0) or not tape_return_type: | ||
# the tape does not have measurements or the gradient is 0.0 | ||
return True | ||
|
||
if tape_return_type is State: | ||
raise QuantumFunctionError( | ||
"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" | ||
) | ||
|
||
return False | ||
|
||
@abstractmethod | ||
def calculate_jacobian(self, tape: QuantumTape): | ||
"""Computes the Jacobian with the adjoint method. | ||
.. code-block:: python | ||
statevector = Lightning [Device] StateVector(num_wires=num_wires) | ||
statevector = statevector.get_final_state(tape) | ||
jacobian = Lightning [Device] AdjointJacobian(statevector).calculate_jacobian(tape) | ||
Args: | ||
tape (QuantumTape): Operations and measurements that represent instructions for execution on Lightning. | ||
Returns: | ||
The Jacobian of a tape. | ||
""" | ||
|
||
# pylint: disable=inconsistent-return-statements | ||
def calculate_vjp(self, tape: QuantumTape, grad_vec): | ||
"""Compute the vector-Jacobian products of a tape. | ||
.. code-block:: python | ||
statevector = Lightning [Device] StateVector(num_wires=num_wires) | ||
statevector = statevector.get_final_state(tape) | ||
vjp = Lightning [Device] AdjointJacobian(statevector).calculate_vjp(tape, grad_vec) | ||
computes :math:`\\pmb{w} = (w_1,\\cdots,w_m)` where | ||
.. math:: | ||
w_k = dy_k \\cdot J_{k,j} | ||
Here, :math:`dy` is the workflow cotangent (grad_vec), and :math:`J` the Jacobian. | ||
Args: | ||
tape (QuantumTape): Operations and measurements that represent instructions for execution on Lightning. | ||
grad_vec (tensor_like): Gradient-output vector, also called `dy` or cotangent. Must have shape matching the output | ||
shape of the corresponding tape, i.e. number of measurements if the return type is expectation. | ||
Returns: | ||
The vector-Jacobian products of a tape. | ||
""" | ||
|
||
empty_array = self._handle_raises(tape, is_jacobian=False, grad_vec=grad_vec) | ||
|
||
if empty_array: | ||
return qml.math.convert_like(np.zeros(len(tape.trainable_params)), grad_vec) | ||
|
||
# Proceed, because tape_return_type is Expectation. | ||
if qml.math.ndim(grad_vec) == 0: | ||
grad_vec = (grad_vec,) | ||
|
||
if len(grad_vec) != len(tape.measurements): | ||
raise ValueError( | ||
"Number of observables in the tape must be the same as the " | ||
"length of grad_vec in the vjp method" | ||
) | ||
|
||
if np.iscomplexobj(grad_vec): | ||
raise ValueError( | ||
"The vjp method only works with a real-valued grad_vec when the " | ||
"tape is returning an expectation value" | ||
) | ||
|
||
ham = qml.simplify(qml.dot(grad_vec, [m.obs for m in tape.measurements])) | ||
|
||
num_params = len(tape.trainable_params) | ||
|
||
if num_params == 0: | ||
return np.array([], dtype=self.qubit_state.dtype) | ||
|
||
new_tape = qml.tape.QuantumScript( | ||
tape.operations, | ||
[qml.expval(ham)], | ||
shots=tape.shots, | ||
trainable_params=tape.trainable_params, | ||
) | ||
|
||
return self.calculate_jacobian(new_tape) |
Oops, something went wrong.