Repository: PennyLaneAI/pennylane
Pull request #2026 (2021-12-14 14:04:01)

Merge rc into `master`

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.
antalszava: [sc-9643]
codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2026?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#2026](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2026?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (3ce3ab1) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/4f385b82a464a1717f2454aec236f510f615fcc7?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (4f385b8) will **decrease** coverage by `2.14%`.
> The diff coverage is `44.59%`.

[![Impacted file tree graph](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2026/graphs/tree.svg?width=650&height=150&src=pr&token=9S3YzlUgGk&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2026?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)

```diff
@@            Coverage Diff             @@
##           master    #2026      +/-   ##
==========================================
- Coverage   99.12%   96.97%   -2.15%     
==========================================
  Files         221      221              
  Lines       16792    16852      +60     
==========================================
- Hits        16645    16343     -302     
- Misses        147      509     +362     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2026?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/devices/default\_qubit\_jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2026/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvZGVmYXVsdF9xdWJpdF9qYXgucHk=) | `0.00% <0.00%> (-95.00%)` | :arrow_down: |
| [pennylane/drawer/style.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2026/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RyYXdlci9zdHlsZS5weQ==) | `100.00% <ø> (ø)` | |
| [pennylane/fourier/reconstruct.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2026/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ZvdXJpZXIvcmVjb25zdHJ1Y3QucHk=) | `100.00% <ø> (ø)` | |
| [pennylane/gradients/general\_shift\_rules.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2026/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2dyYWRpZW50cy9nZW5lcmFsX3NoaWZ0X3J1bGVzLnB5) | `100.00% <ø> (ø)` | |
| [pennylane/hf/molecule.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2026/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2hmL21vbGVjdWxlLnB5) | `100.00% <ø> (ø)` | |
| [pennylane/interfaces/batch/jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2026/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvYmF0Y2gvamF4LnB5) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| [pennylane/measure.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2026/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21lYXN1cmUucHk=) | `99.03% <ø> (ø)` | |
| [pennylane/ops/qubit/state\_preparation.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2026/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9xdWJpdC9zdGF0ZV9wcmVwYXJhdGlvbi5weQ==) | `100.00% <ø> (ø)` | |
| [...ylane/templates/subroutines/commuting\_evolution.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2026/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RlbXBsYXRlcy9zdWJyb3V0aW5lcy9jb21tdXRpbmdfZXZvbHV0aW9uLnB5) | `100.00% <ø> (ø)` | |
| [pennylane/transforms/metric\_tensor.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2026/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvbWV0cmljX3RlbnNvci5weQ==) | `99.48% <ø> (ø)` | |
| ... and [27 more](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2026/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2026?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2026?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Last update [4f385b8...3ce3ab1](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2026?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

antalszava: Note: support for `qml.probs` with the `default.qubit.jax` device was disallowed when providing a shot vector because the logic had a flaw in there.

Patch:
diff --git a/doc/code/qml_operation.rst b/doc/code/qml_operation.rst
index 6004e87e9..c98925707 100644
--- a/doc/code/qml_operation.rst
+++ b/doc/code/qml_operation.rst
@@ -33,6 +33,7 @@ new :class:`~pennylane.ops.qubit.attributes.Attribute` objects.
 
     ~ops.qubit.attributes.Attribute
     ~ops.qubit.attributes.composable_rotations
+    ~ops.qubit.attributes.diagonal_in_z_basis
     ~ops.qubit.attributes.has_unitary_generator
     ~ops.qubit.attributes.self_inverses
     ~ops.qubit.attributes.symmetric_over_all_wires
diff --git a/doc/development/release_notes.md b/doc/development/release_notes.md
index 1cc0f5b64..0ebe8a142 100644
--- a/doc/development/release_notes.md
+++ b/doc/development/release_notes.md
@@ -6,6 +6,8 @@ This page contains the release notes for PennyLane.
 
 .. mdinclude:: ../releases/changelog-dev.md
 
+.. mdinclude:: ../releases/changelog-0.20.0.md
+
 .. mdinclude:: ../releases/changelog-0.19.1.md
 
 .. mdinclude:: ../releases/changelog-0.19.0.md
diff --git a/doc/introduction/measurements.rst b/doc/introduction/measurements.rst
index 8a5465683..3fa402e7c 100644
--- a/doc/introduction/measurements.rst
+++ b/doc/introduction/measurements.rst
@@ -129,10 +129,9 @@ Probability
 -----------
 
 You can also train QNodes on computational basis probabilities, by using
-the :func:`~.pennylane.probs` measurement function. Unlike other
-measurement functions, **this does not accept observables**.
-Instead, it will return a flat array or tensor containing the (marginal)
-probabilities of each quantum state.
+the :func:`~.pennylane.probs` measurement function. The function can
+accept either specified ``wires`` or an observable that rotates the 
+computational basis.
 
 .. code-block:: python3
 
diff --git a/doc/introduction/operations.rst b/doc/introduction/operations.rst
index 7afe87690..58a8d635b 100644
--- a/doc/introduction/operations.rst
+++ b/doc/introduction/operations.rst
@@ -170,6 +170,7 @@ Qubit state preparation
 
     ~pennylane.BasisState
     ~pennylane.QubitStateVector
+    ~pennylane.QubitDensityMatrix
 
 :html:`</div>`
 
diff --git a/doc/introduction/templates.rst b/doc/introduction/templates.rst
index c0b665b9e..6e02126e6 100644
--- a/doc/introduction/templates.rst
+++ b/doc/introduction/templates.rst
@@ -229,7 +229,7 @@ Other useful templates which do not belong to the previous categories can be fou
     :figure: ../_static/templates/subroutines/qft.svg
 
 .. customgalleryitem::
-    :link: :link: ../code/api/pennylane.CommutingEvolution.html
+    :link: ../code/api/pennylane.CommutingEvolution.html
     :description: CommutingEvolution
     :figure: ../_static/templates/subroutines/commuting_evolution.png
 
diff --git a/doc/releases/changelog-0.19.1.md b/doc/releases/changelog-0.19.1.md
index c48bf5415..e64abbd9e 100644
--- a/doc/releases/changelog-0.19.1.md
+++ b/doc/releases/changelog-0.19.1.md
@@ -1,6 +1,6 @@
 :orphan:
 
-# Release 0.19.1 (current release)
+# Release 0.19.1
 
 <h3>Bug fixes</h3>
 
diff --git a/doc/releases/changelog-0.20.0.md b/doc/releases/changelog-0.20.0.md
index 8ed8a0d0f..bf47bc920 100644
--- a/doc/releases/changelog-0.20.0.md
+++ b/doc/releases/changelog-0.20.0.md
@@ -4,46 +4,75 @@
 
 <h3>New features since last release</h3>
 
-* Added `qml.LieAlgebraOptimizer`, a new quantum-aware Lie Algebra optimizer that allows
-  one to perform gradient descent on the special unitary group.
-  [(#1911)](https://github.com/PennyLaneAI/pennylane/pull/1911)
-  
-* Added functions for computing the values of atomic and molecular orbitals at a given position.
-  [(#1867)](https://github.com/PennyLaneAI/pennylane/pull/1867)
+<h4>Shiny new circuit drawer!🎨🖌️ </h4>
 
-  The functions `atomic_orbital` and `molecular_orbital` can be used, as shown in the
-  following codeblock, to evaluate the orbitals. By generating values of the orbitals at different
-  positions, one can plot the spatial shape of a desired orbital.
+* PennyLane now supports drawing a QNode with matplotlib!
+  [(#1803)](https://github.com/PennyLaneAI/pennylane/pull/1803)
+  [(#1811)](https://github.com/PennyLaneAI/pennylane/pull/1811)
+  [(#1931)](https://github.com/PennyLaneAI/pennylane/pull/1931)
+  [(#1954)](https://github.com/PennyLaneAI/pennylane/pull/1954)
 
   ```python
-  symbols  = ['H', 'H']
-  geometry = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]], requires_grad = False) 
-  mol = hf.Molecule(symbols, geometry)
-  hf.generate_scf(mol)()
+  dev = qml.device("default.qubit", wires=4)
 
-  ao = mol.atomic_orbital(0)
-  mo = mol.molecular_orbital(1)
+  @qml.qnode(dev)
+  def circuit(x, z):
+      qml.QFT(wires=(0,1,2,3))
+      qml.Toffoli(wires=(0,1,2))
+      qml.CSWAP(wires=(0,2,3))
+      qml.RX(x, wires=0)
+      qml.CRZ(z, wires=(3,0))
+      return qml.expval(qml.PauliZ(0))
+  fig, ax = qml.draw_mpl(circuit)(1.2345, 1.2345)
+  fig.show()
+  ```
+
+  <img src="https://pennylane.readthedocs.io/en/latest/_images/main_example.png" width=70%/>
+
+<h4>New and improved quantum-aware optimizers</h4>
+
+* Added `qml.LieAlgebraOptimizer`, a new quantum-aware Lie Algebra optimizer
+  that allows one to perform gradient descent on the special unitary group.
+  [(#1911)](https://github.com/PennyLaneAI/pennylane/pull/1911)
+
+  ```python
+  dev = qml.device("default.qubit", wires=2)
+  H = -1.0 * qml.PauliX(0) - qml.PauliZ(1) - qml.PauliY(0) @ qml.PauliX(1)
+
+  @qml.qnode(dev)
+  def circuit():
+      qml.RX(0.1, wires=[0])
+      qml.RY(0.5, wires=[1])
+      qml.CNOT(wires=[0,1])
+      qml.RY(0.6, wires=[0])
+      return qml.expval(H)
+  opt = qml.LieAlgebraOptimizer(circuit=circuit, stepsize=0.1)
   ```
 
+  Note that, unlike other optimizers, the `LieAlgebraOptimizer` accepts a QNode
+  with *no* parameters, and instead grows the circuit by appending operations
+  during the optimization:
+
   ```pycon
