Add the support class for the Adjoint Jacobian to the new device (#907)

### Before submitting

Please complete the following checklist when submitting a PR:

- [X] 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`.

- [ ] 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:**
 Migrate LightningGPU to the new device API

**Description of the Change:**
Create the `adjoint-jacobian` class for the new device API to achieve
the `jacobian` and `vjp` methods

**Benefits:**
Unlocking the `adjoint-jacobian` capabilities with the new device API
for LGPU

**Possible Drawbacks:**

**Related GitHub Issues:**
## **Freezzed PR** ⚠️ ❄️ 
To make a smooth integration of LightningGPU with the new device API, we
set the branch `gpuNewAPI_backend` as the base branch target for future
developments related to this big task.

The branch `gpuNewAPI_backend` has the mock of all classes and methods
necessary for the new API. Also, several tests were disabled with
``` python
if device_name == "lightning.gpu":
    pytest.skip("LGPU new API in WIP.  Skipping.",allow_module_level=True)
```
However, these tests will unblocked as the implementation progresses.

After all the developments for integrating LightningGPU with the new API
have been completed then the PR will be open to merge to `master`

[sc-70936] [sc-70939]

---------

Co-authored-by: Vincent Michaud-Rioux <vincentm@nanoacademic.com>
Co-authored-by: Ali Asadi <10773383+maliasadi@users.noreply.github.com>

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

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)