diff --git a/pennylane_lightning/core/_state_vector_base.py b/pennylane_lightning/core/_state_vector_base.py index bef158bc4..834f8e175 100644 --- a/pennylane_lightning/core/_state_vector_base.py +++ b/pennylane_lightning/core/_state_vector_base.py @@ -136,7 +136,6 @@ def _apply_basis_state(self, state, wires, use_async: Optional[bool] = None): raise ValueError("BasisState parameter and wires must be of equal length.") # Return a computational basis state over all wires. - print("FSX:", use_async) if use_async == None: self._qubit_state.setBasisState(list(state), list(wires)) else: diff --git a/pennylane_lightning/core/lightning_newAPI_base.py b/pennylane_lightning/core/lightning_newAPI_base.py index dcee73fd5..ae3879bcd 100644 --- a/pennylane_lightning/core/lightning_newAPI_base.py +++ b/pennylane_lightning/core/lightning_newAPI_base.py @@ -68,6 +68,7 @@ def __init__( # pylint: disable=too-many-arguments self._c_dtype = c_dtype self._batch_obs = batch_obs + self._sync = None if isinstance(wires, int): self._wire_map = None # should just use wires as is @@ -133,7 +134,7 @@ def jacobian( """ if wire_map is not None: [circuit], _ = qml.map_wires(circuit, wire_map) - state.reset_state() + state.reset_state(self._sync) final_state = state.get_final_state(circuit) return self.LightningAdjointJacobian(final_state, batch_obs=batch_obs).calculate_jacobian( circuit @@ -191,7 +192,7 @@ def vjp( # pylint: disable=too-many-arguments """ if wire_map is not None: [circuit], _ = qml.map_wires(circuit, wire_map) - state.reset_state() + state.reset_state(self._sync) final_state = state.get_final_state(circuit) return self.LightningAdjointJacobian(final_state, batch_obs=batch_obs).calculate_vjp( circuit, cotangents diff --git a/pennylane_lightning/lightning_gpu/_adjoint_jacobian.py b/pennylane_lightning/lightning_gpu/_adjoint_jacobian.py index 822af1f91..89d86515b 100644 --- a/pennylane_lightning/lightning_gpu/_adjoint_jacobian.py +++ b/pennylane_lightning/lightning_gpu/_adjoint_jacobian.py @@ -15,9 +15,37 @@ Internal methods for adjoint Jacobian differentiation method. """ +from warnings import warn + +try: + from pennylane_lightning.lightning_gpu_ops.algorithms import ( + AdjointJacobianC64, + AdjointJacobianC128, + create_ops_listC64, + create_ops_listC128, + ) + + try: + from pennylane_lightning.lightning_gpu_ops.algorithmsMPI import ( + AdjointJacobianMPIC64, + AdjointJacobianMPIC128, + create_ops_listMPIC64, + create_ops_listMPIC128, + ) + + MPI_SUPPORT = True + except ImportError as ex: + warn(str(ex), UserWarning) + MPI_SUPPORT = False + +except ImportError as ex: + warn(str(ex), UserWarning) + pass + import numpy as np -import pennylane as qml +from pennylane.tape import QuantumTape +# pylint: disable=ungrouped-imports from pennylane_lightning.core._adjoint_jacobian_base import LightningBaseAdjointJacobian from ._state_vector import LightningGPUStateVector @@ -31,5 +59,62 @@ class LightningGPUAdjointJacobian(LightningBaseAdjointJacobian): batch_obs(bool): If serialized tape is to be batched or not. """ - def __init__(self, lgpu_state: LightningGPUStateVector, batch_obs: bool = False) -> None: - super().__init__(lgpu_state, batch_obs) + # pylint: disable=too-few-public-methods + + def __init__(self, qubit_state: LightningGPUStateVector, batch_obs: bool = False) -> None: + super().__init__(qubit_state, batch_obs) + # Initialize the C++ binds + self._jacobian_lightning, self._create_ops_list_lightning = self._adjoint_jacobian_dtype() + + def _adjoint_jacobian_dtype(self): + """Binding to Lightning GPU Adjoint Jacobian C++ class. + + Returns: the AdjointJacobian class + """ + jacobian_lightning = ( + AdjointJacobianC64() if self.dtype == np.complex64 else AdjointJacobianC128() + ) + create_ops_list_lightning = ( + create_ops_listC64 if self.dtype == np.complex64 else create_ops_listC128 + ) + return jacobian_lightning, create_ops_list_lightning + + def calculate_jacobian(self, tape: QuantumTape): + """Computes the Jacobian with the adjoint method. + + .. code-block:: python + + statevector = LightningGPUStateVector(num_wires=num_wires) + statevector = statevector.get_final_state(tape) + jacobian = LightningGPUAdjointJacobian(statevector).calculate_jacobian(tape) + + Args: + tape (QuantumTape): Operations and measurements that represent instructions for execution on Lightning. + + Returns: + The Jacobian of a tape. + """ + + empty_array = self._handle_raises(tape, is_jacobian=True) + + if empty_array: + return np.array([], dtype=self.dtype) + + processed_data = self._process_jacobian_tape(tape) + + if not processed_data: # training_params is empty + return np.array([], dtype=self.dtype) + + trainable_params = processed_data["tp_shift"] + jac = self._jacobian_lightning( + processed_data["state_vector"], + processed_data["obs_serialized"], + processed_data["ops_serialized"], + trainable_params, + ) + jac = np.array(jac) + jac = jac.reshape(-1, len(trainable_params)) if len(jac) else jac + jac_r = np.zeros((jac.shape[0], processed_data["all_params"])) + jac_r[:, processed_data["record_tp_rows"]] = jac + + return self._adjoint_jacobian_processing(jac_r) diff --git a/pennylane_lightning/lightning_gpu/lightning_gpu.py b/pennylane_lightning/lightning_gpu/lightning_gpu.py index f43083af7..3be7390f6 100644 --- a/pennylane_lightning/lightning_gpu/lightning_gpu.py +++ b/pennylane_lightning/lightning_gpu/lightning_gpu.py @@ -55,10 +55,7 @@ LGPU_CPP_BINARY_AVAILABLE = True try: - # pylint: disable=no-name-in-module - from pennylane_lightning.lightning_gpu_ops import DevTag, MPIManager - - from ._mpi_handler import LightningGPU_MPIHandler + from ._mpi_handler import MPIHandler MPI_SUPPORT = True except ImportError as ex: @@ -318,6 +315,7 @@ def _set_lightning_classes(self): """Load the LightningStateVector, LightningMeasurements, LightningAdjointJacobian as class attribute""" self.LightningStateVector = LightningGPUStateVector self.LightningMeasurements = LightningGPUMeasurements + self.LightningAdjointJacobian = LightningGPUAdjointJacobian def _setup_execution_config(self, config): """ diff --git a/pennylane_lightning/lightning_kokkos/_state_vector.py b/pennylane_lightning/lightning_kokkos/_state_vector.py index 5e76249de..b629a17db 100644 --- a/pennylane_lightning/lightning_kokkos/_state_vector.py +++ b/pennylane_lightning/lightning_kokkos/_state_vector.py @@ -62,7 +62,6 @@ def __init__( num_wires, dtype=np.complex128, kokkos_args=None, - sync=True, ): # pylint: disable=too-many-arguments super().__init__(num_wires, dtype) @@ -70,7 +69,6 @@ def __init__( self._device_name = "lightning.kokkos" self._kokkos_config = {} - self._sync = sync # Initialize the state vector if kokkos_args is None: diff --git a/pennylane_lightning/lightning_kokkos/lightning_kokkos.py b/pennylane_lightning/lightning_kokkos/lightning_kokkos.py index 4089b3aca..6cf1d0461 100644 --- a/pennylane_lightning/lightning_kokkos/lightning_kokkos.py +++ b/pennylane_lightning/lightning_kokkos/lightning_kokkos.py @@ -302,7 +302,6 @@ def __init__( # pylint: disable=too-many-arguments shots=None, batch_obs=False, # Kokkos arguments - sync=True, kokkos_args=None, ): if not self._CPP_BINARY_AVAILABLE: @@ -324,11 +323,10 @@ def __init__( # pylint: disable=too-many-arguments # Kokkos specific options self._kokkos_args = kokkos_args - self._sync = sync # Creating the state vector self._statevector = self.LightningStateVector( - num_wires=len(self.wires), dtype=c_dtype, kokkos_args=kokkos_args, sync=sync + num_wires=len(self.wires), dtype=c_dtype, kokkos_args=kokkos_args ) if not LightningKokkos.kokkos_config: diff --git a/tests/conftest.py b/tests/conftest.py index ace8debfd..9c4e13c39 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -142,6 +142,9 @@ def get_device(): from pennylane_lightning.lightning_kokkos_ops import LightningException elif device_name == "lightning.gpu": from pennylane_lightning.lightning_gpu import LightningGPU as LightningDevice + from pennylane_lightning.lightning_gpu._adjoint_jacobian import ( + LightningGPUAdjointJacobian as LightningAdjointJacobian, + ) from pennylane_lightning.lightning_gpu._measurements import ( LightningGPUMeasurements as LightningMeasurements, ) @@ -149,8 +152,6 @@ def get_device(): LightningGPUStateVector as LightningStateVector, ) - LightningAdjointJacobian = None - if hasattr(pennylane_lightning, "lightning_gpu_ops"): import pennylane_lightning.lightning_gpu_ops as lightning_ops from pennylane_lightning.lightning_gpu_ops import LightningException diff --git a/tests/lightning_qubit/test_adjoint_jacobian_class.py b/tests/lightning_qubit/test_adjoint_jacobian_class.py index 199f2e4a8..d5b9355bf 100644 --- a/tests/lightning_qubit/test_adjoint_jacobian_class.py +++ b/tests/lightning_qubit/test_adjoint_jacobian_class.py @@ -34,9 +34,6 @@ allow_module_level=True, ) -if device_name == "lightning.gpu": - pytest.skip("LGPU new API in WIP. Skipping.", allow_module_level=True) - if device_name == "lightning.tensor": pytest.skip("Skipping tests for the LightningTensor class.", allow_module_level=True) @@ -423,7 +420,10 @@ def test_multiple_measurements(self, tol, lightning_sv): statevector = lightning_sv(num_wires=2) result_vjp = self.calculate_vjp(statevector, tape1, dy) - statevector.reset_state() + if device_name == "lightning.gpu": + statevector.reset_state(True) + else: + statevector.reset_state() result_jac = self.calculate_jacobian(statevector, tape2) @@ -483,7 +483,11 @@ def test_hermitian_expectation(self, tol, lightning_sv): qml.expval(qml.Hermitian(obs, wires=(0,))) tape.trainable_params = {0} - statevector.reset_state() + if device_name == "lightning.gpu": + statevector.reset_state(True) + else: + statevector.reset_state() + vjp = self.calculate_vjp(statevector, tape, dy) assert np.allclose(vjp, -0.8 * np.sin(x), atol=tol) @@ -500,7 +504,11 @@ def test_hermitian_tensor_expectation(self, tol, lightning_sv): qml.expval(qml.Hermitian(obs, wires=(0,)) @ qml.PauliZ(wires=1)) tape.trainable_params = {0} - statevector.reset_state() + if device_name == "lightning.gpu": + statevector.reset_state(True) + else: + statevector.reset_state() + vjp = self.calculate_vjp(statevector, tape, dy) assert np.allclose(vjp, -0.8 * np.sin(x), atol=tol) diff --git a/tests/lightning_qubit/test_jacobian_method.py b/tests/lightning_qubit/test_jacobian_method.py index d4439ca2b..745feee50 100644 --- a/tests/lightning_qubit/test_jacobian_method.py +++ b/tests/lightning_qubit/test_jacobian_method.py @@ -26,9 +26,6 @@ allow_module_level=True, ) -if device_name == "lightning.gpu": - pytest.skip("LGPU new API in WIP. Skipping.", allow_module_level=True) - if device_name == "lightning.tensor": pytest.skip("Skipping tests for the LightningTensor class.", allow_module_level=True)