-  >>> print(ao(0.0, 0.0, 0.0))
-  >>> print(mo(0.0, 0.0, 0.0))
-  0.6282468778183719
-  0.018251285973461928
+  >>> circuit()
+  tensor(-1.3351865, requires_grad=True)
+  >>> circuit1, cost = opt.step_and_cost()
+  >>> circuit1()
+  tensor(-1.99378872, requires_grad=True)
   ```
 
-* The `metric_tensor` transform can now be used to compute the full
-  tensor, beyond the block diagonal approximation. 
+  For more details, see the
+  [LieAlgebraOptimizer documentation](https://pennylane.readthedocs.io/en/stable/code/api/pennylane.LieAlgebraOptimizer.html).
+
+* The `qml.metric_tensor` transform can now be used to compute the full
+  tensor, beyond the block diagonal approximation.
   [(#1725)](https://github.com/PennyLaneAI/pennylane/pull/1725)
 
-  This is performed using Hadamard tests, and requires an additional wire 
-  on the device to execute the circuits produced by the transform, 
+  This is performed using Hadamard tests, and requires an additional wire
+  on the device to execute the circuits produced by the transform,
   as compared to the number of wires required by the original circuit.
   The transform defaults to computing the full tensor, which can
   be controlled by the `approx` keyword argument.
-  See the 
-  [qml.metric_tensor docstring](https://pennylane.readthedocs.io/en/latest/code/api/pennylane.transforms.metric_tensor.html).
-  for more information and usage details.
 
   As an example, consider the QNode
 
@@ -62,7 +91,7 @@
   ```
 
   Then we can compute the (block) diagonal metric tensor as before, now using the
-  ``approx="block-diag"`` keyword:
+  `approx="block-diag"` keyword:
 
   ```pycon
   >>> qml.metric_tensor(circuit, approx="block-diag")(weights)
@@ -81,6 +110,40 @@
    [-0.23300977  0.01763859  0.21846983]]
   ```
 
