Repository: PennyLaneAI/pennylane
Pull request #3322 (2022-11-11 03:20:14)

Daily rc sync to master 2022-11-11
Automatic sync from the release candidate to master during a feature freeze.

codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3322?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#3322](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3322?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (1acc1d8) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/4dd6088e83c5eed49abd031808ded21c4f841adb?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (4dd6088) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

```diff
@@           Coverage Diff           @@
##           master    #3322   +/-   ##
=======================================
  Coverage   99.70%   99.70%           
=======================================
  Files         299      299           
  Lines       26122    26136   +14     
=======================================
+ Hits        26045    26059   +14     
  Misses         77       77           
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3322?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/\_qubit\_device.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3322/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL19xdWJpdF9kZXZpY2UucHk=) | `99.66% <ø> (ø)` | |
| [pennylane/data/data\_manager.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3322/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RhdGEvZGF0YV9tYW5hZ2VyLnB5) | `100.00% <ø> (ø)` | |
| [pennylane/gradients/finite\_difference.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3322/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2dyYWRpZW50cy9maW5pdGVfZGlmZmVyZW5jZS5weQ==) | `100.00% <ø> (ø)` | |
| [pennylane/gradients/parameter\_shift\_hessian.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3322/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2dyYWRpZW50cy9wYXJhbWV0ZXJfc2hpZnRfaGVzc2lhbi5weQ==) | `100.00% <ø> (ø)` | |
| [pennylane/ops/op\_math/sprod.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3322/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9vcF9tYXRoL3Nwcm9kLnB5) | `100.00% <ø> (ø)` | |
| [pennylane/ops/qubit/hamiltonian.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3322/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9xdWJpdC9oYW1pbHRvbmlhbi5weQ==) | `100.00% <ø> (ø)` | |
| [pennylane/optimize/qnspsa.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3322/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wdGltaXplL3Fuc3BzYS5weQ==) | `100.00% <ø> (ø)` | |
| [pennylane/return\_types.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3322/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3JldHVybl90eXBlcy5weQ==) | `100.00% <ø> (ø)` | |
| [pennylane/transforms/transpile.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3322/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvdHJhbnNwaWxlLnB5) | `100.00% <ø> (ø)` | |
| [pennylane/data/dataset.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3322/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RhdGEvZGF0YXNldC5weQ==) | `100.00% <100.00%> (ø)` | |
| ... and [2 more](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3322/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | |

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/introduction/data.rst b/doc/introduction/data.rst
index 33c3936ab..f77377575 100644
--- a/doc/introduction/data.rst
+++ b/doc/introduction/data.rst
@@ -28,9 +28,19 @@ The :func:`~pennylane.data.load` function returns a ``list`` with the desired da
 
 >>> H2datasets = qml.data.load("qchem", molname="H2", basis="STO-3G", bondlength=1.1)
 >>> print(H2datasets)
-[<pennylane.data.dataset.Dataset object at 0x7f14e4369640>]
+[<Dataset = description: qchem/H2/STO-3G/1.1, attributes: ['molecule', 'hamiltonian', ...]>]
 >>> H2data = H2datasets[0]
 
+We can load datasets for multiple parameter values by providing a list of values instead of a single value.
+To load all possible values, use the special keyword "full".
+
+>>> H2datasets = qml.data.load("qchem", molname="H2", basis="full", bondlength=[0.5, 1.1])
+>>> print(H2datasets)
+[<Dataset = description: qchem/H2/6-31G/0.5, attributes: ['molecule', 'hamiltonian', ...]>,
+ <Dataset = description: qchem/H2/6-31G/1.1, attributes: ['molecule', 'hamiltonian', ...]>,
+ <Dataset = description: qchem/H2/STO-3G/0.5, attributes: ['molecule', 'hamiltonian', ...]>,
+ <Dataset = description: qchem/H2/STO-3G/1.1, attributes: ['molecule', 'hamiltonian', ...]>]
+
 When we only want to download portions of a large dataset, we can specify the desired properties  (referred to as `attributes`).
 For example, we can download or load only the molecule and energy of a dataset as follows:
 
@@ -51,14 +61,18 @@ To determine what attributes are available for a type of dataset, we can use the
 "tapered_hamiltonian",
 "full"]
 
+.. note::
+
+    "full" is the default value for ``attributes``, and it means that all available attributes for the Dataset will be downloaded.
+
 Using Datasets in PennyLane
 ---------------------------
 
 Once loaded, one can access properties of the datasets:
 
->>> H2_dataset.molecule
+>>> H2data.molecule
 <pennylane.qchem.molecule.Molecule object at 0x7f890b409280>
->>> print(H2_dataset.hf_state)
+>>> print(H2data.hf_state)
 [1 1 0 0]
 
 The loaded data items are fully compatible with PennyLane. We can therefore
@@ -67,10 +81,10 @@ use them directly in a PennyLane circuits as follows:
 >>> dev = qml.device("default.qubit",wires=4)
 >>> @qml.qnode(dev)
 ... def circuit():
-...     qml.BasisState(H2_dataset.hf_state, wires = [0, 1, 2, 3])
-...     for op in H2_dataset.vqe_gates:
+...     qml.BasisState(H2data.hf_state, wires = [0, 1, 2, 3])
+...     for op in H2data.vqe_gates:
 ...         qml.apply(op)
-...     return qml.expval(H2_dataset.hamiltonian)
+...     return qml.expval(H2data.hamiltonian)
 >>> print(circuit())
 -1.0791430411076344
 
diff --git a/doc/releases/changelog-0.27.0.md b/doc/releases/changelog-0.27.0.md
index d537a35e1..f5988e674 100644
--- a/doc/releases/changelog-0.27.0.md
+++ b/doc/releases/changelog-0.27.0.md
@@ -56,7 +56,7 @@
 
 <h4>(TODO: title) QChem</h4>
 
-* Grouped coefficients, observables, and basis rotation transformation matrices needed to construct a qubit Hamiltonian in the rotated basis of molecular orbitals are now calculable via `qml.qchem.basis_rotation()`. 
+* Grouped coefficients, observables, and basis rotation transformation matrices needed to construct a qubit Hamiltonian in the rotated basis of molecular orbitals are now calculable via `qml.qchem.basis_rotation()`.
   ([#3011](https://github.com/PennyLaneAI/pennylane/pull/3011))
 
   ```pycon
@@ -199,7 +199,7 @@
   Result value: 1.00; Result type: <class 'jaxlib.xla_extension.DeviceArray'>
   ```
 
-* Wires of operators or entire QNodes can now be mapped to other wires via `qml.map_wires()`. 
+* Wires of operators or entire QNodes can now be mapped to other wires via `qml.map_wires()`.
   [(#3145)](https://github.com/PennyLaneAI/pennylane/pull/3145)
 
   The `qml.map_wires()` function requires a dictionary representing a wire map. Use it with
@@ -221,7 +221,7 @@
     dev = qml.device("default.qubit", wires=["A", "B", "C", "D"])
     wire_map = {0: "A", 1: "B", 2: "C", 3: "D"}
 
-    @qml.qnode(dev) 
+    @qml.qnode(dev)
     def circuit():
         qml.RX(0.54, wires=0)
         qml.PauliX(1)
@@ -271,7 +271,7 @@
 * To download or load only specific properties of a dataset, we can specify the desired attributes in `qml.data.load`:
 
   ```pycon
