Repository: PennyLaneAI/pennylane
Pull request #3374 (2022-11-16 20:33:04)

[Return-types #14] Add shot vector support to autograd interface
**Context:**
Part of the ongoing return types project.

**Description of the Change:**
Adds shot vector support to the autograd interface. Now the user can instantiate a device with a shot vector and compute outputs and gradients of a QNode:
```python
dev = qml.device("default.qubit", wires=1, shots=[1000, 2000, 3000])

@qml.qnode(dev, diff_method="parameter-shift")
def circuit(a):
    qml.RY(a, wires=0)
    qml.RX(0.2, wires=0)
    return qml.expval(qml.PauliZ(0))

def cost(a):
    return qml.math.stack(circuit(a))
```
```
>>> a = np.array(0.4)
>>> circuit(a)
tensor([0.9       , 0.907     , 0.89733333], requires_grad=True)
>>> qml.jacobian(circuit)(a)
array([-0.391     , -0.389     , -0.38433333])
```


github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3374?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#3374](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3374?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (aa599c5) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/3144cced0e7ff5698b92411eb4f366ea786c542a?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (3144cce) will **decrease** coverage by `0.00%`.
> The diff coverage is `100.00%`.

```diff
@@            Coverage Diff             @@
##           master    #3374      +/-   ##
==========================================
- Coverage   99.70%   99.69%   -0.01%     
==========================================
  Files         299      299              
  Lines       26146    26199      +53     
==========================================
+ Hits        26068    26119      +51     
- Misses         78       80       +2     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3374?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/gradients/parameter\_shift\_cv.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3374/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2dyYWRpZW50cy9wYXJhbWV0ZXJfc2hpZnRfY3YucHk=) | `99.59% <ø> (ø)` | |
| [pennylane/gradients/vjp.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3374/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2dyYWRpZW50cy92anAucHk=) | `98.63% <100.00%> (-1.37%)` | :arrow_down: |
| [pennylane/interfaces/autograd.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3374/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvYXV0b2dyYWQucHk=) | `100.00% <100.00%> (ø)` | |
| [pennylane/qchem/\_\_init\_\_.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3374/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3FjaGVtL19faW5pdF9fLnB5) | `100.00% <0.00%> (ø)` | |
| [pennylane/qchem/molecule.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3374/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3FjaGVtL21vbGVjdWxlLnB5) | `100.00% <0.00%> (ø)` | |
| [pennylane/qchem/structure.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3374/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3FjaGVtL3N0cnVjdHVyZS5weQ==) | `100.00% <0.00%> (ø)` | |
| [pennylane/interfaces/execution.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3374/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvZXhlY3V0aW9uLnB5) | `100.00% <0.00%> (ø)` | |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/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 9ebae6325..27fad6095 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -112,11 +112,69 @@
 
 * Remove private `_wires` setter from the `Controlled.map_wires` method.
   [3405](https://github.com/PennyLaneAI/pennylane/pull/3405)
-
+  
 * `QuantumTape._process_queue` has been moved to `qml.queuing.process_queue` to disentangle
   its functionality from the `QuantumTape` class.
   [(#3401)](https://github.com/PennyLaneAI/pennylane/pull/3401)
 
+<h4>Return types project</h4>
+
+* The autograd interface now supports devices with shot vectors.
+  [#3374](https://github.com/PennyLaneAI/pennylane/pull/3374)
+
+  Example with a single measurement:
+  ```python
+  dev = qml.device("default.qubit", wires=1, shots=[1000, 2000, 3000])
+
+  @qml.qnode(dev, diff_method="parameter-shift")
+  def circuit(a):
+      qml.RY(a, wires=0)
+      qml.RX(0.2, wires=0)
+      return qml.expval(qml.PauliZ(0))
+
+  def cost(a):
+      return qml.math.stack(circuit(a))
+  ```
+  ```
+  >>> qml.enable_return()
+  >>> a = np.array(0.4)
+  >>> circuit(a)
+  (array(0.902), array(0.922), array(0.896))
+  >>> cost(a)
+  array([0.9       , 0.907     , 0.89733333])
+  >>> qml.jacobian(cost)(a)
+  array([-0.391     , -0.389     , -0.38433333])
+  ```
+  Example with multiple measurements:
+  ```python
+  dev = qml.device("default.qubit", wires=2, shots=[1000, 2000, 3000])
+
+  @qml.qnode(dev, diff_method="parameter-shift")
+  def circuit(a):
+      qml.RY(a, wires=0)
+      qml.RX(0.2, wires=0)
+      qml.CNOT(wires=[0, 1])
+      return qml.expval(qml.PauliZ(0)), qml.probs([0, 1])
+
+  def cost(a):
+      res = circuit(a)
+      return qml.math.stack([qml.math.hstack(r) for r in res])
+  ```
+  ```
+  >>> circuit(a)
+  ((array(0.904), array([0.952, 0.   , 0.   , 0.048])),
+   (array(0.915), array([0.9575, 0.    , 0.    , 0.0425])),
+   (array(0.902), array([0.951, 0.   , 0.   , 0.049])))
+  >>> cost(a)
+  array([[0.91      , 0.955     , 0.        , 0.        , 0.045     ],
+         [0.895     , 0.9475    , 0.        , 0.        , 0.0525    ],
+         [0.90666667, 0.95333333, 0.        , 0.        , 0.04666667]])
+  >>> qml.jacobian(cost)(a)
+  array([[-0.37      , -0.185     ,  0.        ,  0.        ,  0.185     ],
+         [-0.409     , -0.2045    ,  0.        ,  0.        ,  0.2045    ],
+         [-0.37133333, -0.18566667,  0.        ,  0.        ,  0.18566667]])
+  ```
+
 <h3>Breaking changes</h3>
 
 * The `log_base` attribute has been moved from `MeasurementProcess` to the new `_VnEntropy` and
@@ -163,7 +221,7 @@ Deprecations cycles are tracked at [doc/developement/deprecations.rst](https://d
 
 * Original tape `_obs_sharing_wires` attribute is updated during its expansion.
   [#3293](https://github.com/PennyLaneAI/pennylane/pull/3293)
-  
+
 * Small fix of `MeasurementProcess.map_wires`, where both the `self.obs` and `self._wires`
   attributes were modified.
   [#3292](https://github.com/PennyLaneAI/pennylane/pull/3292)
diff --git a/pennylane/gradients/parameter_shift_cv.py b/pennylane/gradients/parameter_shift_cv.py
index e35760af4..0a94aad1c 100644
--- a/pennylane/gradients/parameter_shift_cv.py
+++ b/pennylane/gradients/parameter_shift_cv.py
@@ -483,6 +483,8 @@ def second_order_param_shift(tape, dev_wires, argnum=None, shifts=None, gradient
     return gradient_tapes, processing_fn
 
 
+# TODO: integration of CV devices with new return types
+# pylint: disable=unused-argument
 @gradient_transform
 def param_shift_cv(
     tape,
@@ -493,6 +495,7 @@ def param_shift_cv(
     fallback_fn=finite_diff,
     f0=None,
     force_order2=False,
+    shots=None,
 ):
     r"""Transform a continuous-variable QNode to compute the parameter-shift gradient of all gate
     parameters with respect to its inputs.
diff --git a/pennylane/gradients/vjp.py b/pennylane/gradients/vjp.py
index 5fe1d7fb6..d9b780869 100644
--- a/pennylane/gradients/vjp.py
+++ b/pennylane/gradients/vjp.py
@@ -16,6 +16,7 @@ This module contains functions for computing the vector-Jacobian product
 of tapes.
 """
 # pylint: disable=no-member, too-many-branches
+from collections.abc import Sequence
 import numpy as np
 import autograd
 
@@ -242,7 +243,7 @@ def compute_vjp(dy, jac, num=None):
     return qml.math.tensordot(jac, dy_row, [[0], [0]])
 
 
-def vjp(tape, dy, gradient_fn, gradient_kwargs=None):
+def vjp(tape, dy, gradient_fn, shots=None, gradient_kwargs=None):
     r"""Generate the gradient tapes and processing function required to compute
     the vector-Jacobian products of a tape.
 
@@ -345,7 +346,7 @@ def vjp(tape, dy, gradient_fn, gradient_kwargs=None):
         return [], lambda _, num=None: None
 
     try:
-        if qml.math.allclose(dy, 0):
+        if not isinstance(dy, tuple) and qml.math.allclose(dy, 0):
             # If the dy vector is zero, then the
             # corresponding element of the VJP will be zero,
             # and we can avoid a quantum computation.
@@ -363,22 +364,30 @@ def vjp(tape, dy, gradient_fn, gradient_kwargs=None):
     except (AttributeError, TypeError, NotImplementedError):
         pass
 
-    gradient_tapes, fn = gradient_fn(tape, **gradient_kwargs)
+    gradient_tapes, fn = gradient_fn(tape, shots=shots, **gradient_kwargs)
 
     def processing_fn(results, num=None):
         # postprocess results to compute the Jacobian
         jac = fn(results)
+        shot_vector = isinstance(shots, Sequence)
+
         if qml.active_return():
             multi = len(tape.measurements) > 1
-            if multi:
-                return compute_vjp_multi_new(dy, jac, num=num)
-            return compute_vjp_single_new(dy, jac, num=num)
+            comp_vjp_fn = compute_vjp_multi_new if multi else compute_vjp_single_new
+
+            if not shot_vector:
+                return comp_vjp_fn(dy, jac, num=num)
+
+            vjp_ = [comp_vjp_fn(dy_, jac_, num=num) for dy_, jac_ in zip(dy, jac)]
+            return qml.math.sum(qml.math.stack(vjp_), 0)
+
         return compute_vjp(dy, jac, num=num)
 
     return gradient_tapes, processing_fn
 
 
-def batch_vjp(tapes, dys, gradient_fn, reduction="append", gradient_kwargs=None):
+# pylint: disable=too-many-arguments
+def batch_vjp(tapes, dys, gradient_fn, shots=None, reduction="append", gradient_kwargs=None):
     r"""Generate the gradient tapes and processing function required to compute
     the vector-Jacobian products of a batch of tapes.
 
@@ -499,7 +508,7 @@ def batch_vjp(tapes, dys, gradient_fn, reduction="append", gradient_kwargs=None)
 
     # Loop through the tapes and dys vector
     for tape, dy in zip(tapes, dys):
-        g_tapes, fn = vjp(tape, dy, gradient_fn, gradient_kwargs)
+        g_tapes, fn = vjp(tape, dy, gradient_fn, shots=shots, gradient_kwargs=gradient_kwargs)
         reshape_info.append(len(g_tapes))
         processing_fns.append(fn)
         gradient_tapes.extend(g_tapes)
diff --git a/pennylane/interfaces/autograd.py b/pennylane/interfaces/autograd.py
index 9d6c18d3d..11e4c3782 100644
--- a/pennylane/interfaces/autograd.py
+++ b/pennylane/interfaces/autograd.py
@@ -16,6 +16,8 @@ This module contains functions for adding the Autograd interface
 to a PennyLane Device class.
 """
 # pylint: disable=too-many-arguments
+from collections.abc import Sequence
+
 import autograd
 from autograd.numpy.numpy_boxes import ArrayBox
 
@@ -424,7 +426,7 @@ def _vjp_new(
 
         jacs = []
         for t in tapes:
-            g_tapes, fn = gradient_fn(t, **gradient_kwargs)
+            g_tapes, fn = gradient_fn(t, shots=device.shot_vector, **gradient_kwargs)
 
             with qml.tape.Unwrap(*g_tapes):
                 res, _ = execute_fn(g_tapes, **gradient_kwargs)
@@ -450,7 +452,7 @@ def _vjp_new(
         if jacs:
             # Jacobians were computed on the forward pass (mode="forward") or the Jacobian was cached
             # No additional quantum evaluations needed; simply compute the VJPs directly.
-            vjps = _compute_vjps_autograd(jacs, dy, multi_measurements)
+            vjps = _compute_vjps_autograd(jacs, dy, multi_measurements, device.shot_vector)
 
         else:
             # Need to compute the Jacobians on the backward pass (accumulation="backward")
@@ -464,6 +466,7 @@ def _vjp_new(
                             tapes,
                             dy,
                             gradient_fn,
+                            shots=device.shot_vector,
                             reduction="append",
                             gradient_kwargs=gradient_kwargs,
                         )
@@ -472,7 +475,12 @@ def _vjp_new(
 
                 else:
                     vjp_tapes, processing_fn = qml.gradients.batch_vjp(
-                        tapes, dy, gradient_fn, reduction="append", gradient_kwargs=gradient_kwargs
+                        tapes,
+                        dy,
+                        gradient_fn,
+                        shots=device.shot_vector,
+                        reduction="append",
+                        gradient_kwargs=gradient_kwargs,
                     )
 
                     # This is where the magic happens. Note that we call ``execute``.
@@ -502,7 +510,7 @@ def _vjp_new(
                 with qml.tape.Unwrap(*tapes):
                     jacs = gradient_fn(tapes, **gradient_kwargs)
 
-                vjps = _compute_vjps_autograd(jacs, dy, multi_measurements)
+                vjps = _compute_vjps_autograd(jacs, dy, multi_measurements, device.shot_vector)
 
         return_vjps = [
             qml.math.to_numpy(v, max_depth=_n) if isinstance(v, ArrayBox) else v for v in vjps
@@ -511,19 +519,29 @@ def _vjp_new(
             # TODO: remove this exceptional case once the source of this issue
             # https://github.com/PennyLaneAI/pennylane-sf/issues/89 is determined
             return (return_vjps,)  # pragma: no cover
+
         return return_vjps
 
     return grad_fn
 
 
-def _compute_vjps_autograd(jacs, dy, multi_measurements):
+def _compute_vjps_autograd(jacs, dy, multi_measurements, shots):
     """Compute the vjps of multiple tapes, directly for a Jacobian and co-tangents dys."""
     vjps = []
     for i, multi in enumerate(multi_measurements):
-        if multi:
-            vjps.append(qml.gradients.compute_vjp_multi_new(dy[i], jacs[i]))
-        else:
-            vjps.append(qml.gradients.compute_vjp_single_new(dy[i], jacs[i]))
+        shot_vector_defined = isinstance(shots, Sequence)
+        dy_ = dy[i] if shot_vector_defined else (dy[i],)
+        jac_ = jacs[i] if shot_vector_defined else (jacs[i],)
+
+        shot_vjps = []
+        for d, j in zip(dy_, jac_):
+            if multi:
+                shot_vjps.append(qml.gradients.compute_vjp_multi_new(d, j))
+            else:
+                shot_vjps.append(qml.gradients.compute_vjp_single_new(d, j))
+
+        vjps.append(qml.math.sum(qml.math.stack(shot_vjps), axis=0))
+
     return vjps
 
 
diff --git a/pennylane/return_types.py b/pennylane/return_types.py
index 161615cb5..aaa1980a3 100644
--- a/pennylane/return_types.py
+++ b/pennylane/return_types.py
@@ -277,6 +277,131 @@ def enable_return():
                        [[ 4.9707499e-01, -2.9999996e-04],
                         [ 6.2500127e-04, -1.2500001e-04]]], dtype=float32)))
 
+    .. details::
+        :title: Usage Details for Autograd
+
+        Autograd only allows differentiating functions that have array or tensor outputs. QNodes that have
+        multiple measurements may output tuples and cause errors with Autograd. Similarly, the outputs of
+        ``qml.grad`` or ``qml.jacobian`` may be tuples if there are multiple trainable arguments and
+        cause errors with higher-order derivatives.
+
+        These issues can be overcome by stacking the QNode results or gradients before computing derivatives.
+
+        **Examples for shot vectors**
+
+        Example for first-order derivatives with a single return:
+
+        .. code-block:: python
+
+            def func(a, b):
+                qml.RY(a, wires=0)
+                qml.RX(b, wires=1)
+                qml.CNOT(wires=[0, 1])
+
+            dev = qml.device("default.qubit", wires=2, shots=[100, 200, 300])
+
+            @qml.qnode(dev, diff_method="parameter-shift")
+            def circuit(a, b):
+                func(a, b)
+                return qml.expval(qml.PauliZ(0))
+
+            def cost(a, b):
+                res = circuit(a, b)
+                return qml.math.stack(res)
+
+        >>> a = np.array(0.4, requires_grad=True)
+        >>> b = np.array(0.6, requires_grad=True)
+        >>> cost(a, b)
+        [0.94       0.93       0.89333333]
+        >>> qml.jacobian(cost)(a, b)
+        (array([-0.44, -0.38, -0.38]), array([-0.01      , -0.01      , -0.00666667]))
+
+        Example for first-order derivatives with multiple returns:
+
+        .. code-block:: python
+
+            @qml.qnode(dev, diff_method="parameter-shift")
+            def circuit(a, b):
+                func(a, b)
+                return qml.expval(qml.PauliZ(0)), qml.probs([0, 1])
+
+            def cost(a, b):
+                res = circuit(a, b)
+                return qml.math.stack([qml.math.hstack(r) for r in res])
+
+        >>> cost(a, b)
+        [[0.92       0.91       0.05       0.03       0.01      ]
+         [0.95       0.9        0.075      0.         0.025     ]
+         [0.95333333 0.89666667 0.08       0.         0.02333333]]
+        >>> qml.jacobian(cost)(a, b)
+        (array([[-0.41  , -0.16  , -0.045 ,  0.01  ,  0.195 ],
+                [-0.325 , -0.1325, -0.03  ,  0.01  ,  0.1525],
+                [-0.39  , -0.175 , -0.02  ,  0.02  ,  0.175 ]]),
+         array([[-0.01      , -0.26      ,  0.255     ,  0.02      , -0.015     ],
+                [ 0.01      , -0.2525    ,  0.2575    ,  0.0075    , -0.0125    ],
+                [-0.02      , -0.26666667,  0.25666667,  0.01666667, -0.00666667]]))
+
+        Example for second-order derivatives with a single return:
+
+        .. code-block:: python
+
+            @qml.qnode(dev, diff_method="parameter-shift", max_diff=2)
+            def circuit(a, b):
+                func(a, b)
+                return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
+
+            def cost(a, b):
+                def cost2(a, b):
+                    res = circuit(a, b)
+                    return qml.math.stack(res)
+
+                return qml.math.stack(qml.jacobian(cost2)(a, b))
+
+        >>> cost(a, b)
+        [[ 0.05       -0.01        0.00666667]
+         [-0.51       -0.48       -0.62      ]]
+        >>> qml.jacobian(cost)(a, b)
+        (array([[-0.03      , -0.02      , -0.00333333],
+                [ 0.01      ,  0.025     , -0.02      ]]),
+         array([[ 0.005     ,  0.03      , -0.02166667],
+                [-0.85      , -0.83      , -0.81333333]]))
+
+        Example for second-order derivatives with multiple returns:
+
+        .. code-block:: python
+
+            @qml.qnode(dev, diff_method="parameter-shift", max_diff=2)
+            def circuit(a, b):
+                func(a, b)
+                return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1)), qml.probs([0, 1])
+
+            def cost(a, b):
+                def cost2(a, b):
+                    res = circuit(a, b)
+                    return qml.math.stack([qml.math.hstack(r) for r in res])
+
+                return qml.math.stack(qml.jacobian(cost2)(a, b))
+
+        >>> cost(a, b)
+        [[[-0.05       -0.2         0.          0.025       0.175     ]
+          [ 0.02       -0.17       -0.0275      0.0175      0.18      ]
+          [-0.00666667 -0.19166667 -0.01333333  0.01666667  0.18833333]]
+         [[-0.52       -0.25        0.265      -0.005      -0.01      ]
+          [-0.54       -0.25        0.2625      0.0075     -0.02      ]
+          [-0.57666667 -0.275       0.28166667  0.00666667 -0.01333333]]]
+        >>> qml.jacobian(cost)(a, b)
+        (array([[[ 0.02      , -0.38      , -0.055     ,  0.045     ,  0.39      ],
+                 [ 0.03      , -0.4275    , -0.04      ,  0.025     ,  0.4425    ],
+                 [-0.00666667, -0.42333333, -0.025     ,  0.02833333,  0.42      ]],
+                [[ 0.        ,  0.0625    , -0.065     ,  0.065     , -0.0625    ],
+                 [ 0.015     ,  0.06875   , -0.05      ,  0.0425    , -0.06125   ],
+                 [-0.03      ,  0.03833333, -0.04583333,  0.06083333, -0.05333333]]]),
+         array([[[-0.06      ,  0.0325    , -0.0525    ,  0.0825    , -0.0625    ],
+                 [ 0.02      ,  0.03      , -0.035     ,  0.025     , -0.02      ],
+                 [ 0.00833333,  0.06166667, -0.05666667,  0.0525    , -0.0575    ]],
+                [[-0.85      , -0.415     ,  0.425     ,  0.        , -0.01      ],
+                 [-0.78      , -0.3775    ,  0.375     ,  0.015     , -0.0125    ],
+                 [-0.80666667, -0.38      ,  0.38666667,  0.01666667, -0.02333333]]]))
     """
 
     global __activated


Test patch:
diff --git a/tests/returntypes/test_autograd_qnode_shot_vector.py b/tests/returntypes/test_autograd_qnode_shot_vector.py
new file mode 100644
index 000000000..8a7d13739
--- /dev/null
+++ b/tests/returntypes/test_autograd_qnode_shot_vector.py
@@ -0,0 +1,687 @@
+# Copyright 2022 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Integration tests for using the Autograd interface with shot vectors and with a QNode"""
+
+import pytest
+
+import pennylane as qml
+from pennylane import numpy as np
+from pennylane import qnode
+
+pytestmark = pytest.mark.autograd
+
+shots = [((5, 2), 1, 10), (1, 10, (5, 2))]
+
+qubit_device_and_diff_method = [
+    ["default.qubit", "finite-diff", {"h": 10e-2}],
+    ["default.qubit", "parameter-shift", {}],
+]
+
+
+@pytest.mark.parametrize("shots", shots)
+@pytest.mark.parametrize("dev_name,diff_method,gradient_kwargs", qubit_device_and_diff_method)
+class TestReturnWithShotVectors:
+    """Class to test the shape of the Grad/Jacobian/Hessian with different return types and shot vectors."""
+
+    def test_jac_single_measurement_param(self, dev_name, diff_method, gradient_kwargs, shots):
+        """For one measurement and one param, the gradient is a float."""
+        dev = qml.device(dev_name, wires=1, shots=shots)
+
+        @qnode(dev, diff_method=diff_method, **gradient_kwargs)
+        def circuit(a):
+            qml.RY(a, wires=0)
+            qml.RX(0.2, wires=0)
+            return qml.expval(qml.PauliZ(0))
+
+        a = np.array(0.1)
+
+        def cost(a):
+            return qml.math.stack(circuit(a))
+
+        jac = qml.jacobian(cost)(a)
+
+        assert isinstance(jac, np.ndarray)
+        num_copies = sum(
+            [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)]
+        )
+        assert jac.shape == (num_copies,)
+
+    def test_jac_single_measurement_multiple_param(
+        self, dev_name, diff_method, gradient_kwargs, shots
+    ):
+        """For one measurement and multiple param, the gradient is a tuple of arrays."""
+        dev = qml.device(dev_name, wires=1, shots=shots)
+
+        @qnode(dev, diff_method=diff_method, **gradient_kwargs)
+        def circuit(a, b):
+            qml.RY(a, wires=0)
+            qml.RX(b, wires=0)
+            return qml.expval(qml.PauliZ(0))
+
+        a = np.array(0.1)
+        b = np.array(0.2)
+
+        def cost(a, b):
+            return qml.math.stack(circuit(a, b))
+
+        jac = qml.jacobian(cost, argnum=[0, 1])(a, b)
+
+        assert isinstance(jac, tuple)
+        num_copies = sum(
+            [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)]
+        )
+        assert len(jac) == 2
+        for j in jac:
+            assert isinstance(j, np.ndarray)
+            assert j.shape == (num_copies,)
+
+    def test_jacobian_single_measurement_multiple_param_array(
+        self, dev_name, diff_method, gradient_kwargs, shots
+    ):
+        """For one measurement and multiple param as a single array params, the gradient is an array."""
+        dev = qml.device(dev_name, wires=1, shots=shots)
+
+        @qnode(dev, diff_method=diff_method, **gradient_kwargs)
+        def circuit(a):
+            qml.RY(a[0], wires=0)
+            qml.RX(a[1], wires=0)
+            return qml.expval(qml.PauliZ(0))
+
+        a = np.array([0.1, 0.2])
+
+        def cost(a):
+            return qml.math.stack(circuit(a))
+
+        jac = qml.jacobian(cost)(a)
+
+        assert isinstance(jac, np.ndarray)
+        num_copies = sum(
+            [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)]
+        )
+        assert jac.shape == (num_copies, 2)
+
+    def test_jacobian_single_measurement_param_probs(
+        self, dev_name, diff_method, gradient_kwargs, shots
+    ):
+        """For a multi dimensional measurement (probs), check that a single array is returned with the correct
+        dimension"""
+        dev = qml.device(dev_name, wires=2, shots=shots)
+
+        @qnode(dev, diff_method=diff_method, **gradient_kwargs)
+        def circuit(a):
+            qml.RY(a, wires=0)
+            qml.RX(0.2, wires=0)
+            return qml.probs(wires=[0, 1])
+
+        a = np.array(0.1)
+
+        def cost(a):
+            return qml.math.stack(circuit(a))
+
+        jac = qml.jacobian(cost)(a)
+
+        assert isinstance(jac, np.ndarray)
+        num_copies = sum(
+            [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)]
+        )
+        assert jac.shape == (num_copies, 4)
+
+    def test_jacobian_single_measurement_probs_multiple_param(
+        self, dev_name, diff_method, gradient_kwargs, shots
+    ):
+        """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with
+        the correct dimension"""
+        dev = qml.device(dev_name, wires=2, shots=shots)
+
+        @qnode(dev, diff_method=diff_method, **gradient_kwargs)
+        def circuit(a, b):
+            qml.RY(a, wires=0)
+            qml.RX(b, wires=0)
+            return qml.probs(wires=[0, 1])
+
+        a = np.array(0.1)
+        b = np.array(0.2)
+
+        def cost(a, b):
+            return qml.math.stack(circuit(a, b))
+
+        jac = qml.jacobian(cost, argnum=[0, 1])(a, b)
+
+        assert isinstance(jac, tuple)
+        num_copies = sum(
+            [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)]
+        )
+        assert len(jac) == 2
+        for j in jac:
+            assert isinstance(j, np.ndarray)
+            assert j.shape == (num_copies, 4)
+
+    def test_jacobian_single_measurement_probs_multiple_param_single_array(
+        self, dev_name, diff_method, gradient_kwargs, shots
+    ):
+        """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with
+        the correct dimension"""
+        dev = qml.device(dev_name, wires=2, shots=shots)
+
+        @qnode(dev, diff_method=diff_method, **gradient_kwargs)
+        def circuit(a):
+            qml.RY(a[0], wires=0)
+            qml.RX(a[1], wires=0)
+            return qml.probs(wires=[0, 1])
+
+        a = np.array([0.1, 0.2])
+
+        def cost(a):
+            return qml.math.stack(circuit(a))
+
+        jac = qml.jacobian(cost)(a)
+
+        assert isinstance(jac, np.ndarray)
+        num_copies = sum(
+            [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)]
+        )
+        assert jac.shape == (num_copies, 4, 2)
+
+    def test_jacobian_expval_expval_multiple_params(
+        self, dev_name, diff_method, gradient_kwargs, shots
+    ):
+        """The gradient of multiple measurements with multiple params return a tuple of arrays."""
+        dev = qml.device(dev_name, wires=2, shots=shots)
+
+        par_0 = np.array(0.1)
+        par_1 = np.array(0.2)
+
+        @qnode(dev, diff_method=diff_method, max_diff=1, **gradient_kwargs)
+        def circuit(x, y):
+            qml.RX(x, wires=[0])
+            qml.RY(y, wires=[1])
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.expval(qml.PauliZ(0))
+
+        def cost(x, y):
+            res = circuit(x, y)
+            return qml.math.stack([qml.math.stack(r) for r in res])
+
+        jac = qml.jacobian(cost, argnum=[0, 1])(par_0, par_1)
+
+        assert isinstance(jac, tuple)
+        num_copies = sum(
+            [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)]
+        )
+        assert len(jac) == 2
+        for j in jac:
+            assert isinstance(j, np.ndarray)
+            assert j.shape == (num_copies, 2)
+
+    def test_jacobian_expval_expval_multiple_params_array(
+        self, dev_name, diff_method, gradient_kwargs, shots
+    ):
+        """The jacobian of multiple measurements with a multiple params array return a single array."""
+        dev = qml.device(dev_name, wires=2, shots=shots)
+
+        @qnode(dev, diff_method=diff_method, **gradient_kwargs)
+        def circuit(a):
+            qml.RY(a[0], wires=0)
+            qml.RX(a[1], wires=0)
+            qml.RY(a[2], wires=0)
+            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.expval(qml.PauliZ(0))
+
+        a = np.array([0.1, 0.2, 0.3])
+
+        def cost(a):
+            res = circuit(a)
+            return qml.math.stack([qml.math.stack(r) for r in res])
+
+        jac = qml.jacobian(cost)(a)
+
+        assert isinstance(jac, np.ndarray)
+        num_copies = sum(
+            [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)]
+        )
+        assert jac.shape == (num_copies, 2, 3)
+
+    def test_jacobian_var_var_multiple_params(self, dev_name, diff_method, gradient_kwargs, shots):
+        """The jacobian of multiple measurements with multiple params return a tuple of arrays."""
+        dev = qml.device(dev_name, wires=2, shots=shots)
+
+        par_0 = np.array(0.1)
+        par_1 = np.array(0.2)
+
+        @qnode(dev, diff_method=diff_method, **gradient_kwargs)
+        def circuit(x, y):
+            qml.RX(x, wires=[0])
+            qml.RY(y, wires=[1])
+            qml.CNOT(wires=[0, 1])
+            return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.var(qml.PauliZ(0))
+
+        def cost(x, y):
+            res = circuit(x, y)
+            return qml.math.stack([qml.math.stack(r) for r in res])
+
+        jac = qml.jacobian(cost, argnum=[0, 1])(par_0, par_1)
+
+        assert isinstance(jac, tuple)
+        num_copies = sum(
+            [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)]
+        )
+        assert len(jac) == 2
+        for j in jac:
+            assert isinstance(j, np.ndarray)
+            assert j.shape == (num_copies, 2)
+
+    def test_jacobian_var_var_multiple_params_array(
+        self, dev_name, diff_method, gradient_kwargs, shots
+    ):
+        """The jacobian of multiple measurements with a multiple params array return a single array."""
+        dev = qml.device(dev_name, wires=2, shots=shots)
+
+        @qnode(dev, diff_method=diff_method, **gradient_kwargs)
+        def circuit(a):
+            qml.RY(a[0], wires=0)
+            qml.RX(a[1], wires=0)
+            qml.RY(a[2], wires=0)
+            return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.var(qml.PauliZ(0))
+
+        a = np.array([0.1, 0.2, 0.3])
+
+        def cost(a):
+            res = circuit(a)
+            return qml.math.stack([qml.math.stack(r) for r in res])
+
+        jac = qml.jacobian(cost)(a)
+
+        assert isinstance(jac, np.ndarray)
+        num_copies = sum(
+            [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)]
+        )
+        assert jac.shape == (num_copies, 2, 3)
+
+    def test_jacobian_multiple_measurement_single_param(
+        self, dev_name, diff_method, gradient_kwargs, shots
+    ):
+        """The jacobian of multiple measurements with a single params return an array."""
+        dev = qml.device(dev_name, wires=2, shots=shots)
+
+        @qnode(dev, diff_method=diff_method, **gradient_kwargs)
+        def circuit(a):
+            qml.RY(a, wires=0)
+            qml.RX(0.2, wires=0)
+            return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1])
+
+        a = np.array(0.1)
+
+        def cost(a):
+            res = circuit(a)
+            return qml.math.stack([qml.math.hstack(r) for r in res])
+
+        jac = qml.jacobian(cost)(a)
+
+        assert isinstance(jac, np.ndarray)
+        num_copies = sum(
+            [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)]
+        )
+        assert jac.shape == (num_copies, 5)
+
+    def test_jacobian_multiple_measurement_multiple_param(
+        self, dev_name, diff_method, gradient_kwargs, shots
+    ):
+        """The jacobian of multiple measurements with a multiple params return a tuple of arrays."""
+        dev = qml.device(dev_name, wires=2, shots=shots)
+
+        @qnode(dev, diff_method=diff_method, **gradient_kwargs)
+        def circuit(a, b):
+            qml.RY(a, wires=0)
+            qml.RX(b, wires=0)
+            return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1])
+
+        a = np.array(0.1, requires_grad=True)
+        b = np.array(0.2, requires_grad=True)
+
+        def cost(a, b):
+            res = circuit(a, b)
+            return qml.math.stack([qml.math.hstack(r) for r in res])
+
+        jac = qml.jacobian(cost, argnum=[0, 1])(a, b)
+
+        assert isinstance(jac, tuple)
+        num_copies = sum(
+            [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)]
+        )
+        assert len(jac) == 2
+        for j in jac:
+            assert isinstance(j, np.ndarray)
+            assert j.shape == (num_copies, 5)
+
+    def test_jacobian_multiple_measurement_multiple_param_array(
+        self, dev_name, diff_method, gradient_kwargs, shots
+    ):
+        """The jacobian of multiple measurements with a multiple params array return a single array."""
+        dev = qml.device(dev_name, wires=2, shots=shots)
+
+        @qnode(dev, diff_method=diff_method, **gradient_kwargs)
+        def circuit(a):
+            qml.RY(a[0], wires=0)
+            qml.RX(a[1], wires=0)
+            return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1])
+
+        a = np.array([0.1, 0.2])
+
+        def cost(a):
+            res = circuit(a)
+            return qml.math.stack([qml.math.hstack(r) for r in res])
+
+        jac = qml.jacobian(cost)(a)
+
+        assert isinstance(jac, np.ndarray)
+        num_copies = sum(
+            [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)]
+        )
+        assert jac.shape == (num_copies, 5, 2)
+
+    def test_hessian_expval_multiple_params(self, dev_name, diff_method, gradient_kwargs, shots):
+        """The hessian of a single measurement with multiple params return a tuple of arrays."""
+        dev = qml.device(dev_name, wires=2, shots=shots)
+
+        par_0 = np.array(0.1)
+        par_1 = np.array(0.2)
+
+        @qnode(dev, diff_method=diff_method, max_diff=2, **gradient_kwargs)
+        def circuit(x, y):
+            qml.RX(x, wires=[0])
+            qml.RY(y, wires=[1])
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1))
+
+        def cost(x, y):
+            def cost2(x, y):
+                res = circuit(x, y)
+                return qml.math.stack(res)
+
+            return qml.math.stack(qml.jacobian(cost2, argnum=[0, 1])(x, y))
+
+        hess = qml.jacobian(cost, argnum=[0, 1])(par_0, par_1)
+
+        assert isinstance(hess, tuple)
+        num_copies = sum(
+            [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)]
+        )
+        assert len(hess) == 2
+        for h in hess:
+            assert isinstance(h, np.ndarray)
+            assert h.shape == (2, num_copies)
+
+    def test_hessian_expval_multiple_param_array(
+        self, dev_name, diff_method, gradient_kwargs, shots
+    ):
+        """The hessian of single measurement with a multiple params array return a single array."""
+        dev = qml.device(dev_name, wires=2, shots=shots)
+
+        params = np.array([0.1, 0.2])
+
+        @qnode(dev, diff_method=diff_method, max_diff=2, **gradient_kwargs)
+        def circuit(x):
+            qml.RX(x[0], wires=[0])
+            qml.RY(x[1], wires=[1])
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1))
+
+        def cost(x):
+            def cost2(x):
+                res = circuit(x)
+                return qml.math.stack(res)
+
+            return qml.jacobian(cost2)(x)
+
+        hess = qml.jacobian(cost)(params)
+
+        assert isinstance(hess, np.ndarray)
+        num_copies = sum(
+            [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)]
+        )
+        assert hess.shape == (num_copies, 2, 2)
+
+    def test_hessian_var_multiple_params(self, dev_name, diff_method, gradient_kwargs, shots):
+        """The hessian of a single measurement with multiple params return a tuple of arrays."""
+        dev = qml.device(dev_name, wires=2, shots=shots)
+
+        par_0 = np.array(0.1)
+        par_1 = np.array(0.2)
+
+        @qnode(dev, diff_method=diff_method, max_diff=2, **gradient_kwargs)
+        def circuit(x, y):
+            qml.RX(x, wires=[0])
+            qml.RY(y, wires=[1])
+            qml.CNOT(wires=[0, 1])
+            return qml.var(qml.PauliZ(0) @ qml.PauliX(1))
+
+        def cost(x, y):
+            def cost2(x, y):
+                res = circuit(x, y)
+                return qml.math.stack(res)
+
+            return qml.math.stack(qml.jacobian(cost2, argnum=[0, 1])(x, y))
+
+        hess = qml.jacobian(cost, argnum=[0, 1])(par_0, par_1)
+
+        assert isinstance(hess, tuple)
+        num_copies = sum(
+            [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)]
+        )
+        assert len(hess) == 2
+        for h in hess:
+            assert isinstance(h, np.ndarray)
+            assert h.shape == (2, num_copies)
+
+    def test_hessian_var_multiple_param_array(self, dev_name, diff_method, gradient_kwargs, shots):
+        """The hessian of single measurement with a multiple params array return a single array."""
+        dev = qml.device(dev_name, wires=2, shots=shots)
+
+        params = np.array([0.1, 0.2])
+
+        @qnode(dev, diff_method=diff_method, max_diff=2, **gradient_kwargs)
+        def circuit(x):
+            qml.RX(x[0], wires=[0])
+            qml.RY(x[1], wires=[1])
+            qml.CNOT(wires=[0, 1])
+            return qml.var(qml.PauliZ(0) @ qml.PauliX(1))
+
+        def cost(x):
+            def cost2(x):
+                res = circuit(x)
+                return qml.math.stack(res)
+
+            return qml.jacobian(cost2)(x)
+
+        hess = qml.jacobian(cost)(params)
+
+        assert isinstance(hess, np.ndarray)
+        num_copies = sum(
+            [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)]
+        )
+        assert hess.shape == (num_copies, 2, 2)
+
+    def test_hessian_probs_expval_multiple_params(
+        self, dev_name, diff_method, gradient_kwargs, shots
+    ):
+        """The hessian of multiple measurements with multiple params return a tuple of arrays."""
+
+        # override shots to speed up the test
+        shots = (5, 10)
+
+        dev = qml.device(dev_name, wires=2, shots=shots)
+
+        par_0 = np.array(0.1)
+        par_1 = np.array(0.2)
+
+        @qnode(dev, diff_method=diff_method, max_diff=2, **gradient_kwargs)
+        def circuit(x, y):
+            qml.RX(x, wires=[0])
+            qml.RY(y, wires=[1])
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0, 1])
+
+        def cost(x, y):
+            def cost2(x, y):
+                res = circuit(x, y)
+                return qml.math.stack([qml.math.hstack(r) for r in res])
+
+            return qml.math.stack(qml.jacobian(cost2, argnum=[0, 1])(x, y))
+
+        hess = qml.jacobian(cost, argnum=[0, 1])(par_0, par_1)
+
+        assert isinstance(hess, tuple)
+        num_copies = sum(
+            [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)]
+        )
+        assert len(hess) == 2
+        for h in hess:
+            assert isinstance(h, np.ndarray)
+            assert h.shape == (2, num_copies, 5)
+
+    def test_hessian_expval_probs_multiple_param_array(
+        self, dev_name, diff_method, gradient_kwargs, shots
+    ):
+        """The hessian of multiple measurements with a multiple param array return a single array."""
+
+        # override shots to speed up the test
+        shots = (5, 10)
+
+        dev = qml.device(dev_name, wires=2, shots=shots)
+
+        params = np.array([0.1, 0.2])
+
+        @qnode(dev, diff_method=diff_method, max_diff=2, **gradient_kwargs)
+        def circuit(x):
+            qml.RX(x[0], wires=[0])
+            qml.RY(x[1], wires=[1])
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0, 1])
+
+        def cost(x):
+            def cost2(x):
+                res = circuit(x)
+                return qml.math.stack([qml.math.hstack(r) for r in res])
+
+            return qml.jacobian(cost2)(x)
+
+        hess = qml.jacobian(cost)(params)
+
+        assert isinstance(hess, np.ndarray)
+        num_copies = sum(
+            [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)]
+        )
+        assert hess.shape == (num_copies, 5, 2, 2)
+
+
+qubit_device_and_diff_method = [
+    ["default.qubit", "finite-diff", {"h": 10e-2}],
+    ["default.qubit", "parameter-shift", {}],
+]
+
+finite_diff_shot_vec_tol = 0.3
+param_shift_shot_vec_tol = 10e-3
+
+shots = [(1000000, 900000, 800000), (1000000, (900000, 2))]
+
+
+@pytest.mark.parametrize("shots", shots)
+@pytest.mark.parametrize("dev_name,diff_method,gradient_kwargs", qubit_device_and_diff_method)
+class TestReturnShotVectorIntegration:
+    """Tests for the integration of shots with the autograd interface."""
+
+    def test_single_expectation_value(self, dev_name, diff_method, gradient_kwargs, shots):
+        """Tests correct output shape and evaluation for a tape
+        with a single expval output"""
+        dev = qml.device(dev_name, wires=2, shots=shots)
+        x = np.array(0.543)
+        y = np.array(-0.654)
+
+        @qnode(dev, diff_method=diff_method, **gradient_kwargs)
+        def circuit(x, y):
+            qml.RX(x, wires=[0])
+            qml.RY(y, wires=[1])
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1))
+
+        def cost(x, y):
+            res = circuit(x, y)
+            return qml.math.stack(res)
+
+        all_res = qml.jacobian(cost, argnum=[0, 1])(x, y)
+
+        assert isinstance(all_res, tuple)
+        num_copies = sum(
+            [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)]
+        )
+        assert len(all_res) == 2
+
+        expected = np.array([-np.sin(y) * np.sin(x), np.cos(y) * np.cos(x)])
+        tol = finite_diff_shot_vec_tol if diff_method == "finite-diff" else param_shift_shot_vec_tol
+
+        for res, exp in zip(all_res, expected):
+            assert isinstance(res, np.ndarray)
+            assert res.shape == (num_copies,)
+            assert np.allclose(res, exp, atol=tol, rtol=0)
+
+    def test_prob_expectation_values(self, dev_name, diff_method, gradient_kwargs, shots):
+        """Tests correct output shape and evaluation for a tape
+        with prob and expval outputs"""
+        dev = qml.device(dev_name, wires=2, shots=shots)
+        x = np.array(0.543)
+        y = np.array(-0.654)
+
+        @qnode(dev, diff_method=diff_method, **gradient_kwargs)
+        def circuit(x, y):
+            qml.RX(x, wires=[0])
+            qml.RY(y, wires=[1])
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1])
+
+        def cost(x, y):
+            res = circuit(x, y)
+            return qml.math.stack([qml.math.hstack(r) for r in res])
+
+        all_res = qml.jacobian(cost, argnum=[0, 1])(x, y)
+
+        assert isinstance(all_res, tuple)
+        num_copies = sum(
+            [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)]
+        )
+        assert len(all_res) == 2
+
+        expected = np.array(
+            [
+                [
+                    -np.sin(x),
+                    -(np.cos(y / 2) ** 2 * np.sin(x)) / 2,
+                    -(np.sin(x) * np.sin(y / 2) ** 2) / 2,
+                    (np.sin(x) * np.sin(y / 2) ** 2) / 2,
+                    (np.cos(y / 2) ** 2 * np.sin(x)) / 2,
+                ],
+                [
+                    0,
+                    -(np.cos(x / 2) ** 2 * np.sin(y)) / 2,
+                    (np.cos(x / 2) ** 2 * np.sin(y)) / 2,
+                    (np.sin(x / 2) ** 2 * np.sin(y)) / 2,
+                    -(np.sin(x / 2) ** 2 * np.sin(y)) / 2,
+                ],
+            ]
+        )
+
+        tol = finite_diff_shot_vec_tol if diff_method == "finite-diff" else param_shift_shot_vec_tol
+
+        for res, exp in zip(all_res, expected):
+            assert isinstance(res, np.ndarray)
+            assert res.shape == (num_copies, 5)
+            assert np.allclose(res, exp, atol=tol, rtol=0)