+  See the
+  [metric tensor documentation](https://pennylane.readthedocs.io/en/stable/code/api/pennylane.transforms.metric_tensor.html).
+  for more information and usage details.
+
+<h4>Faster performance with optimized quantum workflows</h4>
+
+* The QNode has been re-written to support batch execution across the board,
+  custom gradients, better decomposition strategies, and higher-order derivatives.
+  [(#1807)](https://github.com/PennyLaneAI/pennylane/pull/1807)
+  [(#1969)](https://github.com/PennyLaneAI/pennylane/pull/1969)
+
+  - Internally, if multiple circuits are generated for simultaneous execution, they
+    will be packaged into a single job for execution on the device. This can lead to
+    significant performance improvement when executing the QNode on remote
+    quantum hardware or simulator devices with parallelization capabilities.
+
+  - Custom gradient transforms can be specified as the differentiation method:
+
+    ```python
+    @qml.gradients.gradient_transform
+    def my_gradient_transform(tape):
+        ...
+        return tapes, processing_fn
+
+    @qml.qnode(dev, diff_method=my_gradient_transform)
+    def circuit():
+    ```
+
+  For breaking changes related to the use of the new QNode, refer to the
+  Breaking Changes section.
+
+  Note that the old QNode remains accessible at `@qml.qnode_old.qnode`, however this will
+  be removed in the next release.
+
 * Custom decompositions can now be applied to operations at the device level.
   [(#1900)](https://github.com/PennyLaneAI/pennylane/pull/1900)
 
@@ -150,28 +213,51 @@
    2: ──RX(0.6)──RZ(3.14)──RY(1.57)──────────────────────────╰Z──RZ(3.14)──RY(1.57)──╰C──────────────────────┤
   ```
 
-* PennyLane now supports drawing a QNode with matplotlib!
-  [(#1803)](https://github.com/PennyLaneAI/pennylane/pull/1803)
-  [(#1811)](https://github.com/PennyLaneAI/pennylane/pull/1811)
-  [(#1954)](https://github.com/PennyLaneAI/pennylane/pull/1954)
+* Given an operator of the form :math:`U=e^{iHt}`, where :math:`H` has
+  commuting terms and known eigenvalues,
+  `qml.gradients.generate_shift_rule` computes the generalized parameter shift rules for determining
+  the gradient of the expectation value :math:`f(t) = \langle 0|U(t)^\dagger \hat{O} U(t)|0\rangle` on
+  hardware.
+  [(#1788)](https://github.com/PennyLaneAI/pennylane/pull/1788)
+  [(#1932)](https://github.com/PennyLaneAI/pennylane/pull/1932)
 
-  ```python
-  @qml.qnode(qml.device("default.qubit", wires=4))
-  def circuit(x, z):
-      qml.QFT(wires=(0,1,2,3))
-      qml.IsingXX(1.234, wires=(0,2))
-      qml.Toffoli(wires=(0,1,2))
-      qml.CSWAP(wires=(0,2,3))
-      qml.RX(x, wires=0)
-      qml.CRZ(z, wires=(3,0))
-      return qml.expval(qml.PauliZ(0))
+  Given
 
-  qml.drawer.use_style('black_white')
-  fig, ax = qml.draw_mpl(circuit)(1.2345, 1.2345)
-  fig.show()
+  .. math:: H = \sum_i a_i h_i,
+
+  where the eigenvalues of :math:`H` are known and all :math:`h_i` commute, we can compute
+  the *frequencies* (the unique positive differences of any two eigenvalues) using
+  `qml.gradients.eigvals_to_frequencies`.
+
+  `qml.gradients.generate_shift_rule` can then be used to compute the parameter
+  shift rules to compute :math:`f'(t)` using `2R` shifted cost function evaluations.
+  This becomes cheaper than the standard application of the chain rule and
+  two-term shift rule when `R` is less than the
+  number of Pauli words in the generator.
+
+  For example, consider the case where :math:`H` has eigenspectrum `(-1, 0, 1)`:
+
+  ```pycon
+  >>> frequencies = qml.gradients.eigvals_to_frequencies((-1, 0, 1))
+  >>> frequencies
+  (1, 2)
+  >>> coeffs, shifts = qml.gradients.generate_shift_rule(frequencies)
+  >>> coeffs
+  array([ 0.85355339, -0.85355339, -0.14644661,  0.14644661])
+  >>> shifts
+  array([ 0.78539816, -0.78539816,  2.35619449, -2.35619449])
   ```
 
-  <img src="https://pennylane.readthedocs.io/en/latest/_static/draw_mpl/black_white_style.png" width=70%/>
+  As we can see, `generate_shift_rule` returns four coefficients :math:`c_i` and shifts
+  :math:`s_i` corresponding to a four term parameter shift rule. The gradient can then
+  be reconstructed via:
+
+  .. math:: \frac{\partial}{\partial\phi}f = \sum_{i} c_i f(\phi + s_i),
+
+  where :math:`f(\phi) = \langle 0|U(\phi)^\dagger \hat{O} U(\phi)|0\rangle`
+  for some observable :math:`\hat{O}` and the unitary :math:`U(\phi)=e^{iH\phi}`.
+
+<h4>Support for TensorFlow AutoGraph mode with quantum hardware</h4>
 
 * It is now possible to use TensorFlow's [AutoGraph
   mode](https://www.tensorflow.org/guide/function) with QNodes on all devices and with arbitrary
@@ -185,7 +271,7 @@
   ```python
   dev = qml.device("lightning.qubit", wires=2)
 
-  @qml.beta.qnode(dev, diff_method="adjoint", interface="tf", max_diff=1)
+  @qml.qnode(dev, diff_method="adjoint", interface="tf", max_diff=1)
   def circuit(x):
       qml.RX(x[0], wires=0)
       qml.RY(x[1], wires=1)
@@ -213,29 +299,7 @@
   For more details, please see the [TensorFlow AutoGraph
   documentation](https://www.tensorflow.org/guide/function).
 
-
-* `qml.math.scatter_element_add` now supports adding multiple values at
-  multiple indices with a single function call, in all interfaces
-  [(#1864)](https://github.com/PennyLaneAI/pennylane/pull/1864)
-
-  For example, we may set five values of a three-dimensional tensor
-  in the following way:
-
-  ```pycon
-  >>> X = tf.zeros((3, 2, 9), dtype=tf.float64)
-  >>> indices = [(0, 0, 1, 2, 2), (0, 0, 0, 0, 1), (1, 3, 8, 6, 7)]
-  >>> values = [0.1 * i for i in range(5)]
-  >>> qml.math.scatter_element_add(X, indices, values)
-  <tf.Tensor: shape=(3, 2, 9), dtype=float64, numpy=
-  array([[[0., 1., 0., 2., 0., 0., 0., 0., 0.],
-          [0., 0., 0., 0., 0., 0., 0., 0., 0.]],
-
-         [[0., 0., 0., 0., 0., 0., 0., 0., 3.],
-          [0., 0., 0., 0., 0., 0., 0., 0., 0.]],
-
-         [[0., 0., 0., 0., 0., 0., 4., 0., 0.],
-          [0., 0., 0., 0., 0., 0., 0., 5., 0.]]])>
-  ```
+<h4>Characterize your quantum models with classical QNode reconstruction</h4>
 
 * The `qml.fourier.reconstruct` function is added. It can be used to
   reconstruct QNodes outputting expectation values along a specified
@@ -287,59 +351,9 @@
   ```
 
   For more details on usage, reconstruction cost and differentiability support, please see the
-  [fourier.reconstruct docstring](https://pennylane.readthedocs.io/en/latest/code/api/pennylane.fourier.reconstruct.html).
-
-* A thermal relaxation channel is added to the Noisy channels. The channel description can be
-  found on the supplementary information of [Quantum classifier with tailored quantum kernels](https://arxiv.org/abs/1909.02611).
-  [(#1766)](https://github.com/PennyLaneAI/pennylane/pull/1766)
-
-* Added the identity observable to be an operator. Now we can explicitly call the identity
-  operation on our quantum circuits for both qubit and CV devices.
-  [(#1829)](https://github.com/PennyLaneAI/pennylane/pull/1829)
-
-* Given an operator of the form :math:`U=e^{iHt}`, where :math:`H` has
-  commuting terms and known eigenvalues,
-  `qml.gradients.generate_shift_rule` computes the generalized parameter shift rules for determining
-  the gradient of the expectation value :math:`f(t) = \langle 0|U(t)^\dagger \hat{O} U(t)|0\rangle` on
-  hardware.
-  [(#1788)](https://github.com/PennyLaneAI/pennylane/pull/1788)
-  [(#1932)](https://github.com/PennyLaneAI/pennylane/pull/1932)
-
-  Given
-
-  .. math:: H = \sum_i a_i h_i,
-
-  where the eigenvalues of :math:`H` are known and all :math:`h_i` commute, we can compute
-  the *frequencies* (the unique positive differences of any two eigenvalues) using
-  `qml.gradients.eigvals_to_frequencies`.
-
-  `qml.gradients.generate_shift_rule` can then be used to compute the parameter
-  shift rules to compute :math:`f'(t)` using `2R` shifted cost function evaluations.
-  This becomes cheaper than the standard application of the chain rule and
-  two-term shift rule when `R` is less than the
-  number of Pauli words in the generator.
-
-  For example, consider the case where :math:`H` has eigenspectrum ``(-1, 0, 1)``:
-
-  ```pycon
-  >>> frequencies = qml.gradients.eigvals_to_frequencies((-1, 0, 1))
-  >>> frequencies
-  (1, 2)
-  >>> coeffs, shifts = qml.gradients.generate_shift_rule(frequencies)
-  >>> coeffs
-  array([ 0.85355339, -0.85355339, -0.14644661,  0.14644661])
-  >>> shifts
-  array([ 0.78539816, -0.78539816,  2.35619449, -2.35619449])
-  ```
-
-  As we can see, `generate_shift_rule` returns four coefficients :math:`c_i` and shifts
-  :math:`s_i` corresponding to a four term parameter shift rule. The gradient can then
-  be reconstructed via:
-
-  .. math:: \frac{\partial}{\partial\phi}f = \sum_{i} c_i f(\phi + s_i),
+  [fourier.reconstruct docstring](https://pennylane.readthedocs.io/en/stable/code/api/pennylane.fourier.reconstruct.html).
 
-  where :math:`f(\phi) = \langle 0|U(\phi)^\dagger \hat{O} U(\phi)|0\rangle`
-  for some observable :math:`\hat{O}` and the unitary :math:`U(\phi)=e^{iH\phi}`.
+<h4>State-of-the-art operations and templates</h4>
 
 * A circuit template for time evolution under a commuting Hamiltonian utilizing generalized
   parameter shift rules for cost function gradients is available as `qml.CommutingEvolution`.
@@ -350,7 +364,7 @@
   `CommutingEvolution`'s :math:`t` parameter, otherwise the shift rule for a decomposition of
   `CommutingEvolution` will be used.
 
-  The template can be initialized within a `qnode` as:
+  The template can be initialized within QNode as:
 
   ```python
   import pennylane as qml
@@ -361,7 +375,7 @@
   coeffs = [1, -1]
   obs = [qml.PauliX(0) @ qml.PauliY(1), qml.PauliY(0) @ qml.PauliX(1)]
   hamiltonian = qml.Hamiltonian(coeffs, obs)
-  frequencies = [2,4]
+  frequencies = (2,4)
 
   @qml.qnode(dev)
   def circuit(time):
@@ -374,67 +388,145 @@
   and 2) the eigenvalue frequency spectrum is correct, since these checks become
   prohibitively expensive for large Hamiltonians.
 
-* The qml.Barrier() operator has been added. With it we can separate blocks in compilation or use it as a visual tool.
+* The `qml.Barrier()` operator has been added. With it we can separate blocks
+  in compilation or use it as a visual tool.
   [(#1844)](https://github.com/PennyLaneAI/pennylane/pull/1844)
 
-* Added density matrix initialization gate for mixed state simulation. [(#1686)](https://github.com/PennyLaneAI/pennylane/issues/1686)
+* Added the identity observable to be an operator. Now we can explicitly call the identity
+  operation on our quantum circuits for both qubit and CV devices.
+  [(#1829)](https://github.com/PennyLaneAI/pennylane/pull/1829)
+
+* Added the `qml.QubitDensityMatrix` initialization gate for
+  mixed state simulation.
+  [(#1850)](https://github.com/PennyLaneAI/pennylane/pull/1850)
+
+* A thermal relaxation channel is added to the Noisy channels. The channel description can be
+  found on the supplementary information of [Quantum classifier with tailored quantum kernels](https://arxiv.org/abs/1909.02611).
+  [(#1766)](https://github.com/PennyLaneAI/pennylane/pull/1766)
+
+* Added a new `qml.PauliError` channel that allows the application of an
+  arbitrary number of Pauli operators on an arbitrary number of wires.
+  [(#1781)](https://github.com/PennyLaneAI/pennylane/pull/1781)
+
+<h4>Manipulate QNodes to your ❤️s content with new transforms</h4>
 
-* The `merge_amplitude_embedding` transformation has been created to automatically merge all gates of this type into one.
+* The `merge_amplitude_embedding` transformation has been created to
+  automatically merge all gates of this type into one.
   [(#1933)](https://github.com/PennyLaneAI/pennylane/pull/1933)
 
-* The `undo_swaps` transformation has been created to automatically remove all swaps of a circuit.
+  ```python
+  from pennylane.transforms import merge_amplitude_embedding
+
+  dev = qml.device("default.qubit", wires = 3)
+  
+  @qml.qnode(dev)
+  @merge_amplitude_embedding
+  def qfunc():
+      qml.AmplitudeEmbedding([0,1,0,0], wires = [0,1])
+      qml.AmplitudeEmbedding([0,1], wires = 2)
+      return qml.expval(qml.PauliZ(wires = 0))
+  ```
+
+  ```pycon
+  >>> print(qml.draw(qnode)())
+   0: ──╭AmplitudeEmbedding(M0)──┤ ⟨Z⟩
+   1: ──├AmplitudeEmbedding(M0)──┤
+   2: ──╰AmplitudeEmbedding(M0)──┤
+   M0 =
+   [0.+0.j 0.+0.j 0.+0.j 1.+0.j 0.+0.j 0.+0.j 0.+0.j 0.+0.j]
+  ```
+
+* The `undo_swaps` transformation has been created to automatically remove all
+  swaps of a circuit.
   [(#1960)](https://github.com/PennyLaneAI/pennylane/pull/1960)
 
+  ```python
+  dev = qml.device('default.qubit', wires=3)
+
+  @qml.qnode(dev)
+  @qml.transforms.undo_swaps
+  def qfunc():
+      qml.Hadamard(wires=0)
+      qml.PauliX(wires=1)
+      qml.SWAP(wires=[0,1])
+      qml.SWAP(wires=[0,2])
+      qml.PauliY(wires=0)
+      return qml.expval(qml.PauliZ(0))
+  ```
+  
+  ```pycon
+  >>> print(qml.draw(qfunc)())
+   0: ──Y──┤ ⟨Z⟩
+   1: ──H──┤
+   2: ──X──┤
+  ```
+
 <h3>Improvements</h3>
 
-* The PennyLane `qchem` package is now lazily imported; it will only be imported
-  the first time it is accessed.
-  [(#1962)](https://github.com/PennyLaneAI/pennylane/pull/1962)
+* Added functions for computing the values of atomic and molecular orbitals at a given position.
+  [(#1867)](https://github.com/PennyLaneAI/pennylane/pull/1867)
 
-* Change all instances of `"{}".format(..)` to `f"{..}"`.
-  [(#1970)](https://github.com/PennyLaneAI/pennylane/pull/1970)
+  The functions `atomic_orbital` and `molecular_orbital` can be used, as shown in the
+  following codeblock, to evaluate the orbitals. By generating values of the orbitals at different
+  positions, one can plot the spatial shape of a desired orbital.
 
-* Tests do not loop over automatically imported and instantiated operations any more,
+  ```python
+  symbols  = ['H', 'H']
+  geometry = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]], requires_grad = False)
+  mol = hf.Molecule(symbols, geometry)
+  hf.generate_scf(mol)()
 
-* The QNode has been re-written to support batch execution across the board,
-  custom gradients, better decomposition strategies, and higher-order derivatives.
-  [(#1807)](https://github.com/PennyLaneAI/pennylane/pull/1807)
-  [(#1969)](https://github.com/PennyLaneAI/pennylane/pull/1969)
+  ao = mol.atomic_orbital(0)
+  mo = mol.molecular_orbital(1)
+  ```
 
-  - Internally, if multiple circuits are generated for simultaneous execution, they
-    will be packaged into a single job for execution on the device. This can lead to
-    significant performance improvement when executing the QNode on remote
-    quantum hardware or simulator devices with parallelization capabilities.
+  ```pycon
+  >>> print(ao(0.0, 0.0, 0.0))
+  >>> print(mo(0.0, 0.0, 0.0))
+  0.6282468778183719
+  0.018251285973461928
+  ```
 
-  - Custom gradient transforms can be specified as the differentiation method:
+* Added support for Python 3.10.
+  [(#1964)](https://github.com/PennyLaneAI/pennylane/pull/1964)
 
-    ```python
-    @qml.gradients.gradient_transform
-    def my_gradient_transform(tape):
-        ...
-        return tapes, processing_fn
+* The execution of QNodes that have
 
-    @qml.qnode(dev, diff_method=my_gradient_transform)
-    def circuit():
-    ```
+  - multiple return types;
+  - a return type other than Variance and Expectation
 
-  - Arbitrary :math:`n`-th order derivatives are supported on hardware using gradient transforms
-    such as the parameter-shift rule. To specify that an :math:`n`-th order derivative of a QNode
-    will be computed, the `max_diff` argument should be set. By default, this is set to 1
-    (first-order derivatives only). Increasing this value allows for higher order derivatives to be
-    extracted, at the cost of additional (classical) computational overhead during the backwards
-    pass.
+  now raises a descriptive error message when using the JAX interface.
+  [(#2011)](https://github.com/PennyLaneAI/pennylane/pull/2011)
 
-  - When decomposing the circuit, the default decomposition strategy `expansion_strategy="gradient"`
-    will prioritize decompositions that result in the smallest number of parametrized operations
-    required to satisfy the differentiation method. While this may lead to a slight increase in
-    classical processing, it significantly reduces the number of circuit evaluations needed to
-    compute gradients of complicated unitaries.
+* The PennyLane `qchem` package is now lazily imported; it will only be imported
+  the first time it is accessed.
+  [(#1962)](https://github.com/PennyLaneAI/pennylane/pull/1962)
 
-    To return to the old behaviour, `expansion_strategy="device"` can be specified.
+* `qml.math.scatter_element_add` now supports adding multiple values at
+  multiple indices with a single function call, in all interfaces
+  [(#1864)](https://github.com/PennyLaneAI/pennylane/pull/1864)
 
-  Note that the old QNode remains accessible at `@qml.qnode_old.qnode`, however this will
-  be removed in the next release.
+  For example, we may set five values of a three-dimensional tensor
+  in the following way:
+
+  ```pycon
+  >>> X = tf.zeros((3, 2, 9), dtype=tf.float64)
+  >>> indices = [(0, 0, 1, 2, 2), (0, 0, 0, 0, 1), (1, 3, 8, 6, 7)]
+  >>> values = [1 * i for i in range(1,6)]
+  >>> qml.math.scatter_element_add(X, indices, values)
+  <tf.Tensor: shape=(3, 2, 9), dtype=float64, numpy=
+  array([[[0., 1., 0., 2., 0., 0., 0., 0., 0.],
+          [0., 0., 0., 0., 0., 0., 0., 0., 0.]],
+
+         [[0., 0., 0., 0., 0., 0., 0., 0., 3.],
+          [0., 0., 0., 0., 0., 0., 0., 0., 0.]],
+
+         [[0., 0., 0., 0., 0., 0., 4., 0., 0.],
+          [0., 0., 0., 0., 0., 0., 0., 5., 0.]]])>
+  ```
+
+* All instances of `str.format` have been replace with f-strings.
+  [(#1970)](https://github.com/PennyLaneAI/pennylane/pull/1970)
 
 * Tests do not loop over automatically imported and instantiated operations any more,
   which was opaque and created unnecessarily many tests.
@@ -450,29 +542,65 @@
   [RZ(0.3, wires=[0])]
   ```
 
-* ``qml.circuit_drawer.tape_mpl`` produces a matplotlib figure and axes given a tape.
+* `qml.circuit_drawer.tape_mpl` produces a matplotlib figure and axes given a tape.
   [(#1787)](https://github.com/PennyLaneAI/pennylane/pull/1787)
 
-* AngleEmbedding now supports `batch_params` decorator. [(#1812)](https://github.com/PennyLaneAI/pennylane/pull/1812)
+* The `AngleEmbedding`, `BasicEntanglerLayers` and `MottonenStatePreparation`
+  templates now support parameters with batch dimension when using the `@qml.batch_params` decorator.
+  [(#1812)](https://github.com/PennyLaneAI/pennylane/pull/1812)
+  [(#1883)](https://github.com/PennyLaneAI/pennylane/pull/1883)
+  [(#1893)](https://github.com/PennyLaneAI/pennylane/pull/1893)
 
-* Added a new `qml.PauliError` channel that allows the application of an arbitrary number of Pauli operators on an arbitrary number of wires.
-  [(#1781)](https://github.com/PennyLaneAI/pennylane/pull/1781)
+* `qml.draw` now supports a `max_length` argument to help prevent text overflows when printing circuits.
+  [(#1892)](https://github.com/PennyLaneAI/pennylane/pull/1892)
 
-* BasicEntanglerLayers now supports `batch_params` decorator. [(#1883)](https://github.com/PennyLaneAI/pennylane/pull/1883)
+* `Identity` operation is now part of both the `ops.qubit` and `ops.cv`
+  modules.
+  [(#1956)](https://github.com/PennyLaneAI/pennylane/pull/1956)
 
-* MottonenStatePreparation now supports `batch_params` decorator. [(#1893)](https://github.com/PennyLaneAI/pennylane/pull/1893)
+<h3>Breaking changes</h3>
 
-* CircuitDrawer now supports a `max_length` argument to help prevent text overflows when printing circuits to the CLI. [#1841](https://github.com/PennyLaneAI/pennylane/pull/1841)
+* The QNode has been re-written to support batch execution across the board,
+  custom gradients, better decomposition strategies, and higher-order derivatives.
+  [(#1807)](https://github.com/PennyLaneAI/pennylane/pull/1807)
+  [(#1969)](https://github.com/PennyLaneAI/pennylane/pull/1969)
+
+  - Arbitrary :math:`n`-th order derivatives are supported on hardware using gradient transforms
+    such as the parameter-shift rule. To specify that an :math:`n`-th order derivative of a QNode
+    will be computed, the `max_diff` argument should be set. By default, this is set to 1
+    (first-order derivatives only). Increasing this value allows for higher order derivatives to be
+    extracted, at the cost of additional (classical) computational overhead during the backwards
+    pass.
 
-* `Identity` operation is now part of both the `ops.qubit` and `ops.cv` modules.
-   [(#1956)](https://github.com/PennyLaneAI/pennylane/pull/1956)
+  - When decomposing the circuit, the default decomposition strategy `expansion_strategy="gradient"`
+    will prioritize decompositions that result in the smallest number of parametrized operations
+    required to satisfy the differentiation method. While this may lead to a slight increase in
+    classical processing, it significantly reduces the number of circuit evaluations needed to
+    compute gradients of complicated unitaries.
 
-<h3>Breaking changes</h3>
+    To return to the old behaviour, `expansion_strategy="device"` can be specified.
 
-* The default behaviour of the `qml.metric_tensor` transform has been modified:
+  Note that the old QNode remains accessible at `@qml.qnode_old.qnode`, however this will
+  be removed in the next release.
+
+* Certain features deprecated in `v0.19.0` have been removed:
+  [(#1981)](https://github.com/PennyLaneAI/pennylane/pull/1981)
+  [(#1963)](https://github.com/PennyLaneAI/pennylane/pull/1963)
+
+  - The `qml.template` decorator (use a [
+    QuantumTape](https://pennylane.readthedocs.io/en/stable/code/api/pennylane.tape.QuantumTape.html)
+    as a context manager to record operations and its `operations` attribute to
+    return them, see the linked page for examples);
+  - The `default.tensor` and `default.tensor.tf` experimental devices;
+  - The `qml.fourier.spectrum` function (use the `qml.fourier.circuit_spectrum`
+    or `qml.fourier.qnode_spectrum` functions instead);
+  - The `diag_approx` keyword argument of `qml.metric_tensor` and
+    `qml.QNGOptimizer` (pass `approx='diag'` instead).
+
+* The default behaviour of the `qml.metric_tensor` transform has been modified.
   By default, the full metric tensor is computed, leading to higher cost than the previous
   default of computing the block diagonal only. At the same time, the Hadamard tests for
-  the full metric tensor require an additional wire on the device, so that 
+  the full metric tensor require an additional wire on the device, so that
 
   ```pycon
   >>> qml.metric_tensor(some_qnode)(weights)
@@ -488,10 +616,14 @@
 * The `par_domain` attribute in the operator class has been removed.
   [(#1907)](https://github.com/PennyLaneAI/pennylane/pull/1907)
 
-- The `mutable` keyword argument has been removed from the QNode.
+* The `mutable` keyword argument has been removed from the QNode,
+  due to underlying bugs that result in incorrect results being
+  returned from immutable QNodes. This functionality will return
+  in an upcoming release.
   [(#1807)](https://github.com/PennyLaneAI/pennylane/pull/1807)
 
-- The reversible QNode differentiation method has been removed.
+* The reversible QNode differentiation method has been removed; the adjoint
+  differentiation method is preferred instead (`diff_method='adjoint'`).
   [(#1807)](https://github.com/PennyLaneAI/pennylane/pull/1807)
 
 * `QuantumTape.trainable_params` now is a list instead of a set. This
@@ -502,6 +634,7 @@
 * The `num_params` attribute in the operator class is now dynamic. This makes it easier
   to define operator subclasses with a flexible number of parameters.
   [(#1898)](https://github.com/PennyLaneAI/pennylane/pull/1898)
+  [(#1909)](https://github.com/PennyLaneAI/pennylane/pull/1909)
 
 * The static method `decomposition()`, formerly in the `Operation` class, has
   been moved to the base `Operator` class.
@@ -513,20 +646,32 @@
   Instead, devices can check for the diagonal
   property using attributes:
 
-  ``` python
+  ```python
   from pennylane.ops.qubit.attributes import diagonal_in_z_basis
 
   if op in diagonal_in_z_basis:
       # do something
   ```
+  Custom operations can be added to this attribute at runtime via
+  `diagonal_in_z_basis.add("MyCustomOp")`.
 
-<h3>Deprecations</h3>
+<h3>Bug fixes</h3>
 
-* The init module, which contains functions to generate random parameters for 
-  templates, has been removed. Instead, the templates provide a `shape()` method.
-  [(#1963)](https://github.com/PennyLaneAI/pennylane/pull/1963)
+* Fixes a bug with `qml.probs` when using `default.qubit.jax`.
+  [(#1998)](https://github.com/PennyLaneAI/pennylane/pull/1998)
 
-<h3>Bug fixes</h3>
+* Fixes a bug where output tensors of a QNode would always be put on the
+  default GPU with `default.qubit.torch`.
+  [(#1982)](https://github.com/PennyLaneAI/pennylane/pull/1982)
+
+* Device test suite doesn't use empty circuits so that it can also
+  test the IonQ plugin, and it checks if operations are supported in
+  more places.
+  [(#1979)](https://github.com/PennyLaneAI/pennylane/pull/1979)
+
+* Fixes a bug where the metric tensor was computed incorrectly when using
+  gates with `gate.inverse=True`.
+  [(#1987)](https://github.com/PennyLaneAI/pennylane/pull/1987)
 
 * Corrects the documentation of `qml.transforms.classical_jacobian`
   for the Autograd interface (and improves test coverage).
@@ -536,6 +681,9 @@
   interface raised an error.
   [(#1906)](https://github.com/PennyLaneAI/pennylane/pull/1906)
 
+* Fixes a bug with the adjoint of `qml.QFT`.
+  [(#1955)](https://github.com/PennyLaneAI/pennylane/pull/1955)
+
 * Fixes a bug where the `ApproxTimeEvolution` template was not correctly
   computing the operation wires from the input Hamiltonian. This did not
   affect computation with the `ApproxTimeEvolution` template, but did
@@ -545,7 +693,7 @@
 * Fixes a bug where the classical preprocessing Jacobian
   computed by `qml.transforms.classical_jacobian` with JAX
   returned a reduced submatrix of the Jacobian.
-  [(#1935)](https://github.com/PennyLaneAI/pennylane/pull/1935)
+  [(#1948)](https://github.com/PennyLaneAI/pennylane/pull/1948)
 
 * Fixes a bug where the operations are not accessed in the correct order
   in `qml.fourier.qnode_spectrum`, leading to wrong outputs.
@@ -591,6 +739,30 @@
   optimization step updates.
   [(#1929)](https://github.com/PennyLaneAI/pennylane/pull/1929)
 
+* Fixes a bug where differentiating a QNode with multiple array
+  arguments via `qml.gradients.param_shift` throws an error.
+  [(#1989)](https://github.com/PennyLaneAI/pennylane/pull/1989)
+
+* `AmplitudeEmbedding` template no longer produces a `ComplexWarning`
+  when the `features` parameter is batched and provided as a 2D array.
+  [(#1990)](https://github.com/PennyLaneAI/pennylane/pull/1990)
+
+* `qml.circuit_drawer.CircuitDrawer` no longer produces an error
+  when attempting to draw tapes inside of circuits (e.g. from
+  decomposition of an operation or manual placement).
+  [(#1994)](https://github.com/PennyLaneAI/pennylane/pull/1994)
+
+* Fixes a bug where using SciPy sparse matrices with the new QNode
+  could lead to a warning being raised about prioritizing the TensorFlow
+  and PyTorch interfaces.
+  [(#2001)](https://github.com/PennyLaneAI/pennylane/pull/2001)
+
+* Fixed a bug where the `QueueContext` was not empty when first importing PennyLane.
+  [(#1957)](https://github.com/PennyLaneAI/pennylane/pull/1957)
+
+* Fixed circuit drawing problem with `Interferometer` and `CVNeuralNet`.
+  [(#1953)](https://github.com/PennyLaneAI/pennylane/pull/1953)
+
 <h3>Documentation</h3>
 
 * Added examples in documentation for some operations.
@@ -599,8 +771,8 @@
 * Improves the Developer's Guide Testing document.
   [(#1896)](https://github.com/PennyLaneAI/pennylane/pull/1896)
 
-* Add documentation example for AngleEmbedding, BasisEmbedding, StronglyEntanglingLayers, SqueezingEmbedding, DisplacementEmbedding,
-  MottonenStatePreparation and Interferometer.
+* Added documentation examples for `AngleEmbedding`, `BasisEmbedding`, `StronglyEntanglingLayers`,
+  `SqueezingEmbedding`, `DisplacementEmbedding`, `MottonenStatePreparation` and `Interferometer`.
   [(#1910)](https://github.com/PennyLaneAI/pennylane/pull/1910)
   [(#1908)](https://github.com/PennyLaneAI/pennylane/pull/1908)
   [(#1912)](https://github.com/PennyLaneAI/pennylane/pull/1912)
@@ -608,16 +780,11 @@
   [(#1936)](https://github.com/PennyLaneAI/pennylane/pull/1936)
   [(#1937)](https://github.com/PennyLaneAI/pennylane/pull/1937)
 
-* QueueContext was not empty when importing `pennylane`.
-
-* Fixed circuit drawing problem with Interferometer and CVNeuralNet.
-  [(#1953)](https://github.com/PennyLaneAI/pennylane/issues/1953)
-  
 <h3>Contributors</h3>
 
 This release contains contributions from (in alphabetical order):
 
-Guillermo Alonso-Linaje, Juan Miguel Arrazola, Samuel Banning, Benjamin Cordier, Alain Delgado,
-Olivia Di Matteo, David Ittah, Josh Izaac, Soran Jahangiri, Jalani Kanem, Ankit Khandelwal, Shumpei
-Kobayashi, Robert Lang, Christina Lee, Cedric Lin, Alejandro Montanez, Romain Moyard, Antal Száva,
-Maria Schuld, Jay Soni, Rodrigo Vargas, David Wierichs, Roeland Wiersema, Moritz Willmann
+Catalina Albornoz, Guillermo Alonso-Linaje, Juan Miguel Arrazola, Ali Asadi, Utkarsh Azad, Samuel Banning, Benjamin Cordier, Alain Delgado,
+Olivia Di Matteo, Anthony Hayes, David Ittah, Josh Izaac, Soran Jahangiri, Jalani Kanem, Ankit Khandelwal, Nathan Killoran, Shumpei
+Kobayashi, Robert Lang, Christina Lee, Cedric Lin, Alejandro Montanez, Romain Moyard, Lee James O'Riordan, Chae-Yeun Park, Isidor Schoch,
+Maria Schuld, Jay Soni, Antal Száva, Rodrigo Vargas, David Wierichs, Roeland Wiersema, Moritz Willmann.
diff --git a/pennylane/devices/default_qubit_jax.py b/pennylane/devices/default_qubit_jax.py
index caf65ecbd..b54d59a83 100644
--- a/pennylane/devices/default_qubit_jax.py
+++ b/pennylane/devices/default_qubit_jax.py
@@ -16,6 +16,7 @@ reference plugin.
 """
 import pennylane as qml
 from pennylane.devices import DefaultQubit
+from pennylane.wires import Wires
 
 import numpy as np
 
@@ -229,3 +230,51 @@ class DefaultQubitJax(DefaultQubit):
         powers_of_two = 1 << jnp.arange(num_wires, dtype=dtype)
         states_sampled_base_ten = samples[:, None] & powers_of_two
         return (states_sampled_base_ten > 0).astype(dtype)[:, ::-1]
+
+    def estimate_probability(self, wires=None, shot_range=None, bin_size=None):
+        """Return the estimated probability of each computational basis state
+        using the generated samples.
+
+        Args:
+            wires (Iterable[Number, str], Number, str, Wires): wires to calculate
+                marginal probabilities for. Wires not provided are traced out of the system.
+            shot_range (tuple[int]): 2-tuple of integers specifying the range of samples
+                to use. If not specified, all samples are used.
+            bin_size (int): Divides the shot range into bins of size ``bin_size``, and
+                returns the measurement statistic separately over each bin. If not
+                provided, the entire shot range is treated as a single bin.
+
+        Returns:
+            array[float]: list of the probabilities
+        """
+
+        wires = wires or self.wires
+        # convert to a wires object
+        wires = Wires(wires)
+        # translate to wire labels used by device
+        device_wires = self.map_wires(wires)
+        num_wires = len(device_wires)
+
+        sample_slice = Ellipsis if shot_range is None else slice(*shot_range)
+        samples = self._samples[sample_slice, device_wires]
+
+        # convert samples from a list of 0, 1 integers, to base 10 representation
+        powers_of_two = 2 ** np.arange(len(device_wires))[::-1]
+        indices = samples @ powers_of_two
+
+        if bin_size is not None:
+            raise ValueError(
+                "The default.qubit.jax device doesn't support getting probabilities when using a shot vector."
+            )
+
+        basis_states, counts = qml.math.unique(
+            indices, return_counts=True, size=2 ** num_wires, fill_value=-1
+        )
+        prob = np.zeros([2 ** num_wires + 1], dtype=np.float64)
+        prob = qml.math.convert_like(prob, indices)
+
+        for state, count in zip(basis_states, counts):
+            prob = prob.at[state].set(count / len(samples))
+
+        prob = jnp.resize(prob, 2 ** num_wires)  # resize prob which discards the 'filled values'
+        return self._asarray(prob, dtype=self.R_DTYPE)
diff --git a/pennylane/devices/default_qubit_torch.py b/pennylane/devices/default_qubit_torch.py
index cfa6fb7fa..71838cbe0 100644
--- a/pennylane/devices/default_qubit_torch.py
+++ b/pennylane/devices/default_qubit_torch.py
@@ -126,7 +126,8 @@ class DefaultQubitTorch(DefaultQubit):
             If ``shots > 0`` is used, the ``diff_method="backprop"``
             QNode differentiation method is not supported and it is recommended to consider
             switching device to ``default.qubit`` and using ``diff_method="parameter-shift"``.
-        torch_device='cpu' (str): the device on which the computation will be run, ``'cpu'`` or ``'cuda'``
+        torch_device='cpu' (str): the device on which the computation will be
+        run, e.g., ``'cpu'`` or ``'cuda'``
     """
 
     name = "Default qubit (Torch) PennyLane plugin"
@@ -167,38 +168,66 @@ class DefaultQubitTorch(DefaultQubit):
         self._state = self._state.to(self._torch_device)
         self._pre_rotated_state = self._state
 
+    @staticmethod
+    def _get_parameter_torch_device(ops):
+        """An auxiliary function to determine the Torch device specified for
+        the gate parameters of the input operations.
+
+        Returns the first CUDA Torch device found (if any) using a string
+        format. Does not handle tensors put on multiple CUDA Torch devices.
+        Such a case raises an error with Torch.
+
+        If CUDA is not used with any of the parameters, then specifies the CPU
+        if the parameters are on the CPU or None if there were no parametric
+        operations.
+
+        Args:
+            ops (list[Operator]): list of operations to check
+
+        Returns:
+            str or None: The string of the Torch device determined or None if
+            there is no data for any operations.
+        """
+        par_torch_device = None
+        for op in ops:
+            for data in op.data:
+
+                # Using hasattr in case we don't have a Torch tensor as input
+                if hasattr(data, "is_cuda"):
+                    if data.is_cuda:  # pragma: no cover
+                        return ":".join([data.device.type, str(data.device.index)])
+
+                    par_torch_device = "cpu"
+
+        return par_torch_device
+
     def execute(self, circuit, **kwargs):
         ops_and_obs = circuit.operations + circuit.observables
 
-        some_params_cuda = any(
-            data.is_cuda for op in ops_and_obs for data in op.data if hasattr(data, "is_cuda")
-        )
-        if not self._torch_device_specified:
-            if some_params_cuda:
-
-                self._torch_device = "cuda"
-            else:
+        par_torch_device = self._get_parameter_torch_device(ops_and_obs)
 
-                # need to reset in case last execution moved to cuda
-                self._torch_device = "cpu"
+        if not self._torch_device_specified:
+            self._torch_device = par_torch_device
 
-            # If we've changed the device, place the state on the correct
-            # device
+            # If we've changed the device of the parameters between device
+            # executions, need to move the state to the correct Torch device
             if self._state.device != self._torch_device:
                 self._state = self._state.to(self._torch_device)
         else:
-            cuda_params_but_not_cuda_specified = some_params_cuda and self._torch_device != "cuda"
-            not_cuda_params_but_cuda_specified = (
-                not some_params_cuda and self._torch_device == "cuda"
-            )
-            if cuda_params_but_not_cuda_specified or not_cuda_params_but_cuda_specified:
-
-                warnings.warn(
-                    f"Torch device {self._torch_device} specified "
-                    "upon PennyLane device creation does not match the "
-                    "Torch device of the gate parameters; "
-                    f"{self._torch_device} will be used."
-                )
+            if par_torch_device is not None:  # pragma: no cover
+                params_cuda_device = "cuda" in par_torch_device
+                specified_device_cuda = "cuda" in self._torch_device
+
+                # Raise a warning if there's a mismatch between the specified and
+                # used Torch devices
+                if params_cuda_device != specified_device_cuda:
+
+                    warnings.warn(
+                        f"Torch device {self._torch_device} specified "
+                        "upon PennyLane device creation does not match the "
+                        "Torch device of the gate parameters; "
+                        f"{self._torch_device} will be used."
+                    )
 
         return super().execute(circuit, **kwargs)
 
diff --git a/pennylane/drawer/style.py b/pennylane/drawer/style.py
index 270660190..5161b3f75 100644
--- a/pennylane/drawer/style.py
+++ b/pennylane/drawer/style.py
@@ -102,7 +102,7 @@ def use_style(style: str):
 
     .. code-block:: python
 
-        qml.style.use_style('black_white')
+        qml.drawer.use_style('black_white')
 
         @qml.qnode(qml.device('lightning.qubit', wires=(0,1,2,3)))
         def circuit(x, z):
diff --git a/pennylane/fourier/reconstruct.py b/pennylane/fourier/reconstruct.py
index 29290c6be..3e10f2daa 100644
--- a/pennylane/fourier/reconstruct.py
+++ b/pennylane/fourier/reconstruct.py
@@ -438,7 +438,7 @@ def reconstruct(qnode, ids=None, nums_frequency=None, spectra=None, shifts=None)
     >>> recon_Y1 = rec["Y"][(1,)]
     >>> np.isclose(recon_Y0(Y[0]), circuit_value)
     True
-    >>> np.isclose(recon_Y1(Y[1]+1.3), circuit(x, Y+np.eye(2)[1]*1.3)
+    >>> np.isclose(recon_Y1(Y[1]+1.3), circuit(x, Y+np.eye(2)[1]*1.3))
     True
 
     We successfully reconstructed the dependence on the two entries of ``Y`` ,
@@ -459,7 +459,7 @@ def reconstruct(qnode, ids=None, nums_frequency=None, spectra=None, shifts=None)
     dict_keys(['x', 'Y'])
     >>> spectra["x"]
     {(): [-1.0, 0.0, 1.0]}
-    >>> print(*_spectra["Y"].items(), sep="\n")
+    >>> print(*spectra["Y"].items(), sep="\n")
     ((0,), [-1.0, 0.0, 1.0])
     ((1,), [-6.0, -5.0, -4.0, -1.0, 0.0, 1.0, 4.0, 5.0, 6.0])
 
diff --git a/pennylane/gradients/general_shift_rules.py b/pennylane/gradients/general_shift_rules.py
index ab554e7c9..b463042c5 100644
--- a/pennylane/gradients/general_shift_rules.py
+++ b/pennylane/gradients/general_shift_rules.py
@@ -216,8 +216,8 @@ def generate_shift_rule(frequencies, shifts=None, order=1):
     >>> eigvals = (-0.5, 0, 0, 0.5)
     >>> frequencies = eigvals_to_frequencies(eigvals)
     >>> generate_shift_rule(frequencies)
-    [[ 0.85355339, -0.85355339, -0.14644661,  0.14644661],
-     [0.78539816,  -0.78539816,  2.35619449, -2.35619449]]
+    [[ 0.4267767  -0.4267767  -0.0732233   0.0732233 ]
+     [ 1.57079633 -1.57079633  4.71238898 -4.71238898]]
 
     An example with explicitly specified shift values:
 
diff --git a/pennylane/gradients/gradient_transform.py b/pennylane/gradients/gradient_transform.py
index b6e8a00e1..bd9c5ab63 100644
--- a/pennylane/gradients/gradient_transform.py
+++ b/pennylane/gradients/gradient_transform.py
@@ -14,6 +14,8 @@
 """This module contains utilities for defining custom gradient transforms,
 including a decorator for specifying gradient expansions."""
 # pylint: disable=too-few-public-methods
+import warnings
+
 import pennylane as qml
 from pennylane.transforms.tape_expand import expand_invalid_trainable
 
@@ -108,6 +110,27 @@ class gradient_transform(qml.batch_transform):
         self.hybrid = hybrid
         super().__init__(transform_fn, expand_fn=expand_fn, differentiable=differentiable)
 
+    @staticmethod
+    def _jacobian_trainable_args(args, interface):
+        """Return the indices of QNode arguments for which a Jacobian was
+        computed by `qml.transforms.classical_jacobian` with argnum=None.
+        """
+        trainable_args = []
+
+        if interface == "autograd":
+            for idx, arg in enumerate(args):
+                # TODO: make default False once this change is done in qml.jacobian
+                if getattr(arg, "requires_grad", True):
+                    trainable_args.append(idx)
+
+        elif interface == "jax":
+            trainable_args = [0]
+
+        # Torch and Tensorflow interfaces are not considered since `classical_jacobian`
+        # always returns a tuple for them, thus not invoking this function.
+
+        return trainable_args
+
     def default_qnode_wrapper(self, qnode, targs, tkwargs):
         # Here, we overwrite the QNode execution wrapper in order
         # to take into account that classical processing may be present
@@ -144,8 +167,42 @@ class gradient_transform(qml.batch_transform):
                 # is present inside the QNode.
                 return qjac
 
-            # Classical processing of a single argument is present. Return qjac @ cjac.
-            jac = qml.math.squeeze(qml.math.tensordot(qml.math.T(cjac), qjac, [[-1], [-1]]))
-            return qml.math.T(jac)
+            # Classical processing present of either:
+            #   a) a single argument or
+            #   b) multiple arguments of the same shape
+            # The shape of the classical jacobian returned by qml.jacobian (invoked by
+            # classical_jacobian with autograd) depends on the scenario:
+            #   a) (# gate args, qnode arg shape)
+            #   b) (reverse qnode arg shape, # gate args, # qnode args)
+            # It then needs to be contracted with the quantum jacobian of shape:
+            #      (qnode output shape, # gate args)
+            # The result should have shape:
+            #   a) (qnode ouput shape, qnode arg shape)
+            #   b) (reverse qnode arg shape, qnode output shape, # qnode args)
+            num_gate_args = qml.math.shape(qjac)[-1]
+            # Consider only QNode arguments regarded as trainable by the interfaces.
+            trainable_args_idx = self._jacobian_trainable_args(args, qnode.interface)
+            # Since all arguments have the same shape, obtain shape from the first trainable arg
+            qnode_arg_shape = qml.math.shape(args[trainable_args_idx[0]])
+            num_qnode_args = len(trainable_args_idx)
+
+            if qml.math.shape(cjac) == (num_gate_args, *qnode_arg_shape):
+                # single QNode argument
+                jac = qml.math.tensordot(qjac, cjac, [[-1], [0]])
+            elif qml.math.shape(cjac) == (*qnode_arg_shape[::-1], num_gate_args, num_qnode_args):
+                # multiple QNode arguments with stacking
+                cjac = qml.math.swapaxes(cjac, -1, -2)
+                jac = qml.math.tensordot(cjac, qjac, [[-1], [-1]])
+                jac = qml.math.moveaxis(jac, -1, len(qnode_arg_shape))
+            else:  # pragma: no cover
+                jac = ()
+                warnings.warn(
+                    "Unexpected classical Jacobian encoutered, could not compute the hybrid "
+                    "Jacobian of the QNode. You can still attempt to obtain the quantum Jacobian "
+                    "with the `hybrid=False` parameter.",
+                    UserWarning,
+                )
+
+            return qml.math.squeeze(jac)
 
         return jacobian_wrapper
diff --git a/pennylane/hf/molecule.py b/pennylane/hf/molecule.py
index 4955ce176..61171b9a7 100644
--- a/pennylane/hf/molecule.py
+++ b/pennylane/hf/molecule.py
@@ -172,7 +172,7 @@ class Molecule:
         >>> symbols  = ['H', 'H']
         >>> geometry = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]], requires_grad = False)
         >>> mol = qml.hf.Molecule(symbols, geometry)
-        >>> generate_scf(mol)() # run scf to obtain the optimized molecular orbitals
+        >>> qml.hf.generate_scf(mol)() # run scf to obtain the optimized molecular orbitals
         >>> mo = mol.molecular_orbital(1)
         >>> mo(0.0, 0.0, 0.0)
         0.01825128
diff --git a/pennylane/interfaces/batch/jax.py b/pennylane/interfaces/batch/jax.py
index d4e2f1373..1de378589 100644
--- a/pennylane/interfaces/batch/jax.py
+++ b/pennylane/interfaces/batch/jax.py
@@ -23,6 +23,7 @@ from jax.experimental import host_callback
 
 import numpy as np
 import pennylane as qml
+from pennylane.operation import Variance, Expectation
 
 dtype = jnp.float64
 
@@ -86,6 +87,28 @@ def execute(tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n=1, max_d
     )
 
 
+def _validate_tapes(tapes):
+    """Validates that the input tapes are compatible with JAX support.
+
+    Raises:
+        ValueError: if tapes with non-scalar outputs were provided or a return
+        type other than variance and expectation value was used
+    """
+    for t in tapes:
+
+        if len(t.observables) != 1:
+            raise ValueError(
+                "The JAX interface currently only supports quantum nodes with a single return type."
+            )
+
+        for o in t.observables:
+            return_type = o.return_type
+            if return_type is not Variance and return_type is not Expectation:
+                raise ValueError(
+                    f"Only Variance and Expectation returns are supported for the JAX interface, given {return_type}."
+                )
+
+
 def _execute(
     params,
     tapes=None,
@@ -96,6 +119,8 @@ def _execute(
     _n=1,
 ):  # pylint: disable=dangerous-default-value,unused-argument
 
+    _validate_tapes(tapes)
+
     # Only have scalar outputs
     total_size = len(tapes)
     total_params = np.sum([len(p) for p in params])
@@ -207,6 +232,8 @@ def _execute_with_fwd(
     _n=1,
 ):  # pylint: disable=dangerous-default-value,unused-argument
 
+    _validate_tapes(tapes)
+
     # Only have scalar outputs
     total_size = len(tapes)
 
diff --git a/pennylane/measure.py b/pennylane/measure.py
index 1dd1bb278..09507f9da 100644
--- a/pennylane/measure.py
+++ b/pennylane/measure.py
@@ -380,7 +380,8 @@ def sample(op=None, wires=None):
 def probs(wires=None, op=None):
     r"""Probability of each computational basis state.
 
-    This measurement function accepts no observables, and instead
+    This measurement function accepts either a wire specification or
+    an observable. Passing wires to the function
     instructs the QNode to return a flat array containing the
     probabilities :math:`|\langle i | \psi \rangle |^2` of measuring
     the computational basis state :math:`| i \rangle` given the current
diff --git a/pennylane/ops/channel.py b/pennylane/ops/channel.py
index 71acaa40c..5c6969f77 100644
--- a/pennylane/ops/channel.py
+++ b/pennylane/ops/channel.py
@@ -380,7 +380,7 @@ class ResetError(Channel):
 
 class PauliError(Channel):
     r"""PauliError(operators, p, wires)
-    Arbitrary number qubit, arbitrary Pauli operator error channel.
+    Pauli operator error channel for an arbitrary number of qubits.
 
     This channel is modelled by the following Kraus matrices:
 
@@ -650,9 +650,9 @@ class ThermalRelaxationError(Channel):
     * Number of parameters: 4
 
     Args:
-        pe (float): exited state population.
+        pe (float): exited state population. Must be between ``0`` and ``1``.
         t1 (float): the :math:`T_1` relaxation constant.
-        t2 (float): the :math:`T_2` dephasing constant.
+        t2 (float): the :math:`T_2` dephasing constant. Must be less than :math:`2 T_1`.
         tg (float): the gate time for relaxation error.
         wires (Sequence[int] or int): the wire the channel acts on
     """
@@ -670,9 +670,9 @@ class ThermalRelaxationError(Channel):
         t2 = params[2]
         tg = params[3]
         if not 0.0 <= pe <= 1.0:
-            raise ValueError("pe must be between")
+            raise ValueError("pe must be between 0 and 1.")
         if tg < 0:
-            raise ValueError(f"Invalid gate_time ({tg} < 0)")
+            raise ValueError(f"Invalid gate_time tg ({tg} < 0)")
         if t1 <= 0:
             raise ValueError("Invalid T_1 relaxation time parameter: T_1 <= 0.")
         if t2 <= 0:
diff --git a/pennylane/ops/qubit/state_preparation.py b/pennylane/ops/qubit/state_preparation.py
index 09c91cf95..168098904 100644
--- a/pennylane/ops/qubit/state_preparation.py
+++ b/pennylane/ops/qubit/state_preparation.py
@@ -134,7 +134,7 @@ class QubitDensityMatrix(Operation):
         target device.
 
     Args:
-        state (array[complex]): a density matrix of size (2**len(wires), 2**len(wires))
+        state (array[complex]): a density matrix of size ``(2**len(wires), 2**len(wires))``
         wires (Sequence[int] or int): the wire(s) the operation acts on
 
     .. UsageDetails::
@@ -151,8 +151,8 @@ class QubitDensityMatrix(Operation):
             dev = qml.device("default.mixed", wires=2)
             @qml.qnode(dev)
             def circuit():
-                qml.QubitDensityMatrix(initialize_state, wires=[0, 1])
-                 return qml.state()
+                qml.QubitDensityMatrix(rho, wires=[0, 1])
+                return qml.state()
 
         Running this circuit:
 
diff --git a/pennylane/templates/subroutines/commuting_evolution.py b/pennylane/templates/subroutines/commuting_evolution.py
index 094af7c9b..cedc3dba4 100644
--- a/pennylane/templates/subroutines/commuting_evolution.py
+++ b/pennylane/templates/subroutines/commuting_evolution.py
@@ -68,12 +68,12 @@ class CommutingEvolution(Operation):
         time (int or float): The time of evolution, namely the parameter :math:`t` in :math:`e^{- i H t}`.
 
     Keyword args:
-        frequencies (list[int or float]): The unique positive differences between eigenvalues in
+        frequencies (tuple[int or float]): The unique positive differences between eigenvalues in
             the spectrum of the Hamiltonian. If the frequencies are not given, the cost function
             partial derivative will be computed using the standard two-term shift rule applied to
             the constituent Pauli words in the Hamiltonian individually.
 
-        shifts (list[int or float]): The parameter shifts to use in obtaining the
+        shifts (tuple[int or float]): The parameter shifts to use in obtaining the
             generalized parameter shift rules. If unspecified, equidistant shifts are used.
 
     .. UsageDetails::
@@ -90,7 +90,7 @@ class CommutingEvolution(Operation):
             coeffs = [1, -1]
             obs = [qml.PauliX(0) @ qml.PauliY(1), qml.PauliY(0) @ qml.PauliX(1)]
             hamiltonian = qml.Hamiltonian(coeffs, obs)
-            frequencies = [2,4]
+            frequencies = (2,4)
 
             @qml.qnode(dev)
             def circuit(time):
diff --git a/pennylane/transforms/metric_tensor.py b/pennylane/transforms/metric_tensor.py
index bb0d7256c..0bb985639 100644
--- a/pennylane/transforms/metric_tensor.py
+++ b/pennylane/transforms/metric_tensor.py
@@ -183,7 +183,7 @@ def metric_tensor(tape, approx=None, allow_nonunitary=True, aux_wire=None, devic
     For example, we can compute the gradient of the Frobenius norm of the metric tensor
     with respect to the QNode ``weights`` :
 
-    >>> norm_fn = lambda x: qml.math.linalg(mt_fn(x), ord="fro")
+    >>> norm_fn = lambda x: qml.math.linalg.norm(mt_fn(x), ord="fro")
     >>> grad_fn = qml.grad(norm_fn)
     >>> grad_fn(weights)
     array([-0.0282246 ,  0.01340413,  0.        ,  0.        ])
diff --git a/pennylane/transforms/optimization/merge_amplitude_embedding.py b/pennylane/transforms/optimization/merge_amplitude_embedding.py
index 1f2b84ecb..70b34e7c2 100644
--- a/pennylane/transforms/optimization/merge_amplitude_embedding.py
+++ b/pennylane/transforms/optimization/merge_amplitude_embedding.py
@@ -50,7 +50,6 @@ def merge_amplitude_embedding(tape):
     >>> optimized_qfunc = qml.transforms.merge_amplitude_embedding(qfunc)
     >>> optimized_qnode = qml.QNode(optimized_qfunc, dev)
     >>> print(qml.draw(optimized_qnode)())
-
       0: ──╭C───────────────────────╭┤ State
       1: ──╰X───────────────────────├┤ State
       2: ──╭AmplitudeEmbedding(M0)──├┤ State
diff --git a/pennylane/transforms/optimization/remove_barrier.py b/pennylane/transforms/optimization/remove_barrier.py
index 53f7d317a..6b2b1d852 100644
--- a/pennylane/transforms/optimization/remove_barrier.py
+++ b/pennylane/transforms/optimization/remove_barrier.py
@@ -33,7 +33,7 @@ def remove_barrier(tape):
 
     .. code-block:: python
 
-        def qfunc(x, y, z):
+        def qfunc(x, y):
             qml.Hadamard(wires=0)
             qml.Hadamard(wires=1)
             qml.Barrier(wires=[0,1])
diff --git a/qchem/pennylane_qchem/_version.py b/qchem/pennylane_qchem/_version.py
index 51f74f5f4..0ee01e991 100644
--- a/qchem/pennylane_qchem/_version.py
+++ b/qchem/pennylane_qchem/_version.py
@@ -16,4 +16,4 @@
 Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.21.0-dev"
+__version__ = "0.20.0-dev"


Test patch:
diff --git a/tests/devices/test_default_qubit_jax.py b/tests/devices/test_default_qubit_jax.py
index 5f84c0855..9187771b7 100644
--- a/tests/devices/test_default_qubit_jax.py
+++ b/tests/devices/test_default_qubit_jax.py
@@ -1,3 +1,16 @@
+# Copyright 2018-2021 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.
 import pytest
 
 jax = pytest.importorskip("jax", minversion="0.2")
@@ -152,6 +165,49 @@ class TestQNodeIntegration:
         expected = jnp.array([amplitude, 0, jnp.conj(amplitude), 0])
         assert jnp.allclose(state, expected, atol=tol, rtol=0)
 
+    def test_probs_jax(self, tol):
+        """Test that returning probs works with jax"""
+        dev = qml.device("default.qubit.jax", wires=1, shots=100)
+        expected = jnp.array([0.0, 1.0])
+
+        @qml.qnode(dev, interface="jax")
+        def circuit():
+            qml.PauliX(wires=0)
+            return qml.probs()
+
+        result = circuit()
+        assert jnp.allclose(result, expected, atol=tol)
+
+    def test_probs_jax_jit(self, tol):
+        """Test that returning probs works with jax and jit"""
+        dev = qml.device("default.qubit.jax", wires=1, shots=100)
+        expected = jnp.array([0.0, 1.0])
+
+        @jax.jit
+        @qml.qnode(dev, interface="jax")
+        def circuit():
+            qml.PauliX(wires=0)
+            return qml.probs()
+
+        result = circuit()
+        assert jnp.allclose(result, expected, atol=tol)
+
+    def test_probs_jax_jit(self, tol):
+        """Test that returning probs works with jax and jit"""
+        dev = qml.device("default.qubit.jax", wires=1, shots=(2, 2))
+        expected = jnp.array([0.0, 1.0])
+
+        @jax.jit
+        @qml.qnode(dev, interface="jax")
+        def circuit():
+            qml.PauliX(wires=0)
+            return qml.probs()
+
+        with pytest.raises(
+            ValueError, match="doesn't support getting probabilities when using a shot vector"
+        ):
+            result = circuit()
+
     def test_sampling_with_jit(self):
         """Test that sampling works with a jax.jit"""
 
diff --git a/tests/devices/test_default_qubit_torch.py b/tests/devices/test_default_qubit_torch.py
index 193f1f2db..a3605f528 100644
--- a/tests/devices/test_default_qubit_torch.py
+++ b/tests/devices/test_default_qubit_torch.py
@@ -1619,7 +1619,9 @@ class TestHighLevelIntegration:
 
         torch.manual_seed(42)
         weights = torch.rand(
-            qml.templates.StronglyEntanglingLayers.shape(n_wires=2, n_layers=2), requires_grad=True
+            qml.templates.StronglyEntanglingLayers.shape(n_wires=2, n_layers=2),
+            requires_grad=True,
+            device=torch_device,
         )
 
         def cost(weights):
diff --git a/tests/gpu/test_torch.py b/tests/gpu/test_torch.py
index f71347fe2..a925c38cb 100644
--- a/tests/gpu/test_torch.py
+++ b/tests/gpu/test_torch.py
@@ -36,7 +36,7 @@ class TestTorchDevice:
         res = dev.execute(tape)
 
         assert res.is_cuda
-        assert dev._torch_device == "cuda"
+        assert "cuda" in dev._torch_device
 
         res.backward()
         assert x.grad.is_cuda
@@ -57,7 +57,7 @@ class TestTorchDevice:
         res = dev.execute(tape)
 
         assert res.is_cuda
-        assert dev._torch_device == "cuda"
+        assert "cuda" in dev._torch_device
 
         res.backward()
         assert x.grad.is_cuda
@@ -77,7 +77,7 @@ class TestTorchDevice:
 
         res = dev.execute(tape)
         assert res.is_cuda
-        assert dev._torch_device == "cuda"
+        assert "cuda" in dev._torch_device
 
     def test_resets(self):
         """Asserts reverts to cpu after execution on gpu"""
@@ -92,7 +92,7 @@ class TestTorchDevice:
             qml.expval(qml.PauliZ(0))
 
         res1 = dev.execute(tape1)
-        assert dev._torch_device == "cuda"
+        assert "cuda" in dev._torch_device
         assert res1.is_cuda
 
         with qml.tape.QuantumTape() as tape2:
@@ -120,7 +120,7 @@ class TestTorchDevice:
         res = circ(x, y)
         assert res.is_cuda
 
-        assert circ.device._torch_device == "cuda"
+        assert "cuda" in circ.device._torch_device
         res.backward()
         assert x.grad.is_cuda
 
@@ -170,7 +170,7 @@ class TestqnnTorchLayer:
         x = torch.rand((5, n_qubits), dtype=torch.float64).to(torch.device("cuda"))
         res = qlayer(x)
         assert circuit.device.short_name == "default.qubit.torch"
-        assert circuit.device._torch_device == "cuda"
+        assert "cuda" in circuit.device._torch_device
         assert res.is_cuda
 
         loss = torch.sum(res).squeeze()
diff --git a/tests/gradients/test_gradient_transform.py b/tests/gradients/test_gradient_transform.py
index 8da7bf86a..5b1d0c571 100644
--- a/tests/gradients/test_gradient_transform.py
+++ b/tests/gradients/test_gradient_transform.py
@@ -135,6 +135,46 @@ class TestGradientTransformIntegration:
         expected = qml.jacobian(circuit)(w)
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
+    def test_multiple_tensor_arguments(self, tol):
+        """Test that a gradient transform acts on QNodes
+        correctly when multiple tensor QNode arguments are present"""
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev)
+        def circuit(x, y):
+            qml.RX(x[0, 0], wires=0)
+            qml.RY(y[0, 0], wires=0)
+            qml.RZ(x[1, 0], wires=0)
+            return qml.probs(wires=[0, 1])
+
+        x = np.array([[0.1], [0.2]], requires_grad=True)
+        y = np.array([[0.2], [0.3]], requires_grad=True)
+
+        expected = qml.jacobian(circuit)(x, y)
+        res = qml.gradients.param_shift(circuit)(x, y)
+
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+    def test_first_non_trainable_argument(self, tol):
+        """Test that a gradient transform acts on QNodes
+        correctly when the first argument is non-trainable"""
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev)
+        def circuit(x, y):
+            qml.RX(y[0], wires=0)
+            qml.RY(y[0], wires=0)
+            qml.RZ(y[1], wires=0)
+            return qml.probs(wires=[0, 1])
+
+        x = np.array([0.1], requires_grad=False)
+        y = np.array([0.2, 0.3], requires_grad=True)
+
+        expected = qml.jacobian(circuit)(x, y)
+        res = qml.gradients.param_shift(circuit)(x, y)
+
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
     def test_classical_processing_arguments(self, mocker, tol):
         """Test that a gradient transform acts on QNodes
         correctly when the QNode arguments are classically processed"""
diff --git a/tests/interfaces/test_batch_jax_qnode.py b/tests/interfaces/test_batch_jax_qnode.py
index 896832664..fd3e911de 100644
--- a/tests/interfaces/test_batch_jax_qnode.py
+++ b/tests/interfaces/test_batch_jax_qnode.py
@@ -300,6 +300,49 @@ class TestQNode:
         )
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
+    def test_multiple_outputs_raises(self, dev_name, diff_method, mode, tol):
+        """Test executing a QNode that has multiple outputs raises an error."""
+        dev = qml.device(dev_name, wires=2)
+
+        if diff_method == "backprop":
+            pytest.skip("Test is not applicable for backprop")
+
+        @qml.qnode(dev, interface="jax", diff_method=diff_method, mode=mode)
+        def my_circuit(param):
+            qml.RX(param, wires=0)
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))
+
+        with pytest.raises(
+            ValueError,
+            match="JAX interface currently only supports quantum nodes with a single return type",
+        ):
+            my_circuit(1)
+
+    @pytest.mark.parametrize("ret", [qml.probs(wires=0), qml.state()])
+    def test_not_expval_or_var_raises(self, dev_name, diff_method, mode, ret, tol):
+        """Test executing a QNode that has a return type other than expval or
+        var raises an error."""
+        dev = qml.device(dev_name, wires=2)
+
+        if diff_method == "backprop":
+            pytest.skip("Test is not applicable for backprop")
+
+        if diff_method == "adjoint":
+            pytest.skip("Adjoint does not support states")
+
+        @qml.qnode(dev, interface="jax", diff_method=diff_method, mode=mode)
+        def my_circuit(param):
+            qml.RX(param, wires=0)
+            qml.CNOT(wires=[0, 1])
+            return qml.apply(ret)
+
+        with pytest.raises(
+            ValueError,
+            match="Only Variance and Expectation returns are supported for the JAX interface",
+        ):
+            my_circuit(1)
+
 
 class TestShotsIntegration:
     """Test that the QNode correctly changes shot value, and
diff --git a/tests/interfaces/test_batch_torch.py b/tests/interfaces/test_batch_torch.py
index d27522049..ee44dec26 100644
--- a/tests/interfaces/test_batch_torch.py
+++ b/tests/interfaces/test_batch_torch.py
@@ -399,7 +399,7 @@ class TestTorchExecuteIntegration:
             dev = qml.device("default.qubit.autograd", wires=2)
             return dev.batch_execute([tape])[0]
 
-        expected = qml.grad(cost)(0.1)
+        expected = qml.grad(cost, argnum=0)(0.1)
         assert torch.allclose(a.grad, torch.tensor(expected, device=torch_device), atol=tol, rtol=0)
 
     def test_jacobian(self, torch_device, execute_kwargs, tol):