-  >>> part = qml.data.load("qchem", molname="H2", basis="STO-3G", bondlength=1.1, 
+  >>> part = qml.data.load("qchem", molname="H2", basis="STO-3G", bondlength=1.1,
   ...                      attributes=["molecule", "fci_energy"])[0]
   >>> part.molecule
   <pennylane.qchem.molecule.Molecule at 0x7f56c9d78e50>
@@ -394,7 +394,7 @@
 * `OrbitalRotation` is now decomposed into two `SingleExcitation` operations for faster execution and more efficient parameter-shift gradient calculations on devices that natively support `SingleExcitation`.
   [(#3171)](https://github.com/PennyLaneAI/pennylane/pull/3171)
 
-* Reorganized and grouped all functions in PennyLane responsible for manipulation of Pauli operators into a `pauli` 
+* Reorganized and grouped all functions in PennyLane responsible for manipulation of Pauli operators into a `pauli`
   module. Deprecated the `grouping` module and moved logic from `pennylane/grouping` to `pennylane/pauli/grouping`.
   [(#3179)](https://github.com/PennyLaneAI/pennylane/pull/3179)
 
@@ -574,7 +574,7 @@
 
 * `Operator.compute_terms` has been removed. On a specific instance of an operator, use `op.terms()` instead. There is no longer a static method for this.
   [(#3215)](https://github.com/PennyLaneAI/pennylane/pull/3215)
-  
+
 * `qml.tape.QuantumTape.inv()` has been deprecated. Use `qml.tape.QuantumTape.adjoint` instead.
   [(#3237)](https://github.com/PennyLaneAI/pennylane/pull/3237)
 
@@ -651,6 +651,9 @@
   if the Hamiltonian involved some wires that are not present on the device.
   [(#3266)](https://github.com/PennyLaneAI/pennylane/pull/3266)
 
+* Fixed a bug where `qml.tape.QuantumTape.shape()` did not account for the batch dimension of the tape
+  [(#3269)](https://github.com/PennyLaneAI/pennylane/pull/3269)
+
 <h3>Contributors</h3>
 
 This release contains contributions from (in alphabetical order):
diff --git a/doc/requirements.txt b/doc/requirements.txt
index cc3b1bbed..62eebac6a 100644
--- a/doc/requirements.txt
+++ b/doc/requirements.txt
@@ -3,8 +3,8 @@ pip
 appdirs
 autograd
 autoray
-jax==0.3.14
-jaxlib==0.3.14
+jax==0.3.17
+jaxlib==0.3.15
 mistune==0.8.4
 m2r2
 numpy
diff --git a/pennylane/_qubit_device.py b/pennylane/_qubit_device.py
index d8df916ab..558b8ad75 100644
--- a/pennylane/_qubit_device.py
+++ b/pennylane/_qubit_device.py
@@ -1810,7 +1810,7 @@ class QubitDevice(Device):
                 provided, that takes precedence.
 
         Returns:
-            array: the derivative of the tape with respect to trainable parameters.
+            array or tuple[array]: the derivative of the tape with respect to trainable parameters.
             Dimensions are ``(len(observables), len(trainable_params))``.
 
         Raises:
diff --git a/pennylane/data/data_manager.py b/pennylane/data/data_manager.py
index f9f1babca..5a5cb9d34 100644
--- a/pennylane/data/data_manager.py
+++ b/pennylane/data/data_manager.py
@@ -392,7 +392,7 @@ def load_interactive():
         force: False
         dest folder: /Users/jovyan/Downloads/datasets
         Would you like to continue? (Default is yes) [Y/n]:
-        <pennylane.data.dataset.Dataset object at 0x10157ab50>
+        <Dataset = description: qspin/Ising/open/rectangular/4x4, attributes: ['parameters', 'ground_states']>
     """
 
     _refresh_foldermap()
diff --git a/pennylane/data/dataset.py b/pennylane/data/dataset.py
index d496fee22..5b2777960 100644
--- a/pennylane/data/dataset.py
+++ b/pennylane/data/dataset.py
@@ -37,6 +37,7 @@ def _import_zstd_dill():
     return zstd, dill
 
 
+# pylint: disable=too-many-instance-attributes
 class Dataset(ABC):
     """Create a dataset object to store a collection of information describing
     a physical system and its evolution. For example, a dataset for an arbitrary
@@ -108,6 +109,7 @@ class Dataset(ABC):
         self._folder = folder.rstrip("/")
         self._prefix = os.path.join(self._folder, attr_prefix) + "_{}.dat"
         self._prefix_len = len(attr_prefix) + 1
+        self._description = os.path.join(data_name, self._folder.split(data_name)[-1][1:])
         self.__doc__ = docstring
 
         self._fullfile = self._prefix.format("full")
@@ -138,6 +140,15 @@ class Dataset(ABC):
             self._is_standard = False
             self.__base_init__(**kwargs)
 
+    def __repr__(self):
+        attr_str = (
+            str(list(self.attrs))
+            if len(self.attrs) < 3
+            else str(list(self.attrs)[:2])[:-1] + ", ...]"
+        )
+        std_str = f"description: {self._description}, " if self._is_standard else ""
+        return f"<Dataset = {std_str}attributes: {attr_str}>"
+
     @property
     def attrs(self):
         """Returns attributes of the dataset."""
diff --git a/pennylane/gradients/finite_difference.py b/pennylane/gradients/finite_difference.py
index 17eb8303e..4e5ea0f87 100644
--- a/pennylane/gradients/finite_difference.py
+++ b/pennylane/gradients/finite_difference.py
@@ -250,11 +250,11 @@ def _finite_diff_new(
             device shots for the new return types output system.
 
     Returns:
-        tensor_like or tuple[list[QuantumTape], function]:
+        tensor_like or tuple[tensor_like] or tuple[tuple[tensor_like]] or tuple[list[QuantumTape], function]:
 
-        - If the input is a QNode, a tensor
-          representing the output Jacobian matrix of size ``(number_outputs, number_gate_parameters)``
-          is returned.
+        - If the input is a QNode, an object representing the output Jacobian matrix.
+          The type of the object returned is either a tensor, a tuple or a nested tuple depending on the nesting
+          structure of the output.
 
         - If the input is a tape, a tuple containing a list of generated tapes,
           in addition to a post-processing function to be applied to the
@@ -262,7 +262,106 @@ def _finite_diff_new(
 
     **Example**
 
-    #TODO: Add example for new return type.
+    This transform can be registered directly as the quantum gradient transform
+    to use during autodifferentiation:
+
+    >>> dev = qml.device("default.qubit", wires=2)
+    >>> @qml.qnode(dev, interface="autograd", diff_method="finite-diff")
+    ... def circuit(params):
+    ...     qml.RX(params[0], wires=0)
+    ...     qml.RY(params[1], wires=0)
+    ...     qml.RX(params[2], wires=0)
+    ...     return qml.expval(qml.PauliZ(0))
+    >>> params = np.array([0.1, 0.2, 0.3], requires_grad=True)
+    >>> qml.jacobian(circuit)(params)
+    tensor([-0.38751725, -0.18884792, -0.38355708], requires_grad=True)
+
+    When differentiating QNodes with multiple measurements using Autograd or TensorFlow, the outputs of the QNode first
+    need to be stacked. The reason is that those two frameworks only allow differentiating functions with array or
+    tensor outputs, instead of functions that output sequences. In contrast, Jax and Torch require no additional
+    post-processing.
+
+    >>> import jax
+    >>> dev = qml.device("default.qubit", wires=2)
+    >>> @qml.qnode(dev, interface="jax", diff_method="finite-diff")
+    ... def circuit(params):
+    ...     qml.RX(params[0], wires=0)
+    ...     qml.RY(params[1], wires=0)
+    ...     qml.RX(params[2], wires=0)
+    ...     return qml.expval(qml.PauliZ(0)), qml.var(qml.PauliZ(0))
+    >>> params = jax.numpy.array([0.1, 0.2, 0.3])
+    >>> jax.jacobian(circuit)(params)
+    (DeviceArray([-0.38751727, -0.18884793, -0.3835571 ], dtype=float32), DeviceArray([0.6991687 , 0.34072432, 0.6920237 ], dtype=float32))
+
+
+    .. details::
+        :title: Usage Details
+
+        This gradient transform can also be applied directly to :class:`QNode <pennylane.QNode>` objects:
+
+        >>> @qml.qnode(dev)
+        ... def circuit(params):
+        ...     qml.RX(params[0], wires=0)
+        ...     qml.RY(params[1], wires=0)
+        ...     qml.RX(params[2], wires=0)
+        ...     return qml.expval(qml.PauliZ(0)), qml.var(qml.PauliZ(0))
+        >>> qml.gradients.finite_diff(circuit)(params)
+        ((tensor(-0.38751724, requires_grad=True),
+          tensor(-0.18884792, requires_grad=True),
+          tensor(-0.38355709, requires_grad=True)),
+         (tensor(0.69916868, requires_grad=True),
+          tensor(0.34072432, requires_grad=True),
+          tensor(0.69202366, requires_grad=True)))
+
+        This quantum gradient transform can also be applied to low-level
+        :class:`~.QuantumTape` objects. This will result in no implicit quantum
+        device evaluation. Instead, the processed tapes, and post-processing
+        function, which together define the gradient are directly returned:
+
+        >>> with qml.tape.QuantumTape() as tape:
+        ...     qml.RX(params[0], wires=0)
+        ...     qml.RY(params[1], wires=0)
+        ...     qml.RX(params[2], wires=0)
+        ...     qml.expval(qml.PauliZ(0))
+        ...     qml.var(qml.PauliZ(0))
+        >>> gradient_tapes, fn = qml.gradients.finite_diff(tape)
+        >>> gradient_tapes
+        [<QuantumTape: wires=[0], params=3>,
+         <QuantumTape: wires=[0], params=3>,
+         <QuantumTape: wires=[0], params=3>,
+         <QuantumTape: wires=[0], params=3>]
+
+        This can be useful if the underlying circuits representing the gradient
+        computation need to be analyzed.
+
+        The output tapes can then be evaluated and post-processed to retrieve
+        the gradient:
+
+        >>> dev = qml.device("default.qubit", wires=2)
+        >>> fn(qml.execute(gradient_tapes, dev, None))
+        ((array(-0.38751724), array(-0.18884792), array(-0.38355709)),
+         (array(0.69916868), array(0.34072432), array(0.69202366)))
+
+        Devices that have a shot vector defined can also be used for execution, provided
+        the ``shots`` argument was passed to the transform:
+
+        >>> shots = (10, 100, 1000)
+        >>> dev = qml.device("default.qubit", wires=2, shots=shots)
+        >>> @qml.qnode(dev)
+        ... def circuit(params):
+        ...     qml.RX(params[0], wires=0)
+        ...     qml.RY(params[1], wires=0)
+        ...     qml.RX(params[2], wires=0)
+        ...     return qml.expval(qml.PauliZ(0)), qml.var(qml.PauliZ(0))
+        >>> params = np.array([0.1, 0.2, 0.3], requires_grad=True)
+        >>> qml.gradients.finite_diff(circuit, shots=shots, h=10e-2)(params)
+        (((array(-2.), array(-2.), array(0.)), (array(3.6), array(3.6), array(0.))),
+         ((array(1.), array(0.4), array(1.)),
+          (array(-1.62), array(-0.624), array(-1.62))),
+         ((array(-0.48), array(-0.34), array(-0.46)),
+          (array(0.84288), array(0.6018), array(0.80868))))
+
+        The outermost tuple contains results corresponding to each element of the shot vector.
     """
     if argnum is None and not tape.trainable_params:
         return _no_trainable_grad_new(tape, shots)
diff --git a/pennylane/gradients/parameter_shift.py b/pennylane/gradients/parameter_shift.py
index 2cb4bf586..78edc7f0d 100644
--- a/pennylane/gradients/parameter_shift.py
+++ b/pennylane/gradients/parameter_shift.py
@@ -1157,7 +1157,6 @@ def var_param_shift(
     return gradient_tapes, processing_fn
 
 
-# TODO: docstrings & mention shots arg
 @gradient_transform
 def _param_shift_new(
     tape,
@@ -1210,11 +1209,11 @@ def _param_shift_new(
             device shots for the new return types output system.
 
     Returns:
-        tensor_like or tuple[list[QuantumTape], function]:
+        tensor_like or tuple[tensor_like] or tuple[tuple[tensor_like]] or tuple[list[QuantumTape], function]:
 
-        - If the input is a QNode, a tensor
-          representing the output Jacobian matrix of size ``(number_outputs, number_gate_parameters)``
-          is returned.
+        - If the input is a QNode, an object representing the output Jacobian matrix.
+          The type of the object returned is either a tensor, a tuple or a nested tuple depending on the nesting
+          structure of the output.
 
         - If the input is a tape, a tuple containing a list of generated tapes,
           in addition to a post-processing function to be applied to the
@@ -1282,16 +1281,32 @@ def _param_shift_new(
     to use during autodifferentiation:
 
     >>> dev = qml.device("default.qubit", wires=2)
-    >>> @qml.qnode(dev, gradient_fn=qml.gradients.param_shift)
+    >>> @qml.qnode(dev, interface="autograd", diff_method="parameter-shift")
     ... def circuit(params):
     ...     qml.RX(params[0], wires=0)
     ...     qml.RY(params[1], wires=0)
     ...     qml.RX(params[2], wires=0)
-    ...     return qml.expval(qml.PauliZ(0)), qml.var(qml.PauliZ(0))
+    ...     return qml.expval(qml.PauliZ(0))
     >>> params = np.array([0.1, 0.2, 0.3], requires_grad=True)
     >>> qml.jacobian(circuit)(params)
-    tensor([[-0.38751725, -0.18884792, -0.38355708],
-            [ 0.69916868,  0.34072432,  0.69202365]], requires_grad=True)
+    tensor([-0.38751725, -0.18884792, -0.38355708], requires_grad=True)
+
+    When differentiating QNodes with multiple measurements using Autograd or TensorFlow, the outputs of the QNode first
+    need to be stacked. The reason is that those two frameworks only allow differentiating functions with array or
+    tensor outputs, instead of functions that output sequences. In contrast, Jax and Torch require no additional
+    post-processing.
+
+    >>> import jax
+    >>> dev = qml.device("default.qubit", wires=2)
+    >>> @qml.qnode(dev, interface="jax", diff_method="parameter-shift")
+    ... def circuit(params):
+    ...     qml.RX(params[0], wires=0)
+    ...     qml.RY(params[1], wires=0)
+    ...     qml.RX(params[2], wires=0)
+    ...     return qml.expval(qml.PauliZ(0)), qml.var(qml.PauliZ(0))
+    >>> params = jax.numpy.array([0.1, 0.2, 0.3])
+    >>> jax.jacobian(circuit)(params)
+    (DeviceArray([-0.38751727, -0.18884793, -0.3835571 ], dtype=float32), DeviceArray([0.6991687 , 0.34072432, 0.6920237 ], dtype=float32))
 
     .. note::
 
@@ -1335,8 +1350,12 @@ def _param_shift_new(
         ...     qml.RX(params[2], wires=0)
         ...     return qml.expval(qml.PauliZ(0)), qml.var(qml.PauliZ(0))
         >>> qml.gradients.param_shift(circuit)(params)
-        tensor([[-0.38751725, -0.18884792, -0.38355708],
-                [ 0.69916868,  0.34072432,  0.69202365]], requires_grad=True)
+        ((tensor(-0.38751724, requires_grad=True),
+          tensor(-0.18884792, requires_grad=True),
+          tensor(-0.38355709, requires_grad=True)),
+         (tensor(0.69916868, requires_grad=True),
+          tensor(0.34072432, requires_grad=True),
+          tensor(0.69202366, requires_grad=True)))
 
         This quantum gradient transform can also be applied to low-level
         :class:`~.QuantumTape` objects. This will result in no implicit quantum
@@ -1366,8 +1385,30 @@ def _param_shift_new(
 
         >>> dev = qml.device("default.qubit", wires=2)
         >>> fn(qml.execute(gradient_tapes, dev, None))
-        [[-0.38751721 -0.18884787 -0.38355704]
-         [ 0.69916862  0.34072424  0.69202359]]
+        ((array(-0.3875172), array(-0.18884787), array(-0.38355704)),
+         (array(0.69916862), array(0.34072424), array(0.69202359)))
+
+        Devices that have a shot vector defined can also be used for execution, provided
+        the ``shots`` argument was passed to the transform:
+
+        >>> shots = (10, 100, 1000)
+        >>> dev = qml.device("default.qubit", wires=2, shots=shots)
+        >>> @qml.qnode(dev)
+        ... def circuit(params):
+        ...     qml.RX(params[0], wires=0)
+        ...     qml.RY(params[1], wires=0)
+        ...     qml.RX(params[2], wires=0)
+        ...     return qml.expval(qml.PauliZ(0)), qml.var(qml.PauliZ(0))
+        >>> params = np.array([0.1, 0.2, 0.3], requires_grad=True)
+        >>> qml.gradients.param_shift(circuit, shots=shots)(params)
+        (((array(-0.6), array(-0.1), array(-0.1)),
+          (array(1.2), array(0.2), array(0.2))),
+         ((array(-0.39), array(-0.24), array(-0.49)),
+          (array(0.7488), array(0.4608), array(0.9408))),
+         ((array(-0.36), array(-0.191), array(-0.37)),
+          (array(0.65808), array(0.349148), array(0.67636))))
+
+        The outermost tuple contains results corresponding to each element of the shot vector.
 
         When setting the keyword argument ``broadcast`` to ``True``, the shifted
         circuit evaluations for each operation are batched together, resulting in
@@ -1388,10 +1429,16 @@ def _param_shift_new(
         The postprocessing function will know that broadcasting is used and handle
         the results accordingly:
         >>> fn(qml.execute(gradient_tapes, dev, None))
-        array([[-0.3875172 , -0.18884787, -0.38355704]])
+        (array(-0.3875172), array(-0.18884787), array(-0.38355704))
 
         An advantage of using ``broadcast=True`` is a speedup:
 
+        >>> @qml.qnode(dev)
+        ... def circuit(params):
+        ...     qml.RX(params[0], wires=0)
+        ...     qml.RY(params[1], wires=0)
+        ...     qml.RX(params[2], wires=0)
+        ...     return qml.expval(qml.PauliZ(0))
         >>> number = 100
         >>> serial_call = "qml.gradients.param_shift(circuit, broadcast=False)(params)"
         >>> timeit.timeit(serial_call, globals=globals(), number=number) / number
@@ -1768,6 +1815,12 @@ def param_shift(
 
         An advantage of using ``broadcast=True`` is a speedup:
 
+        >>> @qml.qnode(dev)
+        ... def circuit(params):
+        ...     qml.RX(params[0], wires=0)
+        ...     qml.RY(params[1], wires=0)
+        ...     qml.RX(params[2], wires=0)
+        ...     return qml.expval(qml.PauliZ(0))
         >>> number = 100
         >>> serial_call = "qml.gradients.param_shift(circuit, broadcast=False)(params)"
         >>> timeit.timeit(serial_call, globals=globals(), number=number) / number
diff --git a/pennylane/gradients/parameter_shift_hessian.py b/pennylane/gradients/parameter_shift_hessian.py
index c03cecd8f..7f588d95b 100644
--- a/pennylane/gradients/parameter_shift_hessian.py
+++ b/pennylane/gradients/parameter_shift_hessian.py
@@ -704,9 +704,7 @@ def param_shift_hessian(tape, argnum=None, diagonal_shifts=None, off_diagonal_sh
 def _param_shift_hessian_tuple(
     tape, argnum=None, diagonal_shifts=None, off_diagonal_shifts=None, f0=None
 ):
-    r"""TODO: Update docstring examples
-
-    Transform a QNode to compute the parameter-shift Hessian with respect to its trainable
+    r"""Transform a QNode to compute the parameter-shift Hessian with respect to its trainable
     parameters. This is the Hessian transform to replace the old one in the new return types system
 
     Use this transform to explicitly generate and explore parameter-shift circuits for computing
@@ -773,8 +771,8 @@ def _param_shift_hessian_tuple(
 
     >>> x = np.array([0.5, 0.2], requires_grad=True)
     >>> qml.gradients.param_shift_hessian(circuit)(x)
-    tensor([[-0.86883595,  0.04762358],
-            [ 0.04762358,  0.05998862]], requires_grad=True)
+    ((array(-0.86883595), array(0.04762358)),
+     (array(0.04762358), array(0.05998862)))
 
     .. details::
         :title: Usage Details
@@ -791,8 +789,8 @@ def _param_shift_hessian_tuple(
         >>> all(isinstance(tape, qml.tape.QuantumTape) for tape in hessian_tapes)
         True
         >>> postproc_fn(qml.execute(hessian_tapes, dev, None))
-        array([[-0.86883595,  0.04762358],
-               [ 0.04762358,  0.05998862]])
+        ((array(-0.86883595), array(0.04762358)),
+         (array(0.04762358), array(0.05998862)))
 
         The Hessian tapes can be inspected via their draw function, which reveals the different
         gate arguments generated from parameter-shift rules (we only draw the first four out of
@@ -844,8 +842,7 @@ def _param_shift_hessian_tuple(
 
         >>> hessian_tapes, postproc_fn = qml.gradients.param_shift_hessian(tape, argnum=(1,))
         >>> postproc_fn(qml.execute(hessian_tapes, dev, None))
-        array([[0.        , 0.        ],
-               [0.        , 0.05998862]])
+        ((array(0.), array(0.)), (array(0.), array(0.05998862)))
 
     """
 
diff --git a/pennylane/ops/op_math/sprod.py b/pennylane/ops/op_math/sprod.py
index 83bb3bf39..8fa21ce9b 100644
--- a/pennylane/ops/op_math/sprod.py
+++ b/pennylane/ops/op_math/sprod.py
@@ -264,13 +264,10 @@ class SProd(SymbolicOp):
         return SProd(scalar=self.scalar, base=new_base)
 
     def _check_scalar_is_valid(self):
-        """Check that the given scalar is valid.
-
-        Args:
-            scalar (Union[int, float, complex]): scalar value to check
+        """Check that ``self.scalar`` is valid.
 
         Raises:
-            ValueError: if the scalar is not valid
+            ValueError: if ``self.scalar`` is not valid
         """
         backend = autoray.infer_backend(self.scalar)
         # TODO: Remove shape check when supporting batching
diff --git a/pennylane/ops/qubit/hamiltonian.py b/pennylane/ops/qubit/hamiltonian.py
index 7628f06fa..cf666e88b 100644
--- a/pennylane/ops/qubit/hamiltonian.py
+++ b/pennylane/ops/qubit/hamiltonian.py
@@ -488,12 +488,11 @@ class Hamiltonian(Observable):
 
         **Examples**
 
-        >>> A = np.array([[1, 0], [0, -1]])
         >>> H = qml.Hamiltonian(
         ...     [0.5, 0.5],
-        ...     [qml.Hermitian(A, 0) @ qml.PauliY(1), qml.PauliY(1) @ qml.Hermitian(A, 0) @ qml.Identity("a")]
+        ...     [qml.PauliZ(0) @ qml.PauliY(1), qml.PauliY(1) @ qml.PauliZ(0) @ qml.Identity("a")]
         ... )
-        >>> obs = qml.Hermitian(A, 0) @ qml.PauliY(1)
+        >>> obs = qml.PauliZ(0) @ qml.PauliY(1)
         >>> print(H.compare(obs))
         True
 
diff --git a/pennylane/optimize/qnspsa.py b/pennylane/optimize/qnspsa.py
index 2e634f8c4..ae2265816 100644
--- a/pennylane/optimize/qnspsa.py
+++ b/pennylane/optimize/qnspsa.py
@@ -95,9 +95,9 @@ class QNSPSAOptimizer:
     >>> params = np.random.rand(2)
     >>> opt = QNSPSAOptimizer(stepsize=5e-2)
     >>> for i in range(51):
-    >>> params, loss = opt.step_and_cost(cost, params)
-    >>> if i % 10 == 0:
-    ...     print(f"Step {i}: cost = {loss:.4f}")
+    >>>     params, loss = opt.step_and_cost(cost, params)
+    >>>     if i % 10 == 0:
+    ...         print(f"Step {i}: cost = {loss:.4f}")
     Step 0: cost = 0.9987
     Step 10: cost = 0.9841
     Step 20: cost = 0.8921
diff --git a/pennylane/return_types.py b/pennylane/return_types.py
index 9a65c3b0c..7f9c68014 100644
--- a/pennylane/return_types.py
+++ b/pennylane/return_types.py
@@ -19,9 +19,11 @@ __activated = False
 
 
 def enable_return():
-    """Function that turns on the new return type system. The new system guarantees intuitive return types such that a
-    sequence (e.g., list or tuple) is returned based on the ``return`` statement of the quantum function. This system
-    avoids the creation of ragged arrays, where multiple measurements are stacked together.
+    """Function that turns on the experimental return type system that prefers the use of sequences over arrays.
+
+    The new system guarantees that a sequence (e.g., list or tuple) is returned based on the ``return`` statement of the
+    quantum function. This system avoids the creation of ragged arrays, where multiple measurements are stacked
+    together.
 
     **Example**
 
@@ -38,8 +40,7 @@ def enable_return():
 
         qnode = qml.QNode(circuit, dev)
 
-    >>> res = qnode(0.5)
-    >>> res
+    >>> qnode(0.5)
     tensor([0.5       , 0.5       , 0.08014815, 0.96939564, 0.03060436,
         0.93879128], requires_grad=True)
 
@@ -58,45 +59,224 @@ def enable_return():
 
         qnode = qml.QNode(circuit, dev)
 
-    >>> res = qnode(0.5)
-    >>> res
+    >>> qnode(0.5)
     (tensor([0.5, 0.5], requires_grad=True), tensor(0.08014815, requires_grad=True), tensor([0.96939564, 0.03060436], requires_grad=True), tensor(0.93879128, requires_grad=True))
 
-    The new return types system unlocks the use of ``probs`` mixed with different measurements in backpropagation with JAX:
+    .. note::
 
-    .. code-block:: python
+        This is an experimental feature and may not support every feature in PennyLane. The list of supported features
+        from PennyLane include:
 
-        import jax
+        * :func:`~.pennylane.execute`
+        * Gradient transforms
 
-        qml.enable_return()
+          #. :func:`~.pennylane.gradients.param_shift`;
+          #. :func:`~.pennylane.gradients.finite_diff`;
+          #. :class:`~.pennylane.gradients.hessian_transform`;
+          #. :func:`~.pennylane.gradients.param_shift_hessian`.
 
-        dev = qml.device("default.qubit", wires=2)
-        qml.enable_return()
+        * Interfaces
 
-        @qml.qnode(dev, interface="jax")
-        def circuit(a):
-          qml.RX(a[0], wires=0)
-          qml.CNOT(wires=(0, 1))
-          qml.RY(a[1], wires=1)
-          qml.RZ(a[2], wires=1)
-          return qml.expval(qml.PauliZ(wires=0)), qml.probs(wires=[0, 1]), qml.vn_entropy(wires=1)
+          #. Autograd;
+          #. TensorFlow;
+          #. JAX (without jitting);
 
-        x = jax.numpy.array([0.1, 0.2, 0.3])
+        * PennyLane optimizers
+        * :meth:`~.pennylane.tape.QuantumTape.shape`
 
-    >>> res = jax.jacobian(circuit)(x)
-    >>> res
-    (DeviceArray([-9.9833414e-02, -7.4505806e-09, -3.9932679e-10], dtype=float32),
-    DeviceArray([[-4.9419206e-02, -9.9086545e-02,  3.4938008e-09],
-               [-4.9750542e-04,  9.9086538e-02,  1.2768372e-10],
-               [ 4.9750548e-04,  2.4812977e-04,  4.8371929e-13],
-               [ 4.9419202e-02, -2.4812980e-04,  2.6696912e-11]],            dtype=float32),
-    DeviceArray([ 2.9899091e-01, -4.4703484e-08,  9.5104014e-10], dtype=float32))
+    Note that this is an experimental feature and may not support every feature in PennyLane. See the ``Usage Details``
+    section for more details.
+
+    .. details::
+        :title: Usage Details
+
+        **Gotcha: Autograd and TensorFlow can only compute gradients of tensor-valued functions**
+
+        Autograd and TensorFlow only allow differentiating functions that have array or tensor outputs. QNodes that
+        have multiple measurements may output other sequences with the new return types may cause errors with Autograd
+        or TensorFlow.
+
+        This issue can be overcome by stacking the QNode results before computing derivatives:
+
+        .. code-block:: python
+
+            qml.enable_return()
+
+            a = np.array(0.1, requires_grad=True)
+            b = np.array(0.2, requires_grad=True)
+
+            dev = qml.device("lightning.qubit", wires=2)
+
+            @qml.qnode(dev, diff_method="parameter-shift")
+            def circuit(a, b):
+                qml.RY(a, wires=0)
+                qml.RX(b, wires=1)
+                qml.CNOT(wires=[0, 1])
+                return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))
+
+            def cost(x, y):
+                return qml.numpy.hstack(circuit(x, y))
+
+        >>> qml.jacobian(cost)(a, b)
+        (array([-0.09983342,  0.01983384]), array([-5.54649074e-19, -9.75170327e-01]))
+
+        If no stacking is performed, Autograd raises the following error:
+
+        .. code-block:: python
+
+            TypeError: 'ArrayVSpace' object cannot be interpreted as an integer
+
+        The solution with TensorFlow is similar with the difference that stacking happens within the
+        ``tf.GradientTape()`` context:
+
+        .. code-block:: python
+
+            a = tf.Variable(0.1, dtype=tf.float64)
+            b = tf.Variable(0.2, dtype=tf.float64)
+
+            dev = qml.device("default.qubit", wires=2)
+
+            @qml.qnode(dev, diff_method="parameter-shift", interface="tf")
+            def circuit(a, b):
+                qml.RY(a, wires=0)
+                qml.RX(b, wires=1)
+                qml.CNOT(wires=[0, 1])
+                return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))
+
+            with tf.GradientTape() as tape:
+                res = circuit(a, b)
+                res = tf.stack(res)
+
+            assert circuit.qtape.trainable_params == [0, 1]
+
+            tape.jacobian(res, [a, b])
+
+        If the measurements do not have the same shape then you need to use concatenation:
+
+        .. code-block:: python
+
+            a = tf.Variable(0.1, dtype=tf.float64)
+            b = tf.Variable(0.2, dtype=tf.float64)
 
-    where before the following error was raised:
+            dev = qml.device("default.qubit", wires=2)
 
-    .. code-block:: text
+            @qml.qnode(dev, diff_method="parameter-shift", interface="tf")
+            def circuit(a, b):
+                qml.RY(a, wires=0)
+                qml.RX(b, wires=1)
+                qml.CNOT(wires=[0, 1])
+                return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1])
+
+            with tf.GradientTape() as tape:
+                res = circuit(a, b)
+                res = tf.concat([tf.reshape(i, [-1]) for i in res], 0)
+
+            assert circuit.qtape.trainable_params == [0, 1]
+
+            tape.jacobian(res, [a, b])
+
+        If no stacking is performed, TensorFlow raises the following error:
+
+        .. code-block:: python
+
+            AttributeError: 'tuple' object has no attribute 'shape'
+
+        **JAX interface upgrades: higher-order derivatives and mixing measurements**
+
+        Higher-order derivatives can now be computed with the JAX interface:
+
+        .. code-block:: python
+
+            import jax
+
+            qml.enable_return()
+
+            dev = qml.device("lightning.qubit", wires=2)
+
+            par_0 = jax.numpy.array(0.1)
+            par_1 = jax.numpy.array(0.2)
+
+            @qml.qnode(dev, interface="jax", diff_method="parameter-shift", max_diff=2)
+            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))
+
+        >>> jax.hessian(circuit, argnums=[0, 1])(par_0, par_1)
+        ((DeviceArray(-0.19767681, dtype=float32, weak_type=True),
+          DeviceArray(-0.09784342, dtype=float32, weak_type=True)),
+         (DeviceArray(-0.09784339, dtype=float32, weak_type=True),
+          DeviceArray(-0.19767687, dtype=float32, weak_type=True)))
+
+        The new return types system also unlocks the use of ``probs`` mixed with different measurements with JAX:
+
+        .. code-block:: python
+
+            import jax
+
+            qml.enable_return()
+
+            dev = qml.device("default.qubit", wires=2)
+            qml.enable_return()
+
+            @qml.qnode(dev, interface="jax")
+            def circuit(a):
+              qml.RX(a[0], wires=0)
+              qml.CNOT(wires=(0, 1))
+              qml.RY(a[1], wires=1)
+              qml.RZ(a[2], wires=1)
+              return qml.expval(qml.PauliZ(wires=0)), qml.probs(wires=[0, 1])
+
+            x = jax.numpy.array([0.1, 0.2, 0.3])
+
+        >>> jax.jacobian(circuit)(x)
+        (DeviceArray([-9.9833414e-02, -7.4505806e-09,  6.9285655e-10], dtype=float32),
+         DeviceArray([[-4.9419206e-02, -9.9086545e-02,  3.4938008e-09],
+                      [-4.9750542e-04,  9.9086538e-02,  1.2768372e-10],
+                      [ 4.9750548e-04,  2.4812977e-04,  4.8371929e-13],
+                      [ 4.9419202e-02, -2.4812980e-04,  2.6696912e-11]],            dtype=float32))
+
+        where before the following error was raised:
+
+        .. code-block:: python
+
+            ValueError: All input arrays must have the same shape.
+
+        The new return types system also unlocks the use of shot vectors with all the previous features. For example you
+        can take the second derivative and multiple measurement with with JAX:
+
+        .. code-block:: python
+
+            import jax
+
+            qml.enable_return()
+
+            dev = qml.device("default.qubit", wires=2, shots=(1, 10000))
+
+            params = jax.numpy.array([0.1, 0.2])
+
+            @qml.qnode(dev, interface="jax", diff_method="parameter-shift", max_diff=2)
+            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)), qml.probs(wires=[0])
+
+        >>> jax.hessian(circuit)(params)
+        ((DeviceArray([[ 0.,  0.],
+                        [ 2., -3.]], dtype=float32),
+          DeviceArray([[[-0.5,  0. ],
+                       [ 0. ,  0. ]],
+                      [[ 0.5,  0. ],
+                       [ 0. ,  0. ]]], dtype=float32)),
+         (DeviceArray([[ 0.07677898,  0.0563341 ],
+                       [ 0.07238522, -1.830669  ]], dtype=float32),
+          DeviceArray([[[-4.9707499e-01,  2.9999996e-04],
+                        [-6.2500127e-04,  1.2500001e-04]],
+                       [[ 4.9707499e-01, -2.9999996e-04],
+                        [ 6.2500127e-04, -1.2500001e-04]]], dtype=float32)))
 
-        ValueError: All input arrays must have the same shape.
     """
 
     global __activated
@@ -137,6 +317,9 @@ def disable_return():
 def active_return():
     """Function that checks if the new return types system is activated.
 
+    Returns:
+        bool: Returns ``True`` if the new return types system is activated.
+
     **Example**
 
     By default, the new return types system is turned off:
diff --git a/pennylane/tape/qscript.py b/pennylane/tape/qscript.py
index d9cc8390b..f3d19622f 100644
--- a/pennylane/tape/qscript.py
+++ b/pennylane/tape/qscript.py
@@ -19,8 +19,8 @@ executed by a device.
 
 import contextlib
 import copy
-from typing import List, Union
 from collections import Counter, defaultdict
+from typing import List, Union
 
 import pennylane as qml
 from pennylane.measurements import (
@@ -86,7 +86,6 @@ class QuantumScript:
 
     Keyword Args:
         name (str): a name given to the quantum script
-        do_queue=False (bool): Whether or not to queue. Defaults to ``False`` for ``QuantumScript``.
         _update=True (bool): Whether or not to set various properties on initialization. Setting
             ``_update=False`` reduces computations if the script is only an intermediary step.
 
@@ -146,7 +145,7 @@ class QuantumScript:
 
     >>> dev = qml.device('default.qubit', wires=(0,'a'))
     >>> qml.execute([qscript], dev, gradient_fn=None)
-    [array([0.77750694])]
+    [array([-0.77750694])]
 
     ``ops``, ``measurements``, and ``prep`` are converted to lists upon initialization,
     so those arguments accept any iterable object:
@@ -743,7 +742,8 @@ class QuantumScript:
 
             elif ret_type == qml.measurements.Sample:
 
-                shape = (len(mps), device.shots)
+                dim = mps[0].shape(device)
+                shape = (len(mps),) + dim[1:]
 
             # No other measurement type to check
 
@@ -803,24 +803,23 @@ class QuantumScript:
         return shape
 
     def shape(self, device):
-        """Produces the output shape of the tape by inspecting its measurements
+        """Produces the output shape of the quantum script by inspecting its measurements
         and the device used for execution.
 
         .. note::
 
-            The computed shape is not stored because the output shape may be
-            dependent on the device used for execution.
+            The computed shape is not stored because the output shape may be dependent on the device
+                used for execution.
 
         Args:
-            device (.Device): the device that will be used for the tape execution
+            device (.Device): the device that will be used for the script execution
 
         Raises:
-            ValueError: raised for unsupported cases for
-                example when the tape contains heterogeneous measurements
+            ValueError: raised for unsupported cases for example when the script contains
+                heterogeneous measurements
 
         Returns:
-            Union[tuple[int], list[tuple[int]]]: the output shape(s) of the
-            tape result
+            Union[tuple[int], list[tuple[int]]]: the output shape(s) of the quantum script result
 
         **Example:**
 
@@ -844,12 +843,18 @@ class QuantumScript:
                 output_shape = self._multi_homogenous_measurement_shape(self.measurements, device)
             else:
                 raise ValueError(
-                    "Getting the output shape of a tape that contains multiple types of measurements is unsupported."
+                    "Getting the output shape of a quantum script that contains multiple types of "
+                    "measurements is unsupported."
                 )
+
+        if device._shot_vector is None and self.batch_size is not None:
+            # insert the batch dimension
+            output_shape = output_shape[0:1] + (self.batch_size,) + output_shape[1:]
+
         return output_shape
 
     def _shape_new(self, device):
-        """Produces the output shape of the tape by inspecting its measurements
+        """Produces the output shape of the quantum script by inspecting its measurements
         and the device used for execution.
 
         .. note::
@@ -858,11 +863,10 @@ class QuantumScript:
             dependent on the device used for execution.
 
         Args:
-            device (.Device): the device that will be used for the tape execution
+            device (.Device): the device that will be used for the script execution
 
         Returns:
-            Union[tuple[int], tuple[tuple[int]]]: the output shape(s) of the
-            tape result
+            Union[tuple[int], tuple[tuple[int]]]: the output shape(s) of the quantum script result
 
         **Examples**
 
@@ -878,8 +882,16 @@ class QuantumScript:
             >>> qs.shape(dev)
             ((4,), (), (4,))
         """
+        if device.shot_vector is not None and self.batch_size is not None:
+            raise NotImplementedError(
+                "Parameter broadcasting when using a shot vector is not supported yet."
+            )
+
         shapes = tuple(meas_process.shape(device) for meas_process in self.measurements)
 
+        if self.batch_size is not None:
+            shapes = tuple((self.batch_size,) + shape for shape in shapes)
+
         if len(shapes) == 1:
             return shapes[0]
 
@@ -925,12 +937,12 @@ class QuantumScript:
 
     @property
     def _numeric_type_new(self):
-        """Returns the expected numeric type of the tape result by inspecting
+        """Returns the expected numeric type of the quantum script result by inspecting
         its measurements.
 
         Returns:
             Union[type, Tuple[type]]: The numeric type corresponding to the result type of the
-            tape, or a tuple of such types if the tape contains multiple measurements
+            quantum script, or a tuple of such types if the script contains multiple measurements
 
         **Example:**
 
@@ -947,7 +959,7 @@ class QuantumScript:
         return types[0] if len(types) == 1 else types
 
     # ========================================================
-    # Transforms: Tape to Tape
+    # Transforms: QuantumScript to QuantumScript
     # ========================================================
 
     def copy(self, copy_operations=False):
@@ -964,15 +976,15 @@ class QuantumScript:
         """
         if copy_operations:
             # Perform a shallow copy of all operations in the state prep, operation, and measurement
-            # queues. The operations will continue to share data with the original tape operations
+            # queues. The operations will continue to share data with the original script operations
             # unless modified.
             _prep = [copy.copy(op) for op in self._prep]
             _ops = [copy.copy(op) for op in self._ops]
             _measurements = [copy.copy(op) for op in self.measurements]
         else:
             # Perform a shallow copy of the state prep, operation, and measurement queues. The
-            # operations within the queues will be references to the original tape operations;
-            # changing the original operations will always alter the operations on the copied tape.
+            # operations within the queues will be references to the original script operations;
+            # changing the original operations will always alter the operations on the copied script.
             _prep = self._prep.copy()
             _ops = self._ops.copy()
             _measurements = self.measurements.copy()
@@ -1090,7 +1102,7 @@ class QuantumScript:
         return qml.tape.UnwrapTape(self)
 
     # ========================================================
-    # Transforms: Tape to Information
+    # Transforms: QuantumScript to Information
     # ========================================================
 
     @property
diff --git a/pennylane/transforms/transpile.py b/pennylane/transforms/transpile.py
index 95ddbfc95..7daf80252 100644
--- a/pennylane/transforms/transpile.py
+++ b/pennylane/transforms/transpile.py
@@ -69,7 +69,7 @@ def transpile(tape: QuantumTape, coupling_map: Union[List, nx.Graph]):
 
     >>> dev = qml.device('default.qubit', wires=[0, 1, 2, 3])
     >>> transpiled_circuit = qml.transforms.transpile(coupling_map=[(0, 1), (1, 3), (3, 2), (2, 0)])(circuit)
-    >>> transpiled_qnode = qml.QNode(circuit, dev)
+    >>> transpiled_qnode = qml.QNode(transpiled_circuit, dev)
     >>> print(qml.draw(transpiled_qnode)())
     0: ─╭●────────────────╭●─┤ ╭Probs
     1: ─╰X─╭●───────╭●────│──┤ ├Probs


Test patch:
diff --git a/tests/data/test_dataset.py b/tests/data/test_dataset.py
index 1e2ef3f0b..1b893dedd 100644
--- a/tests/data/test_dataset.py
+++ b/tests/data/test_dataset.py
@@ -209,3 +209,33 @@ def test_import_zstd_dill(monkeypatch):
 
         with pytest.raises(ImportError, match="This feature requires zstd and dill"):
             qml.data.dataset._import_zstd_dill()
+
+
+def test_repr_standard(tmp_path):
+    """Test that __repr__ for standard Datasets look as expected."""
+    folder = tmp_path / "qchem" / "H2" / "STO-3G" / "1.02"
+    os.makedirs(folder)
+    qml.data.Dataset._write_file(
+        {"molecule": 1, "hf_state": 2}, str(folder / "H2_STO-3G_1.02_full.dat")
+    )
+
+    dataset = qml.data.Dataset("qchem", str(folder), "H2_STO-3G_1.02", "", standard=True)
+    assert (
+        repr(dataset)
+        == "<Dataset = description: qchem/H2/STO-3G/1.02, attributes: ['molecule', 'hf_state']>"
+    )
+
+    dataset.vqe_energy = 1.1
+    assert (
+        repr(dataset)
+        == "<Dataset = description: qchem/H2/STO-3G/1.02, attributes: ['molecule', 'hf_state', ...]>"
+    )
+
+
+def test_repr_non_standard():
+    """Test that __repr__ for non-standard Datasets look as expected."""
+    dataset = qml.data.Dataset(foo=1, bar=2)
+    assert repr(dataset) == "<Dataset = attributes: ['foo', 'bar']>"
+
+    dataset.baz = 3
+    assert repr(dataset) == "<Dataset = attributes: ['foo', 'bar', ...]>"
diff --git a/tests/interfaces/test_jax_qnode.py b/tests/interfaces/test_jax_qnode.py
index 8a9f9920a..94c2e43d8 100644
--- a/tests/interfaces/test_jax_qnode.py
+++ b/tests/interfaces/test_jax_qnode.py
@@ -1121,6 +1121,31 @@ class TestQubitIntegration:
             # to the first parameter of circuit1.
             assert circuit1.qtape.trainable_params == [1, 2]
 
+    def test_param_broadcasting_fwd_pass(self, dev_name, diff_method, mode, interface, tol):
+        """Test that parameter-broadcasting works well with JAX on the forward pass."""
+        if mode == "forward":
+            pytest.skip(
+                "This test doesn't compute Jacobians - forward mode executions involves Jacobian computation."
+            )
+
+        num_layers = 2
+        num_qubits = 3
+        weight_shapes = (2 + num_layers, num_qubits)
+
+        dev = qml.device(dev_name, wires=num_qubits)
+
+        @qml.qnode(dev, interface=interface, diff_method=diff_method, mode=mode)
+        def circuit(params):
+            qml.AngleEmbedding(params[0, :], wires=range(num_qubits), rotation="X")
+            qml.BasicEntanglerLayers(
+                weights=params[1:-1, :], rotation=qml.RY, wires=range(num_qubits)
+            )
+            qml.AngleEmbedding(params[-1, :], wires=range(num_qubits), rotation="Y")
+            return qml.expval(qml.PauliZ(0))
+
+        params = np.random.random((5, 3, 3))
+        assert circuit(params).shape == (3,)
+
 
 @pytest.mark.parametrize(
     "diff_method,kwargs",
@@ -1581,3 +1606,29 @@ class TestJIT:
 
         assert np.allclose(g0, expected_g[0][idx], atol=tol, rtol=0)
         assert np.allclose(g1, expected_g[1][idx], atol=tol, rtol=0)
+
+    def test_param_broadcasting_fwd_pass(self, dev_name, diff_method, mode, tol):
+        """Test that parameter-broadcasting works well with JAX JIT on the forward pass."""
+        if mode == "forward":
+            pytest.skip(
+                "This test doesn't compute Jacobians - forward mode executions involves Jacobian computation."
+            )
+
+        num_layers = 2
+        num_qubits = 3
+        weight_shapes = (2 + num_layers, num_qubits)
+
+        dev = qml.device(dev_name, wires=num_qubits)
+
+        @jax.jit
+        @qnode(dev, diff_method=diff_method, interface="jax", mode=mode)
+        def circuit(params):
+            qml.AngleEmbedding(params[0, :], wires=range(num_qubits), rotation="X")
+            qml.BasicEntanglerLayers(
+                weights=params[1:-1, :], rotation=qml.RY, wires=range(num_qubits)
+            )
+            qml.AngleEmbedding(params[-1, :], wires=range(num_qubits), rotation="Y")
+            return qml.expval(qml.PauliZ(0))
+
+        params = np.random.random((5, 3, 3))
+        assert circuit(params).shape == (3,)
diff --git a/tests/returntypes/test_new_qscript.py b/tests/returntypes/test_new_qscript.py
index 77513ff39..46c8b51c1 100644
--- a/tests/returntypes/test_new_qscript.py
+++ b/tests/returntypes/test_new_qscript.py
@@ -327,6 +327,62 @@ class TestOutputShape:
 
         assert res_shape == expected
 
+    @pytest.mark.autograd
+    @pytest.mark.parametrize("measurement, expected_shape", measures)
+    @pytest.mark.parametrize("shots", [None, 1, 10])
+    def test_broadcasting_single(self, measurement, expected_shape, shots):
+        """Test that the output shape produced by the tape matches the expected
+        output shape for a single measurement and parameter broadcasting"""
+        if shots is None and measurement.return_type is qml.measurements.Sample:
+            pytest.skip("Sample doesn't support analytic computations.")
+
+        if measurement.return_type is qml.measurements.State and measurement.wires is not None:
+            pytest.skip("Density matrix does not support parameter broadcasting")
+
+        num_wires = 3
+        dev = qml.device("default.qubit", wires=num_wires, shots=shots)
+
+        a = np.array([0.1, 0.2, 0.3])
+        b = np.array([0.4, 0.5, 0.6])
+
+        with qml.tape.QuantumTape() as tape:
+            qml.RY(a, wires=0)
+            qml.RX(b, wires=0)
+            qml.apply(measurement)
+
+        expected_shape = qml.execute([tape], dev, gradient_fn=None)[0].shape
+
+        assert tape.shape(dev) == expected_shape
+
+    @pytest.mark.autograd
+    @pytest.mark.parametrize("measurement, expected", measures)
+    @pytest.mark.parametrize("shots", [None, 1, 10])
+    def test_broadcasting_multi(self, measurement, expected, shots):
+        """Test that the output shape produced by the tape matches the expected
+        output shape for multiple measurements and parameter broadcasting"""
+        if shots is None and measurement.return_type is qml.measurements.Sample:
+            pytest.skip("Sample doesn't support analytic computations.")
+
+        if measurement.return_type is qml.measurements.State:
+            pytest.skip("State does not support multiple measurements")
+
+        dev = qml.device("default.qubit", wires=3, shots=shots)
+
+        a = np.array([0.1, 0.2, 0.3])
+        b = np.array([0.4, 0.5, 0.6])
+
+        with qml.tape.QuantumTape() as tape:
+            qml.RY(a, wires=0)
+            qml.RX(b, wires=0)
+            for _ in range(2):
+                qml.apply(measurement)
+
+        expected = qml.execute([tape], dev, gradient_fn=None)[0]
+        actual = tape.shape(dev)
+
+        for exp, act in zip(expected, actual):
+            assert exp.shape == act
+
     @pytest.mark.autograd
     def test_multi_measure_sample_obs_shot_vector(self):
         """Test that the expected output shape is obtained when using multiple
@@ -374,6 +430,21 @@ class TestOutputShape:
 
         assert res == expected_shape
 
+    def test_raises_broadcasting_shot_vector(self):
+        """Test that getting the output shape of a tape that uses parameter
+        broadcasting along with a device with a shot vector raises an error."""
+        dev = qml.device("default.qubit", wires=3, shots=(1, 2, 3))
+
+        with qml.tape.QuantumTape() as tape:
+            qml.RY(np.array([0.1, 0.2]), wires=0)
+            qml.RX(np.array([0.3, 0.4]), wires=0)
+            qml.expval(qml.PauliZ(0))
+
+        msg = "Parameter broadcasting when using a shot vector is not supported yet"
+
+        with pytest.raises(NotImplementedError, match=msg):
+            tape.shape(dev)
+
 
 class TestNumericType:
     """Tests for determining the numeric type of the tape output."""
diff --git a/tests/returntypes/test_parameter_shift_new.py b/tests/returntypes/test_parameter_shift_new.py
index 2c827a306..42e1ef742 100644
--- a/tests/returntypes/test_parameter_shift_new.py
+++ b/tests/returntypes/test_parameter_shift_new.py
@@ -66,9 +66,11 @@ class TestGetOperationRecipe:
             qml.RX(x, 0)
             qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliX(1))
 
-        # This is not a correct gradient recipe, but that's not the issue
+        # Incorrect gradient recipe, but this test only checks execution with an unshifted term.
         recipes = ([[-1e7, 1, 0], [1e7, 1, 1e7]],)
         tapes, fn = qml.gradients.param_shift(tape, gradient_recipes=recipes)
+        assert len(tapes) == 2
+
         res = fn(qml.execute(tapes, dev, None))
         assert len(res) == 2
         assert isinstance(res, tuple)
diff --git a/tests/tape/test_tape.py b/tests/tape/test_tape.py
index 9ddf52d14..42e607298 100644
--- a/tests/tape/test_tape.py
+++ b/tests/tape/test_tape.py
@@ -2229,6 +2229,60 @@ class TestOutputShape:
         for r, e in zip(res, execution_results):
             assert r == e.shape
 
+    @pytest.mark.autograd
+    @pytest.mark.parametrize("measurement, expected_shape", measures)
+    @pytest.mark.parametrize("shots", [None, 1, 10])
+    def test_broadcasting_single(self, measurement, expected_shape, shots):
+        """Test that the output shape produced by the tape matches the expected
+        output shape for a single measurement and parameter broadcasting"""
+        if shots is None and measurement.return_type is qml.measurements.Sample:
+            pytest.skip("Sample doesn't support analytic computations.")
+
+        if measurement.return_type is qml.measurements.State and measurement.wires is not None:
+            pytest.skip("Density matrix does not support parameter broadcasting")
+
+        num_wires = 3
+        dev = qml.device("default.qubit", wires=num_wires, shots=shots)
+
+        a = np.array([0.1, 0.2, 0.3])
+        b = np.array([0.4, 0.5, 0.6])
+
+        with qml.tape.QuantumTape() as tape:
+            qml.RY(a, wires=0)
+            qml.RX(b, wires=0)
+            qml.apply(measurement)
+
+        expected_shape = qml.execute([tape], dev, gradient_fn=None)[0].shape
+
+        assert tape.shape(dev) == expected_shape
+
+    @pytest.mark.autograd
+    @pytest.mark.parametrize("measurement, expected", measures)
+    @pytest.mark.parametrize("shots", [None, 1, 10])
+    def test_broadcasting_multi(self, measurement, expected, shots):
+        """Test that the output shape produced by the tape matches the expected
+        output shape for multiple measurements and parameter broadcasting"""
+        if shots is None and measurement.return_type is qml.measurements.Sample:
+            pytest.skip("Sample doesn't support analytic computations.")
+
+        if measurement.return_type is qml.measurements.State:
+            pytest.skip("State does not support multiple measurements")
+
+        dev = qml.device("default.qubit", wires=3, shots=shots)
+
+        a = np.array([0.1, 0.2, 0.3])
+        b = np.array([0.4, 0.5, 0.6])
+
+        with qml.tape.QuantumTape() as tape:
+            qml.RY(a, wires=0)
+            qml.RX(b, wires=0)
+            for _ in range(2):
+                qml.apply(measurement)
+
+        expected = qml.execute([tape], dev, gradient_fn=None)[0].shape
+
+        assert tape.shape(dev) == expected
+
     def test_multi_measure_probs_shot_vector_errors(self):
         """Test that getting the output shape of a tape containing multiple
         probability measurements with different number of wires errors when