Add the support class for the Adjoint Jacobian to the new device

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:

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

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 @@
         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)

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):
         """

pennylane_lightning/lightning_kokkos/_state_vector.py

@@ -62,15 +62,13 @@ def __init__(
         num_wires,
         dtype=np.complex128,
         kokkos_args=None,
-        sync=True,
     ):  # pylint: disable=too-many-arguments
 
         super().__init__(num_wires, dtype)
 
         self._device_name = "lightning.kokkos"
 
         self._kokkos_config = {}
-        self._sync = sync
 
         # Initialize the state vector
         if kokkos_args is None:

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:

tests/conftest.py

@@ -142,15 +142,16 @@ 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,
     )
     from pennylane_lightning.lightning_gpu._state_vector import (
         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

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)

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)