Add DeviceDerivatives jacobian product calculator (#4527)

This PR follows on from #4435 , and adds a child class that can handle
device-provided derivatives.

It has two private caches, `_results_cache` and `_jacs_cache`. This
handles storing the jacobian between the forward pass and the backward
pass. The caches will also store results and jacobians between
successive calls to the same instance.

The class can handle both device interfaces, `qml.Device` and
`qml.devices.Device`.

Name for the class is still up in the air. This class takes the full
jacobian from the device and then performs the dot product. The class in
#4637 takes the jacobian product itself from the device.

---------

Co-authored-by: David Wierichs <[email protected]>
Co-authored-by: Mudit Pandey <[email protected]>

doc/releases/changelog-dev.md

@@ -139,9 +139,10 @@
   [(#4628)](https://github.com/PennyLaneAI/pennylane/pull/4628)
   [(#4649)](https://github.com/PennyLaneAI/pennylane/pull/4649)
 
-* The `JacobianProductCalculator` abstract base class and implementation `TransformJacobianProducts`
-  have been added to `pennylane.interfaces.jacobian_products`.
+* The `JacobianProductCalculator` abstract base class and implementations `TransformJacobianProducts`
+  and `DeviceDerivatives` have been added to `pennylane.interfaces.jacobian_products`.
   [(#4435)](https://github.com/PennyLaneAI/pennylane/pull/4435)
+  [(#4527)](https://github.com/PennyLaneAI/pennylane/pull/4527)
 
 * 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

pennylane/_device.py

@@ -560,6 +560,9 @@ def execute_and_gradients(self, circuits, method="jacobian", **kwargs):
             tuple[list[array[float]], list[array[float]]]: Tuple containing list of measured value(s)
             and list of Jacobians. Returned Jacobians should be of shape ``(output_shape, num_params)``.
         """
+        if self.tracker.active:
+            self.tracker.update(execute_and_derivative_batches=1, derivatives=len(circuits))
+            self.tracker.record()
         gradient_method = getattr(self, method)
 
         res = []
@@ -593,6 +596,9 @@ def gradients(self, circuits, method="jacobian", **kwargs):
             list[array[float]]: List of Jacobians. Returned Jacobians should be of
             shape ``(output_shape, num_params)``.
         """
+        if self.tracker.active:
+            self.tracker.update(derivatives=len(circuits))
+            self.tracker.record()
         gradient_method = getattr(self, method)
         return [gradient_method(circuit, **kwargs) for circuit in circuits]
 

pennylane/gradients/hadamard_gradient.py

@@ -207,6 +207,10 @@ def hadamard_grad(
     assert_no_state_returns(tape.measurements, transform_name)
     assert_no_variance(tape.measurements, transform_name)
     assert_no_tape_batching(tape, transform_name)
+    if len(tape.measurements) > 1 and tape.shots.has_partitioned_shots:
+        raise NotImplementedError(
+            "hadamard gradient does not support multiple measurements with partitioned shots."
+        )
 
     if argnum is None and not tape.trainable_params:
         return _no_trainable_grad(tape)
@@ -323,12 +327,14 @@ def _expval_hadamard_grad(tape, argnum, aux_wire):
     multi_params = len(tape.trainable_params) > 1
 
     def processing_fn(results):  # pylint: disable=too-many-branches
-        final_res = [
-            [qml.math.convert_like(2 * coeff * r, r) for r in res]
-            if isinstance(res, tuple)
-            else qml.math.convert_like(2 * coeff * res, res)
-            for coeff, res in zip(coeffs, results)
-        ]
+        """Post processing function for computing a hadamard gradient."""
+        final_res = []
+        for coeff, res in zip(coeffs, results):
+            if isinstance(res, tuple):
+                new_val = [qml.math.convert_like(2 * coeff * r, r) for r in res]
+            else:
+                new_val = qml.math.convert_like(2 * coeff * res, res)
+            final_res.append(new_val)
 
         # Post process for probs
         if measurements_probs:

pennylane/interfaces/__init__.py

@@ -48,6 +48,7 @@
 
     ~interfaces.jacobian_products.JacobianProductCalculator
     ~interfaces.jacobian_products.TransformJacobianProducts
+    ~interfaces.jacobian_products.DeviceDerivatives
 
 """
 from .execution import cache_execute, execute, INTERFACE_MAP, SUPPORTED_INTERFACES