Repository: PennyLaneAI/pennylane
Pull request #4637 (2023-10-02 13:52:32)

Add `DeviceJacobianProducts` jacobian product calculator
This PR is a follow on to #4527 . [sc-44325]

It adds `DeviceJacobianProducts` that can bind device provided jvps and vjps to the execution pipeline.

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4637?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified lines are covered by tests :white_check_mark:
> Comparison is base [(`b42adc5`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/b42adc5de674290f2b0b3563fa07564aa198d9f2?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.63% compared to head [(`cddd699`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4637?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.63%.

<details><summary>Additional details and impacted files</summary>


```diff
@@             Coverage Diff              @@
##           jpc-devices    #4637   +/-   ##
============================================
  Coverage        99.63%   99.63%           
============================================
  Files              376      376           
  Lines            33693    33712   +19     
============================================
+ Hits             33569    33588   +19     
  Misses             124      124           
```


| [Files](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4637?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/devices/qubit/adjoint\_jacobian.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4637?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvcXViaXQvYWRqb2ludF9qYWNvYmlhbi5weQ==) | `100.00% <100.00%> (ø)` | |
| [pennylane/interfaces/jacobian\_products.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4637?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvamFjb2JpYW5fcHJvZHVjdHMucHk=) | `100.00% <100.00%> (ø)` | |


</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4637?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 30617b48e..e20ba24c7 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -140,9 +140,10 @@
   [(#4649)](https://github.com/PennyLaneAI/pennylane/pull/4649)
 
 * The `JacobianProductCalculator` abstract base class and implementations `TransformJacobianProducts`
-  and `DeviceDerivatives` have been added to `pennylane.interfaces.jacobian_products`.
+  `DeviceDerivatives`, and `DeviceJacobianProducts` have been added to `pennylane.interfaces.jacobian_products`.
   [(#4435)](https://github.com/PennyLaneAI/pennylane/pull/4435)
   [(#4527)](https://github.com/PennyLaneAI/pennylane/pull/4527)
+  [(#4637)](https://github.com/PennyLaneAI/pennylane/pull/4637)
 
 * Extended ``qml.qchem.import_state`` to import wavefunctions from MPS DMRG and SHCI classical
   calculations performed with the Block2 and Dice libraries, incorporating new tests and wavefunction
diff --git a/pennylane/devices/device_api.py b/pennylane/devices/device_api.py
index c8bfd53a8..9c07e6c4d 100644
--- a/pennylane/devices/device_api.py
+++ b/pennylane/devices/device_api.py
@@ -599,8 +599,9 @@ class Device(abc.ABC):
 
         Args:
             circuits (Union[QuantumTape, Sequence[QuantumTape]]): the circuit or batch of circuits
-            cotangents (Tuple[Number]): Gradient-output vector. Must have shape matching the output shape of the
-                corresponding circuit
+            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:
@@ -638,8 +639,9 @@ class Device(abc.ABC):
 
         Args:
             circuits (Union[QuantumTape, Sequence[QuantumTape]]): the circuit or batch of circuits to be executed
-            cotangents Tuple[Number]: Gradient-output vector. Must have shape matching the output shape of the
-                corresponding circuit
+            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:
diff --git a/pennylane/devices/qubit/adjoint_jacobian.py b/pennylane/devices/qubit/adjoint_jacobian.py
index 921599135..af7cb930a 100644
--- a/pennylane/devices/qubit/adjoint_jacobian.py
+++ b/pennylane/devices/qubit/adjoint_jacobian.py
@@ -218,6 +218,8 @@ def adjoint_vjp(tape: QuantumTape, cotangents: Tuple[Number], state=None):
 
     ket = state if state is not None else get_final_state(tape)[0]
 
+    if np.shape(cotangents) == tuple():
+        cotangents = (cotangents,)
     obs = qml.dot(cotangents, tape.observables)
     bra = apply_operation(obs, ket)
 
diff --git a/pennylane/interfaces/__init__.py b/pennylane/interfaces/__init__.py
index b0698fd83..ebb489002 100644
--- a/pennylane/interfaces/__init__.py
+++ b/pennylane/interfaces/__init__.py
@@ -49,6 +49,7 @@ Jacobian Product Calculation
     ~interfaces.jacobian_products.JacobianProductCalculator
     ~interfaces.jacobian_products.TransformJacobianProducts
     ~interfaces.jacobian_products.DeviceDerivatives
+    ~interfaces.jacobian_products.DeviceJacobianProducts
 
 """
 from .execution import cache_execute, execute, INTERFACE_MAP, SUPPORTED_INTERFACES
diff --git a/pennylane/interfaces/jacobian_products.py b/pennylane/interfaces/jacobian_products.py
index 8854dc218..f00e17828 100644
--- a/pennylane/interfaces/jacobian_products.py
+++ b/pennylane/interfaces/jacobian_products.py
@@ -547,3 +547,56 @@ class DeviceDerivatives(JacobianProductCalculator):
         jacs = self._dev_compute_derivatives(tapes)
         self._jacs_cache[tapes] = jacs
         return jacs
+
+
+class DeviceJacobianProducts(JacobianProductCalculator):
+    """Compute jacobian products using the native device methods.
+
+    Args:
+        device (pennylane.devices.Device): the device for execution and derivatives.
+            Must define both the vjp and jvp.
+        execution_config (pennylane.devices.ExecutionConfig): a datastructure containing the options needed to fully
+           describe the execution.
+
+    >>> dev = qml.device('default.qubit')
+    >>> config = qml.devices.ExecutionConfig(gradient_method="adjoint")
+    >>> jpc = DeviceJacobianProducts(dev, config)
+
+    This class relies on :meth:`~.devices.Device.compute_vjp` and :meth:`~.devices.Device.execute_and_compute_jvp`,
+    and works strictly for the newer device interface :class:`~.devices.Device`.  This contrasts :class:`~.DeviceDerivatives`
+    which works for both device interfaces and requests the full jacobian from the device.
+
+    """
+
+    def __repr__(self):
+        return f"<DeviceJacobianProducts: {self._device.name}, {self._execution_config}>"
+
+    def __init__(
+        self, device: qml.devices.Device, execution_config=qml.devices.DefaultExecutionConfig
+    ):
+        if logger.isEnabledFor(logging.DEBUG):  # pragma: no cover
+            logger.debug("DeviceJacobianProducts created with (%s, %s)", device, execution_config)
+        self._device = device
+        self._execution_config = execution_config
+
+    def execute_and_compute_jvp(
+        self, tapes: Batch, tangents: Tuple[Tuple[TensorLike]]
+    ) -> Tuple[ResultBatch, Tuple]:
+        if logger.isEnabledFor(logging.DEBUG):  # pragma: no cover
+            logger.debug("execute_and_compute_jvp called with (%s, %s)", tapes, tangents)
+        numpy_tapes = tuple(qml.transforms.convert_to_numpy_parameters(t) for t in tapes)
+        tangents = qml.math.unwrap(tangents)
+        return self._device.execute_and_compute_jvp(numpy_tapes, tangents, self._execution_config)
+
+    def compute_vjp(self, tapes: Batch, dy: Tuple[Tuple[TensorLike]]) -> Tuple:
+        if logger.isEnabledFor(logging.DEBUG):  # pragma: no cover
+            logger.debug("compute_vjp called with (%s, %s)", tapes, dy)
+        numpy_tapes = tuple(qml.transforms.convert_to_numpy_parameters(t) for t in tapes)
+        dy = qml.math.unwrap(dy)
+        return self._device.compute_vjp(numpy_tapes, dy, self._execution_config)
+
+    def compute_jacobian(self, tapes: Batch):
+        if logger.isEnabledFor(logging.DEBUG):  # pragma: no cover
+            logger.debug("compute_jacobian called with %s", tapes)
+        numpy_tapes = tuple(qml.transforms.convert_to_numpy_parameters(t) for t in tapes)
+        return self._device.compute_derivatives(numpy_tapes, self._execution_config)


Test patch:
diff --git a/tests/devices/qubit/test_adjoint_jacobian.py b/tests/devices/qubit/test_adjoint_jacobian.py
index dd5b2e234..eb411d219 100644
--- a/tests/devices/qubit/test_adjoint_jacobian.py
+++ b/tests/devices/qubit/test_adjoint_jacobian.py
@@ -399,7 +399,7 @@ class TestAdjointJVP:
 class TestAdjointVJP:
     """Test for adjoint_vjp"""
 
-    @pytest.mark.parametrize("cotangents", [(0,), (1.232,)])
+    @pytest.mark.parametrize("cotangents", [(0,), (1.232,), 5.2])
     def test_single_param_single_obs(self, cotangents, tol):
         """Test VJP is correct for a single parameter and observable"""
         x = np.array(0.654)
@@ -409,7 +409,8 @@ class TestAdjointVJP:
         actual = adjoint_vjp(qs, cotangents)
         assert isinstance(actual, np.ndarray)
 
-        expected = -cotangents[0] * np.sin(x)
+        iterable_cotangent = cotangents if isinstance(cotangents, tuple) else (cotangents,)
+        expected = -iterable_cotangent[0] * np.sin(x)
         assert np.allclose(actual, expected, atol=tol)
 
     @pytest.mark.parametrize("cotangents", [(0, 0), (0, 0.653), (1.232, 2.963)])
diff --git a/tests/interfaces/test_jacobian_products.py b/tests/interfaces/test_jacobian_products.py
index aae81d457..3a5d5fe9f 100644
--- a/tests/interfaces/test_jacobian_products.py
+++ b/tests/interfaces/test_jacobian_products.py
@@ -26,6 +26,7 @@ from pennylane.interfaces.jacobian_products import (
     JacobianProductCalculator,
     TransformJacobianProducts,
     DeviceDerivatives,
+    DeviceJacobianProducts,
 )
 
 dev = qml.device("default.qubit")
@@ -46,16 +47,26 @@ hadamard_grad_jpc = TransformJacobianProducts(
 device_jacs = DeviceDerivatives(dev, adjoint_config)
 legacy_device_jacs = DeviceDerivatives(dev_old, gradient_kwargs={"method": "adjoint_jacobian"})
 device_ps_jacs = DeviceDerivatives(dev_ps, ps_config)
+device_native_jps = DeviceJacobianProducts(dev, adjoint_config)
+device_ps_native_jps = DeviceJacobianProducts(dev_ps, ps_config)
 
 transform_jpc_matrix = [param_shift_jpc, hadamard_grad_jpc]
 dev_jpc_matrix = [device_jacs, legacy_device_jacs, device_ps_jacs]
-jpc_matrix = [param_shift_jpc, hadamard_grad_jpc, device_jacs, legacy_device_jacs, device_ps_jacs]
+jpc_matrix = [
+    param_shift_jpc,
+    hadamard_grad_jpc,
+    device_jacs,
+    legacy_device_jacs,
+    device_ps_jacs,
+    device_native_jps,
+    device_ps_native_jps,
+]
 
 
 def _accepts_finite_shots(jpc):
     if isinstance(jpc, TransformJacobianProducts):
         return True
-    if isinstance(jpc, DeviceDerivatives):
+    if isinstance(jpc, (DeviceDerivatives, DeviceJacobianProducts)):
         return isinstance(jpc._device, ParamShiftDerivativesDevice)
     return False
 
@@ -136,6 +147,23 @@ class TestBasics:
 
         assert repr(jpc) == expected
 
+    def test_device_jacobian_products_repr(self):
+        """Test the repr method for device jacobian products."""
+
+        device = qml.device("default.qubit")
+        config = qml.devices.ExecutionConfig(gradient_method="adjoint")
+
+        jpc = DeviceJacobianProducts(device, config)
+
+        expected = (
+            r"<DeviceJacobianProducts: default.qubit,"
+            r" ExecutionConfig(grad_on_execution=None, use_device_gradient=None, "
+            r"gradient_method='adjoint', gradient_keyword_arguments={}, device_options={},"
+            r" interface=None, derivative_order=1)>"
+        )
+
+        assert repr(jpc) == expected
+
 
 @pytest.mark.parametrize("jpc", jpc_matrix)
 @pytest.mark.parametrize("shots", (None, 10000, (10000, 10000)))
diff --git a/tests/param_shift_dev.py b/tests/param_shift_dev.py
index e1d95f6b2..1acbfd0d0 100644
--- a/tests/param_shift_dev.py
+++ b/tests/param_shift_dev.py
@@ -73,3 +73,83 @@ class ParamShiftDerivativesDevice(qml.devices.DefaultQubit):
         results = all_results[: len(circuits)]
         jacs = fn(all_results[len(circuits) :])
         return (results[0], jacs[0]) if is_single_circuit else (results, jacs)
+
+    def compute_jvp(self, circuits, tangents, execution_config=qml.devices.DefaultExecutionConfig):
+        is_single_circuit = False
+        if isinstance(circuits, qml.tape.QuantumScript):
+            is_single_circuit = True
+            circuits = (circuits,)
+            tangents = (tangents,)
+
+        if self.tracker.active:
+            self.tracker.update(jvp_batches=1, jvps=len(circuits))
+            self.tracker.record()
+
+        batch, fn = qml.gradients.batch_jvp(circuits, tangents, qml.gradients.param_shift)
+        results = self.execute(batch)
+        return fn(results)[0] if is_single_circuit else fn(results)
+
+    def execute_and_compute_jvp(
+        self, circuits, tangents, execution_config=qml.devices.DefaultExecutionConfig
+    ):
+        is_single_circuit = False
+        if isinstance(circuits, qml.tape.QuantumScript):
+            is_single_circuit = True
+            circuits = [circuits]
+            tangents = [tangents]
+
+        if self.tracker.active:
+            for c in circuits:
+                self.tracker.update(resources=c.specs["resources"])
+            self.tracker.update(
+                execute_and_jvp_batches=1, executions=len(circuits), jvps=len(circuits)
+            )
+            self.tracker.record()
+
+        batch, fn = qml.gradients.batch_jvp(circuits, tangents, qml.gradients.param_shift)
+        full_batch = tuple(circuits) + tuple(batch)
+        all_results = self.execute(full_batch)
+        results = all_results[: len(circuits)]
+        jvps = fn(all_results[len(circuits) :])
+        return (results[0], jvps[0]) if is_single_circuit else results, jvps
+
+    def compute_vjp(
+        self, circuits, cotangents, execution_config=qml.devices.DefaultExecutionConfig
+    ):
+        is_single_circuit = False
+        if isinstance(circuits, qml.tape.QuantumScript):
+            is_single_circuit = True
+            circuits = [circuits]
+            cotangents = [cotangents]
+
+        if self.tracker.active:
+            self.tracker.update(vjp_batches=1, vjps=len(circuits))
+            self.tracker.record()
+
+        batch, fn = qml.gradients.batch_vjp(circuits, cotangents, qml.gradients.param_shift)
+        results = self.execute(batch)
+        return fn(results)[0] if is_single_circuit else fn(results)
+
+    def execute_and_compute_vjp(
+        self, circuits, cotangents, execution_config=qml.devices.DefaultExecutionConfig
+    ):
+        is_single_circuit = False
+        if isinstance(circuits, qml.tape.QuantumScript):
+            is_single_circuit = True
+            circuits = [circuits]
+            cotangents = [cotangents]
+
+        if self.tracker.active:
+            for c in circuits:
+                self.tracker.update(resources=c.specs["resources"])
+            self.tracker.update(
+                execute_and_vjp_batches=1, executions=len(circuits), vjps=len(circuits)
+            )
+            self.tracker.record()
+
+        batch, fn = qml.gradients.batch_vjp(circuits, cotangents, qml.gradients.param_shift)
+        full_batch = tuple(circuits) + tuple(batch)
+        all_results = self.execute(full_batch)
+        results = all_results[: len(circuits)]
+        vjps = fn(all_results[len(circuits) :])
+        return (results[0], vjps[0]) if is_single_circuit else results, vjps