Repository: PennyLaneAI/pennylane
Pull request #1864 (2021-11-05 07:18:15)

Add `qml.fourier.reconstruct`
**Context:**
Full reconstructions of one-dimensional Fourier series with arbitrary frequencies can be performed via discrete Fourier transforms. `qml.fourier.reconstruct` implements this both for equidistant frequencies and sampling points and for non-equidistant frequencies and sampling points.

**Description of the Change:**
Addition of `qml.fourier.reconstruct`

**Benefits:**
Reconstruction is now possible to do within PennyLane, for example in conjunction with `qml.fourier.circuit_spectrum` or qml.fourier.qnode_spectrum`.
This is an addition allowing us to implement an even more generalized version of Rotosolve than done in #1489

**Possible Drawbacks:**
N/A

**Related GitHub Issues:**
Part of a PR series that replaces #1635, #1681 being the latest of the series.


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

[![Impacted file tree graph](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1864/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/1864?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)

```diff
@@            Coverage Diff             @@
##           master    #1864      +/-   ##
==========================================
- Coverage   98.98%   98.59%   -0.40%     
==========================================
  Files         218      219       +1     
  Lines       16426    16543     +117     
==========================================
+ Hits        16260    16311      +51     
- Misses        166      232      +66     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1864?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/fourier/reconstruct.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1864/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ZvdXJpZXIvcmVjb25zdHJ1Y3QucHk=) | `43.10% <43.10%> (ø)` | |
| [pennylane/fourier/\_\_init\_\_.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1864/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ZvdXJpZXIvX19pbml0X18ucHk=) | `100.00% <100.00%> (ø)` | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1864?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/1864?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Last update [5443339...c98a9f4](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1864?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).


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 9bf246fb5..d971b8dad 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -4,18 +4,94 @@
 
 <h3>New features since last release</h3>
 
+* `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.]]])>
+  ```
+
+* The `qml.fourier.reconstruct` function is added. It can be used to
+  reconstruct QNodes outputting expectation values along a specified
+  parameter dimension, with a minimal number of calls to the
+  original QNode. The returned
+  reconstruction is exact and purely classical, and can be evaluated
+  without any quantum executions.
+  [(#1864)](https://github.com/PennyLaneAI/pennylane/pull/1864)
+
+  The reconstruction technique differs for functions with equidistant frequencies
+  that are reconstructed using the function value at equidistant sampling points, and
+  for functions with arbitrary frequencies reconstructed using arbitrary sampling points.
+
+  As an example, consider the following QNode:
+
+  ```python
+  dev = qml.device("default.qubit", wires=2)
+
+  @qml.qnode(dev)
+  def circuit(x, Y, f=1.0):
+      qml.RX(f * x, wires=0)
+      qml.RY(Y[0], wires=0)
+      qml.RY(Y[1], wires=1)
+      qml.CNOT(wires=[0, 1])
+      qml.RY(3 * Y[1], wires=1)
+      return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
+  ```
+
+  It has three variational parameters overall: A scalar input `x`
+  and an array-valued input `Y` with two entries. Additionally, we can
+  tune the dependence on `x` with the frequency `f`.
+  We then can reconstruct the QNode output function with respect to `x` via
+
+  ```pycon
+  >>> x = 0.3
+  >>> Y = np.array([0.1, -0.9])
+  >>> rec = qml.fourier.reconstruct(circuit, ids="x", nums_frequency={"x": {0: 1}})(x, Y)
+  >>> rec
+  {'x': {0: <function pennylane.fourier.reconstruct._reconstruct_equ.<locals>._reconstruction(x)>}}
+  ```
+
+  As we can see, we get a nested dictionary in the format of the input `nums_frequency`
+  with functions as values. These functions are simple float-to-float callables:
+
+  ```pycon
+  >>> univariate = rec["x"][0]
+  >>> univariate(x)
+  -0.880208251507
+  ```
+
+  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 
+
+* 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) 
+  [(#1829)](https://github.com/PennyLaneAI/pennylane/pull/1829)
 
 * Added density matrix initialization gate for mixed state simulation. [(#1686)](https://github.com/PennyLaneAI/pennylane/issues/1686)
 
 <h3>Improvements</h3>
 
+
 * 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)
@@ -97,10 +173,10 @@
   been moved to the base `Operator` class.
   [(#1873)](https://github.com/PennyLaneAI/pennylane/pull/1873)
 
-* `DiagonalOperation` is not a separate subclass any more. 
-  [(#1889)](https://github.com/PennyLaneAI/pennylane/pull/1889) 
+* `DiagonalOperation` is not a separate subclass any more.
+  [(#1889)](https://github.com/PennyLaneAI/pennylane/pull/1889)
 
-  Instead, devices can check for the diagonal 
+  Instead, devices can check for the diagonal
   property using attributes:
 
   ``` python
@@ -108,7 +184,7 @@
 
   if op in diagonal_in_z_basis:
       # do something
-  ``` 
+  ```
 
 <h3>Deprecations</h3>
 
@@ -123,12 +199,12 @@
 * `ExpvalCost` now returns corrects results shape when `optimize=True` with 
   shots batch.
   [(#1897)](https://github.com/PennyLaneAI/pennylane/pull/1897)
-  
+
 * `qml.circuit_drawer.MPLDrawer` was slightly modified to work with
   matplotlib version 3.5.
   [(#1899)](https://github.com/PennyLaneAI/pennylane/pull/1899)
 
-* `qml.CSWAP` and `qml.CRot` now define `control_wires`, and `qml.SWAP` 
+* `qml.CSWAP` and `qml.CRot` now define `control_wires`, and `qml.SWAP`
   returns the default empty wires object.
   [(#1830)](https://github.com/PennyLaneAI/pennylane/pull/1830)
 
diff --git a/pennylane/fourier/__init__.py b/pennylane/fourier/__init__.py
index 8afab5000..f6fe4a9bd 100644
--- a/pennylane/fourier/__init__.py
+++ b/pennylane/fourier/__init__.py
@@ -17,6 +17,7 @@ import warnings
 from .coefficients import coefficients
 from .circuit_spectrum import circuit_spectrum
 from .qnode_spectrum import qnode_spectrum
+from .reconstruct import reconstruct
 from .utils import join_spectra, get_spectrum
 
 
diff --git a/pennylane/fourier/reconstruct.py b/pennylane/fourier/reconstruct.py
new file mode 100644
index 000000000..29290c6be
--- /dev/null
+++ b/pennylane/fourier/reconstruct.py
@@ -0,0 +1,666 @@
+# 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.
+"""Contains a function that computes the fourier series of
+a quantum expectation value."""
+from functools import wraps
+from inspect import signature
+import warnings
+
+import numpy as np
+from autoray import numpy as anp
+import pennylane as qml
+
+
+def _reconstruct_equ(fun, num_frequency, x0=None, f0=None, interface=None):
+    r"""Reconstruct a univariate Fourier series with consecutive integer
+    frequencies, using trigonometric interpolation and equidistant shifts.
+
+    This technique is based on
+    `Dirichlet kernels <https://en.wikipedia.org/wiki/Dirichlet_kernel>`_, see
+    `Vidal and Theis (2018) <https://arxiv.org/abs/1812.06323>`_ or
+    `Wierichs et al. (2021) <https://arxiv.org/abs/2107.12390>`_.
+
+    Args:
+        fun (callable): Univariate finite Fourier series to reconstruct.
+            It must have signature ``float -> float`` .
+        num_frequency (int): Number of integer frequencies in ``fun``.
+            All integer frequencies below ``num_frequency`` are assumed
+            to be present in ``fun`` as well; if they are not, the output
+            is correct put the reconstruction could have been performed
+            with fewer evaluations of ``fun`` .
+        x0 (float): Center to which to shift the reconstruction.
+            The points at which ``fun`` is evaluated are *not* affected
+            by ``x0`` .
+        f0 (float): Value of ``fun`` at zero; Providing ``f0`` saves one
+            evaluation of ``fun``.
+        interface (str): Which auto-differentiation framework to use as
+            interface. This determines in which interface the output
+            reconstructed function is intended to be used.
+
+    Returns:
+        callable: Reconstructed Fourier series with ``num_frequency`` frequencies.
+        This function is a purely classical function. Furthermore, it is fully
+        differentiable.
+    """
+    if not abs(int(num_frequency)) == num_frequency:
+        raise ValueError(f"num_frequency must be a non-negative integer, got {num_frequency}")
+
+    a = (num_frequency + 0.5) / np.pi
+    b = 0.5 / np.pi
+
+    shifts_pos = qml.math.arange(1, num_frequency + 1) / a
+    shifts_neg = -shifts_pos[::-1]
+    shifts = qml.math.concatenate([shifts_neg, [0.0], shifts_pos])
+    shifts = anp.asarray(shifts, like=interface)
+    f0 = fun(0.0) if f0 is None else f0
+    evals = (
+        list(map(fun, shifts[:num_frequency])) + [f0] + list(map(fun, shifts[num_frequency + 1 :]))
+    )
+    evals = anp.asarray(evals, like=interface)
+
+    x0 = anp.asarray(np.float64(0.0), like=interface) if x0 is None else x0
+
+    def _reconstruction(x):
+        """Univariate reconstruction based on equidistant shifts and Dirichlet kernels.
+        The derivative at of ``sinc`` are not well-implemented in TensorFlow and Autograd,
+        use the Fourier transform reconstruction if this derivative is needed.
+        """
+        _x = x - x0 - shifts
+        return qml.math.tensordot(
+            qml.math.sinc(a * _x) / qml.math.sinc(b * _x),
+            evals,
+            axes=[[0], [0]],
+        )
+
+    return _reconstruction
+
+
+_warn_text_f0_ignored = (
+    "The provided value of the function at zero will be ignored due to the "
+    "provided shift values. This may lead to additional evaluations of the "
+    "function to be reconstructed."
+)
+
+
+def _reconstruct_gen(fun, spectrum, shifts=None, x0=None, f0=None, interface=None):
+    r"""Reconstruct a univariate (real-valued) Fourier series with given spectrum.
+
+    Args:
+        fun (callable): Univariate finite Fourier series to reconstruct.
+            It must have signature ``float -> float`` .
+        spectrum (Collection): Frequency spectrum of the Fourier series;
+            non-positive frequencies are ignored.
+        shifts (Sequence): Shift angles at which to evaluate ``fun`` for the reconstruction.
+            Chosen equidistantly within the interval :math:`[0, 2\pi/f_\text{max}]`
+            if ``shifts=None`` , where :math:`f_\text{max}` is the biggest
+            frequency in ``spectrum``.
+        x0 (float): Center to which to shift the reconstruction.
+            The points at which ``fun`` is evaluated are *not* affected
+            by ``x0`` .
+        f0 (float): Value of ``fun`` at zero; If :math:`0` is among the ``shifts``
+            and ``f0`` is provided, one evaluation of ``fun`` is saved.
+        interface (str): Which auto-differentiation framework to use as
+            interface. This determines in which interface the output
+            reconstructed function is intended to be used.
+
+    Returns:
+        callable: Reconstructed Fourier series with :math:`R` frequencies in ``spectrum`` .
+        This function is a purely classical function. Furthermore, it is fully differentiable.
+    """
+    # pylint: disable=unused-argument, too-many-arguments
+
+    have_f0 = f0 is not None
+    have_shifts = shifts is not None
+
+    spectrum = anp.asarray(spectrum, like=interface)
+    spectrum = spectrum[spectrum > 0]
+    f_max = qml.math.max(spectrum)
+
+    # If no shifts are provided, choose equidistant ones
+    if not have_shifts:
+        R = qml.math.shape(spectrum)[0]
+        shifts = qml.math.arange(-R, R + 1) * 2 * np.pi / (f_max * (2 * R + 1)) * R
+        zero_idx = R
+        need_f0 = True
+    elif have_f0:
+        zero_idx = qml.math.where(qml.math.isclose(shifts, qml.math.zeros_like(shifts[0])))
+        zero_idx = zero_idx[0][0] if (len(zero_idx) > 0 and len(zero_idx[0]) > 0) else None
+        need_f0 = zero_idx is not None
+
+    # Take care of shifts close to zero if f0 was provided
+    if have_f0 and need_f0:
+        # Only one shift may be zero at a time
+        shifts = qml.math.concatenate(
+            [shifts[zero_idx : zero_idx + 1], shifts[:zero_idx], shifts[zero_idx + 1 :]]
+        )
+        shifts = anp.asarray(shifts, like=interface)
+        evals = anp.asarray([f0] + list(map(fun, shifts[1:])), like=interface)
+    else:
+        shifts = anp.asarray(shifts, like=interface)
+        if have_f0 and not need_f0:
+            warnings.warn(_warn_text_f0_ignored)
+        evals = anp.asarray(list(map(fun, shifts)), like=interface)
+
+    L = len(shifts)
+    # Construct the coefficient matrix case by case
+    C1 = qml.math.ones((L, 1))
+    C2 = qml.math.cos(qml.math.tensordot(shifts, spectrum, axes=0))
+    C3 = qml.math.sin(qml.math.tensordot(shifts, spectrum, axes=0))
+    C = qml.math.hstack([C1, C2, C3])
+
+    # Solve the system of linear equations
+    cond = qml.math.linalg.cond(C)
+    if cond > 1e8:
+        warnings.warn(
+            f"The condition number of the Fourier transform matrix is very large: {cond}.",
+            UserWarning,
+        )
+    W = qml.math.linalg.solve(C, evals)
+
+    # Extract the Fourier coefficients
+    R = (L - 1) // 2
+    a0 = W[0]
+    a = anp.asarray(W[1 : R + 1], like=interface)
+    b = anp.asarray(W[R + 1 :], like=interface)
+
+    x0 = anp.asarray(np.float64(0.0), like=interface) if x0 is None else x0
+    # Construct the Fourier series
+    def _reconstruction(x):
+        """Univariate reconstruction based on arbitrary shifts."""
+        x = x - x0
+        return (
+            a0
+            + qml.math.tensordot(qml.math.cos(spectrum * x), a, axes=[[0], [0]])
+            + qml.math.tensordot(qml.math.sin(spectrum * x), b, axes=[[0], [0]])
+        )
+
+    return _reconstruction
+
+
+def _parse_ids(ids, info_dict):
+    """Parse different formats of ``ids`` into the right dictionary format,
+    potentially using the information in ``info_dict`` to complete it.
+    """
+    if ids is None:
+        # Infer all id information from info_dict
+        return {outer_key: inner_dict.keys() for outer_key, inner_dict in info_dict.items()}
+    if isinstance(ids, str):
+        # ids only provides a single argument name but no parameter indices
+        return {ids: info_dict[ids].keys()}
+    if not isinstance(ids, dict):
+        # ids only provides argument names but no parameter indices
+        return {_id: info_dict[_id].keys() for _id in ids}
+
+    return ids
+
+
+def _parse_shifts(shifts, R, arg_name, par_idx, atol, need_f0):
+    """Processes shifts for a single reconstruction and determines
+    wheter the function at the reconstruction point, ``f0`` will be
+    needed.
+    """
+    # pylint: disable=too-many-arguments
+    _shifts = shifts.get(arg_name)
+    if _shifts is not None:
+        _shifts = _shifts.get(par_idx)
+    if _shifts is not None:
+        # Check whether the _shifts have the correct size
+        if len(_shifts) != 2 * R + 1:
+            raise ValueError(
+                f"The number of provided shifts ({len(_shifts)}) does not fit to the "
+                f"number of frequencies (2R+1={2*R+1}) for parameter {par_idx} in "
+                f"argument {arg_name}."
+            )
+        if any(qml.math.isclose(_shifts, qml.math.zeros_like(_shifts), rtol=0, atol=atol)):
+            # If 0 is among the shifts, f0 is needed
+            return _shifts, True
+        # If 0 is not among the shifts, f0 is not needed
+        return _shifts, (False or need_f0)
+    # If no shifts are given, f0 is needed always
+    return _shifts, True
+
+
+def _prepare_jobs(ids, nums_frequency, spectra, shifts, atol):
+    r"""For inputs to reconstruct, determine how the given information yields
+    function reconstruction tasks and collect them into a dictionary ``jobs``.
+    Also determine whether the function at zero is needed.
+
+    Args:
+        ids (dict or Sequence or str): Indices for the QNode parameters with respect to which
+            the QNode should be reconstructed as a univariate function, per QNode argument.
+            Each key of the dict, entry of the list, or the single ``str`` has to be the name
+            of an argument of ``qnode`` .
+            If a ``dict`` , the values of ``ids`` have to contain the parameter indices
+            for the respective array-valued QNode argument represented by the key.
+            These indices always are tuples, i.e. ``()`` for scalar and ``(i,)`` for
+            one-dimensional arguments.
+            If a ``list`` , the parameter indices are inferred from ``nums_frequency`` if
+            given or ``spectra`` else.
+            If ``None``, all keys present in ``nums_frequency`` / ``spectra`` are considered.
+        nums_frequency (dict[dict]): Numbers of integer frequencies -- and biggest
+            frequency -- per QNode parameter. The keys have to be argument names of ``qnode``
+            and the inner dictionaries have to be mappings from parameter indices to the
+            respective integer number of frequencies. If the QNode frequencies are not contiguous
+            integers, the argument ``spectra`` should be used to save evaluations of ``qnode`` .
+            Takes precedence over ``spectra`` and leads to usage of equidistant shifts.
+        spectra (dict[dict]): Frequency spectra per QNode parameter.
+            The keys have to be argument names of ``qnode`` and the inner dictionaries have to
+            be mappings from parameter indices to the respective frequency spectrum for that
+            parameter. Ignored if ``nums_frequency!=None``.
+        shifts (dict[dict]): Shift angles for the reconstruction per QNode parameter.
+            The keys have to be argument names of ``qnode`` and the inner dictionaries have to
+            be mappings from parameter indices to the respective shift angles to be used for that
+            parameter. For :math:`R` non-zero frequencies, there must be :math:`2R+1` shifts
+            given. Ignored if ``nums_frequency!=None``.
+        atol (float): Absolute tolerance used to analyze shifts lying close to 0.
+
+    Returns:
+        dict[dict]: Indices for the QNode parameters with respect to which the QNode
+            will be reconstructed. Cast to the dictionary structure explained above.
+            If the input ``ids`` was a dictionary, it is returned unmodified.
+        callable: The reconstruction method to use, one out of two internal methods.
+        dict[dict[dict]]: Keyword arguments for the reconstruction method specifying
+            how to carry out the reconstruction. The outer-most keys are QNode argument
+            names, the middle keys are parameter indices like the inner keys of
+            ``nums_frequency`` or ``spectra`` and the inner-most dictionary contains the
+            keyword arguments, i.e. the keys are keyword argument names for the
+            reconstruction method
+        bool: Whether any of the reconstruction jobs will require the evaluation
+            of the function at the position of reconstruction itself.
+    """
+    if nums_frequency is None:
+        if spectra is None:
+            raise ValueError("Either nums_frequency or spectra must be given.")
+
+        ids = _parse_ids(ids, spectra)
+
+        if shifts is None:
+            shifts = {}
+
+        need_f0 = False
+        recon_fn = _reconstruct_gen
+
+        jobs = {}
+
+        # If no shifts are provided, compute them
+        for arg_name, inner_dict in ids.items():
+            _jobs = {}
+
+            for par_idx in inner_dict:
+
+                # Determine spectrum and number of frequencies, discounting for 0
+                _spectrum = spectra[arg_name][par_idx]
+                R = len(_spectrum) - 1
+                _shifts, need_f0 = _parse_shifts(shifts, R, arg_name, par_idx, atol, need_f0)
+
+                # Store job
+                if R > 0:
+                    _jobs[par_idx] = {"shifts": _shifts, "spectrum": _spectrum}
+                else:
+                    # R=0 belongs to a constant function
+                    _jobs[par_idx] = None
+
+            jobs[arg_name] = _jobs
+
+    else:
+        jobs = {}
+        need_f0 = True
+
+        ids = _parse_ids(ids, nums_frequency)
+
+        recon_fn = _reconstruct_equ
+
+        for arg_name, inner_dict in ids.items():
+            _jobs = {}
+
+            for par_idx in inner_dict:
+                _num_frequency = nums_frequency[arg_name][par_idx]
+                _jobs[par_idx] = {"num_frequency": _num_frequency} if _num_frequency > 0 else None
+
+            jobs[arg_name] = _jobs
+
+    return ids, recon_fn, jobs, need_f0
+
+
+def reconstruct(qnode, ids=None, nums_frequency=None, spectra=None, shifts=None):
+    r"""Reconstruct an expectation value QNode along a single parameter direction.
+    This means we restrict the QNode to vary only one parameter, a univariate restriction.
+    For common quantum gates, such restrictions are finite Fourier series with known
+    frequency spectra. Thus they may be reconstructed using Dirichlet kernels or
+    a non-uniform Fourier transform.
+
+    Args:
+        qnode (pennylane.QNode): Quantum node to be reconstructed, representing a
+            circuit that outputs an expectation value.
+        ids (dict or Sequence or str): Indices for the QNode parameters with respect to which
+            the QNode should be reconstructed as a univariate function, per QNode argument.
+            Each key of the dict, entry of the list, or the single ``str`` has to be the name
+            of an argument of ``qnode`` .
+            If a ``dict`` , the values of ``ids`` have to contain the parameter indices
+            for the respective array-valued QNode argument represented by the key.
+            These indices always are tuples, i.e., ``()`` for scalar and ``(i,)`` for
+            one-dimensional arguments.
+            If a ``list`` , the parameter indices are inferred from ``nums_frequency`` if
+            given or ``spectra`` else.
+            If ``None``, all keys present in ``nums_frequency`` / ``spectra`` are considered.
+        nums_frequency (dict[dict]): Numbers of integer frequencies -- and biggest
+            frequency -- per QNode parameter. The keys have to be argument names of ``qnode``
+            and the inner dictionaries have to be mappings from parameter indices to the
+            respective integer number of frequencies. If the QNode frequencies are not contiguous
+            integers, the argument ``spectra`` should be used to save evaluations of ``qnode`` .
+            Takes precedence over ``spectra`` and leads to usage of equidistant shifts.
+        spectra (dict[dict]): Frequency spectra per QNode parameter.
+            The keys have to be argument names of ``qnode`` and the inner dictionaries have to
+            be mappings from parameter indices to the respective frequency spectrum for that
+            parameter. Ignored if ``nums_frequency!=None``.
+        shifts (dict[dict]): Shift angles for the reconstruction per QNode parameter.
+            The keys have to be argument names of ``qnode`` and the inner dictionaries have to
+            be mappings from parameter indices to the respective shift angles to be used for that
+            parameter. For :math:`R` non-zero frequencies, there must be :math:`2R+1` shifts
+            given. Ignored if ``nums_frequency!=None``.
+
+    Returns:
+        function: Function which accepts the same arguments as the QNode.
+        When called, this function will return a dictionary of dictionaries,
+        formatted like ``nums_frequency`` or ``spectra`` ,
+        that contains the univariate reconstructions per QNode parameter.
+
+    For each provided ``id`` in ``ids``, the QNode is restricted to varying the single QNode
+    parameter corresponding to the ``id`` . This univariate function is then reconstructed
+    via a Fourier transform or Dirichlet kernels, depending on the provided input.
+    Either the frequency ``spectra`` of the QNode with respect to its input parameters or
+    the numbers of frequencies, ``nums_frequency`` , per parameter must be provided.
+
+    For quantum-circuit specific details, we refer the reader to
+    `Vidal and Theis (2018) <https://arxiv.org/abs/1812.06323>`__ ,
+    `Vidal and Theis (2020) <https://www.frontiersin.org/articles/10.3389/fphy.2020.00297/full>`__ ,
+    `Schuld, Sweke and Meyer (2021) <https://journals.aps.org/pra/abstract/10.1103/PhysRevA.103.032430>`__ ,
+    and
+    `Wierichs, Izaac, Wang and Lin (2021) <https://arxiv.org/abs/2107.12390>`__ .
+    An introduction to the concept of quantum circuits as Fourier series can also be found in
+    the
+    `Quantum models as Fourier series <https://pennylane.ai/qml/demos/tutorial_expressivity_fourier_series.html>`__
+    and
+    `General parameter-shift rules <https://pennylane.ai/qml/demos/tutorial_general_parshift.html>`__
+    demos as well as the
+    :mod:`qml.fourier <pennylane.fourier>` module docstring.
+
+    **Example**
+
+    Consider the following QNode:
+
+    .. code-block:: python
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev)
+        def circuit(x, Y):
+            qml.RX(x, wires=0)
+            qml.RY(Y[0], wires=0)
+            qml.RY(Y[1], wires=1)
+            qml.CNOT(wires=[0, 1])
+            qml.RY(5*  Y[1], wires=1)
+            return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
+
+        x = 0.4
+        Y = np.array([1.9, -0.5])
+        f = 2.3
+
+        circuit_value = circuit(x, Y)
+
+    It has three variational parameters ``x`` (a scalar) and two entries of ``Y``
+    (an array-like).
+    A reconstruction job could then be with respect to the two entries of ``Y``,
+    which enter the circuit with one and six integer frequencies, respectively
+    (see the additional examples below for details on how to obtain the frequency
+    spectrum if it is not known):
+
+    >>> nums_frequency = {"Y": {(0,): 1, (1,): 6}}
+    >>> with qml.Tracker(circuit.device) as tracker:
+    ...     rec = qml.fourier.reconstruct(circuit, {"Y": [(0,), (1,)]}, nums_frequency)(x, Y)
+    >>> rec.keys()
+    dict_keys(['Y'])
+    >>> print(*rec["Y"].items(), sep="\n")
+    ((0,), <function _reconstruct_equ.<locals>._reconstruction at 0x7fbd685aee50>)
+    ((1,), <function _reconstruct_equ.<locals>._reconstruction at 0x7fbd6866eee0>)
+    >>> recon_Y0 = rec["Y"][(0,)]
+    >>> 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)
+    True
+
+    We successfully reconstructed the dependence on the two entries of ``Y`` ,
+    keeping ``x`` and the respective other entry in ``Y`` at their initial values.
+    Let us also see how many executions of the device were used to obtain the
+    reconstructions:
+
+    >>> tracker.totals
+    {'executions': 15}
+
+    The example above used that we already knew the frequency spectra of the
+    QNode of interest. However, this is in general not the case and we may need
+    to compute the spectrum first. This can be done with
+    :func:`.fourier.qnode_spectrum` :
+
+    >>> spectra = qml.fourier.qnode_spectrum(circuit)(x, Y)
+    >>> spectra.keys()
+    dict_keys(['x', 'Y'])
+    >>> spectra["x"]
+    {(): [-1.0, 0.0, 1.0]}
+    >>> 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])
+
+    For more detailed explanations, usage details and additional examples, see
+    the usage details section below.
+
+    .. UsageDetails::
+
+        **Input formatting**
+
+        As described briefly above, the essential inputs to ``reconstruct`` that provide information
+        about the QNode are given as dictionaries of dictionaries, where the outer keys reference
+        the argument names of ``qnode`` and the inner keys reference the parameter indices within
+        each array-valued QNode argument. These parameter indices always are tuples, so that
+        for scalar-valued QNode parameters, the parameter index is ``()`` by convention and the
+        ``i`` -th parameter of a one-dimensional array can be accessed via ``(i,)`` .
+        For example, providing ``nums_frequency``
+
+        - for a scalar argument: ``nums_frequency = {"x": {(): 4}}``
+        - for a one-dimensional argument: ``nums_frequency = {"Y": {(0,): 2, (1,): 9, (4,): 1}}``
+        - for a three-dimensional argument: ``nums_frequency = {"Z": {(0, 2, 5): 2, (1, 1, 4): 1}}``
+
+        This applies to ``nums_frequency`` , ``spectra`` , and ``shifts`` .
+
+        Note that the information provided in ``nums_frequency`` / ``spectra`` is essential for
+        the correctness of the reconstruction.
+
+        On the other hand, the input format for ``ids`` is flexible and allows a collection of
+        parameter indices for each QNode argument name (as a ``dict`` ), a collection of argument
+        names (as a ``list``, ``set``, ``tuple`` or similar), or a single argument name
+        (as a ``str`` ) to be defined. For ``ids=None`` , all argument names contained in
+        ``nums_frequency`` -- or ``spectra`` if ``nums_frequency`` is not used -- are considered.
+        For inputs that do not specify parameter indices per QNode argument name (all formats but
+        ``dict`` ), these parameter indices are inferred from ``nums_frequency`` / ``spectra`` .
+
+        **Reconstruction cost**
+
+        The reconstruction cost -- in terms of calls to ``qnode`` -- depend on the number of
+        frequencies given via ``nums_frequency`` or ``spectra`` . A univariate reconstruction
+        for :math:`R` frequencies takes :math:`2R+1` evaluations. If multiple univariate
+        reconstructions are performed at the same point with various numbers of frequencies
+        :math:`R_k` , the cost are :math:`1+2\sum_k R_k` if the shift :math:`0` is used in all
+        of them. This is in particular the case if ``nums_frequency`` or ``spectra`` with
+        ``shifts=None`` is used.
+
+        If the number of frequencies is too large or the given frequency spectrum contains
+        more than the spectrum of ``qnode`` , the reconstruction is performed suboptimally
+        but remains correct.
+        For integer-valued spectra with gaps, the equidistant reconstruction is thus suboptimal
+        and the non-equidistant version method be used (also see the examples below).
+
+        **Numerical stability**
+
+        In general, the reconstruction with equidistant shifts for equidistant frequencies
+        (used if ``nums_frequency`` is provided) is more stable numerically than the more
+        general Fourier reconstruction (used if ``nums_frequency=None`` ).
+        If the system of equations to be solved in the Fourier transform is
+        ill-conditioned, a warning is raised as the output might become unstable.
+        Examples for this are shift values or frequencies that lie very close to each other.
+
+        **Differentiability**
+
+        The returned scalar functions are differentiable in all interfaces with respect
+        to their scalar input variable. They expect these inputs to be in the same
+        interface as the one used by the QNode. More advanced differentiability, for example
+        of the reconstructions with respect to QNode properties, is not supported
+        reliably yet.
+
+        .. warning::
+
+            When using ``TensorFlow`` or ``Autograd`` *and* ``nums_frequency`` ,
+            the reconstructed functions are not differentiable at the point of
+            reconstruction. One workaround for this is to use ``spectra`` as
+            input instead and to thereby use the Fourier transform instead of
+            Dirichlet kernels. Alternatively, the original QNode evaluation can
+            be used.
+
+        **More examples**
+
+        Consider the QNode from the example above, now with an additional, tunable frequency
+        ``f`` for the Pauli-X rotation that is controlled by ``x`` :
+
+        .. code-block:: python
+
+            @qml.qnode(dev)
+            def circuit(x, Y, f=1.0):
+                qml.RX(f * x, wires=0)
+                qml.RY(Y[0], wires=0)
+                qml.RY(Y[1], wires=1)
+                qml.CNOT(wires=[0, 1])
+                qml.RY(5*  Y[1], wires=1)
+                return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
+
+            f = 2.3
+
+            circuit_value = circuit(x, Y)
+
+
+        We repeat the reconstruction job for the dependence on ``Y[1]`` .
+        Note that even though information about ``Y[0]`` is contained in ``nums_frequency`` ,
+        ``ids`` determines which reconstructions are performed.
+
+        >>> with qml.Tracker(circuit.device) as tracker:
+        ...     rec = qml.fourier.reconstruct(circuit, {"Y": [(1,)]}, nums_frequency)(x, Y)
+        >>> tracker.totals
+        {'executions': 13}
+
+        As expected, we required :math:`2R+1=2\cdot 6+1=13` circuit executions. However, not
+        all frequencies below :math:`f_\text{max}=6` are present in the circuit, so that
+        a reconstruction using knowledge of the full frequency spectrum will be cheaper:
+
+        >>> spectra = {"Y": {(1,): [0., 1., 4., 5., 6.]}}
+        >>> with tracker:
+        ...     rec = qml.fourier.reconstruct(circuit, {"Y": [(1,)]}, None, spectra)(x, Y)
+        >>> tracker.totals
+        {'executions': 9}
+
+        We again obtain the full univariate dependence on ``Y[1]`` but with considerably
+        fewer executions on the quantum device.
+        Once we obtained the classical function that describes the dependence, no
+        additional circuit evaluations are performed:
+
+        >>> with tracker:
+        ...     for Y1 in np.arange(-np.pi, np.pi, 20):
+        ...         rec["Y"][(1,)](-2.1)
+        >>> tracker.totals
+        {}
+
+        If we want to reconstruct the dependence of ``circuit`` on ``x`` , we cannot use
+        ``nums_frequency`` if ``f`` is not an integer. One could rescale ``x`` to obtain
+        the frequency :math:`1` again, or directly use ``spectra`` . We will combine the
+        latter with another reconstruction with respect to ``Y[0]`` :
+
+        >>> spectra = {"x": {(): [0., f]}, "Y": {(0,): [0., 1.]}}
+        >>> with tracker:
+        ...     rec = qml.fourier.reconstruct(circuit, None, None, spectra)(x, Y, f=f)
+        >>> tracker.totals
+        {'executions': 5}
+        >>> recon_x = rec["x"][()]
+        >>> np.isclose(recon_x(x+0.5), circuit(x+0.5, Y, f=f)
+        True
+
+        Note that by convention, the parameter index for a scalar variable is ``()`` and
+        that the frequency :math:`0` always needs to be included in the spectra.
+        Furthermore, we here skipped the input ``ids`` so that the reconstruction
+        was performed for all keys in ``spectra`` .
+        The reconstruction with a single non-zero frequency
+        costs three evaluations of ``circuit`` for each, ``x`` and ``Y[0]`` . Performing
+        both reconstructions at the same position allowed us to save one of the
+        evaluations and reduce the number of calls to :math:`5`.
+
+    """
+    # pylint: disable=cell-var-from-loop, unused-argument
+
+    atol = 1e-8
+    ids, recon_fn, jobs, need_f0 = _prepare_jobs(ids, nums_frequency, spectra, shifts, atol)
+    arg_names = list(signature(qnode.func).parameters.keys())
+    arg_idx_from_names = {arg_name: i for i, arg_name in enumerate(arg_names)}
+
+    @wraps(qnode)
+    def wrapper(*args, **kwargs):
+        if need_f0:
+            f0 = qnode(*args, **kwargs)
+        else:
+            f0 = None
+
+        interface = qml.math.get_interface(args[0])
+
+        def constant_fn(x):
+            """Univariate reconstruction of a constant Fourier series."""
+            return f0
+
+        # Carry out the reconstruction jobs
+        reconstructions = {}
+        for arg_name, inner_dict in jobs.items():
+            _reconstructions = {}
+            arg_idx = arg_idx_from_names[arg_name]
+
+            for par_idx, job in inner_dict.items():
+                if job is None:
+                    _reconstructions[par_idx] = constant_fn
+                else:
+                    if len(qml.math.shape(args[arg_idx])) == 0:
+                        shift_vec = qml.math.ones_like(args[arg_idx])
+                        x0 = args[arg_idx]
+                    else:
+                        shift_vec = qml.math.zeros_like(args[arg_idx])
+                        shift_vec = qml.math.scatter_element_add(shift_vec, par_idx, 1.0)
+                        x0 = args[arg_idx][par_idx]
+
+                    def _univariate_fn(x):
+                        new_arg = args[arg_idx] + shift_vec * x
+                        new_args = args[:arg_idx] + (new_arg,) + args[arg_idx + 1 :]
+                        return qnode(*new_args, **kwargs)
+
+                    _reconstructions[par_idx] = recon_fn(
+                        _univariate_fn, **job, x0=x0, f0=f0, interface=interface
+                    )
+
+            reconstructions[arg_name] = _reconstructions
+
+        return reconstructions
+
+    return wrapper
diff --git a/pennylane/math/multi_dispatch.py b/pennylane/math/multi_dispatch.py
index 9cd0878e3..5723c46e3 100644
--- a/pennylane/math/multi_dispatch.py
+++ b/pennylane/math/multi_dispatch.py
@@ -235,7 +235,7 @@ def dot(tensor1, tensor2):
         if x.ndim <= 2 and y.ndim <= 2:
             return x @ y
 
-        return np.tensordot(x, y, dims=[[-1], [-2]], like=interface)
+        return np.tensordot(x, y, axes=[[-1], [-2]], like=interface)
 
     if interface == "tensorflow":
         if len(np.shape(x)) == 0 and len(np.shape(y)) == 0:
@@ -365,28 +365,73 @@ def stack(values, axis=0):
     return np.stack(values, axis=axis, like=interface)
 
 
-def where(condition, x, y):
-    """Returns elements chosen from x or y depending on a boolean tensor condition.
+def where(condition, x=None, y=None):
+    """Returns elements chosen from x or y depending on a boolean tensor condition,
+    or the indices of entries satisfying the condition.
 
     The input tensors ``condition``, ``x``, and ``y`` must all be broadcastable to the same shape.
 
     Args:
-        condition (tensor_like[bool]): A boolean tensor. Where True, elements from
-            ``x`` will be chosen, otherwise ``y``.
+        condition (tensor_like[bool]): A boolean tensor. Where ``True`` , elements from
+            ``x`` will be chosen, otherwise ``y``. If ``x`` and ``y`` are ``None`` the
+            indices where ``condition==True`` holds will be returned.
         x (tensor_like): values from which to choose if the condition evaluates to ``True``
         y (tensor_like): values from which to choose if the condition evaluates to ``False``
 
     Returns:
-        tensor_like: A tensor with elements from ``x`` where the condition is ``True``, and
-        ``y`` otherwise. The output tensor has the same shape as the input tensors.
+        tensor_like or tuple[tensor_like]: If ``x is None`` and ``y is None``, a tensor
+        or tuple of tensors with the indices where ``condition`` is ``True`` .
+        Else, a tensor with elements from ``x`` where the ``condition`` is ``True``,
+        and ``y`` otherwise. In this case, the output tensor has the same shape as
+        the input tensors.
 
-    **Example**
+    **Example with three arguments**
 
     >>> a = torch.tensor([0.6, 0.23, 0.7, 1.5, 1.7], requires_grad=True)
     >>> b = torch.tensor([-1., -2., -3., -4., -5.], requires_grad=True)
     >>> math.where(a < 1, a, b)
     tensor([ 0.6000,  0.2300,  0.7000, -4.0000, -5.0000], grad_fn=<SWhereBackward>)
+
+    .. warning::
+
+        The output format for ``x=None`` and ``y=None`` follows the respective
+        interface and differs between TensorFlow and all other interfaces:
+        For TensorFlow, the output is a tensor with shape
+        ``(num_true, len(condition.shape))`` where ``num_true`` is the number
+        of entries in ``condition`` that are ``True`` .
+        The entry at position ``(i, j)`` is the ``j`` th entry of the ``i`` th
+        index.
+        For all other interfaces, the output is a tuple of tensor-like objects,
+        with the ``j`` th object indicating the ``j`` th entries of all indices.
+        Also see the examples below.
+
+    **Example with single argument**
+
+    For Torch, Autograd, JAX and NumPy, the output formatting is as follows:
+
+    >>> a = [[0.6, 0.23, 1.7],[1.5, 0.7, -0.2]]
+    >>> math.where(torch.tensor(a) < 1)
+    (tensor([0, 0, 1, 1]), tensor([0, 1, 1, 2]))
+
+    This is not a single tensor-like object but corresponds to the shape
+    ``(2, 4)`` . For TensorFlow, on the other hand:
+
+    >>> math.where(tf.constant(a) < 1)
+    tf.Tensor(
+    [[0 0]
+     [0 1]
+     [1 1]
+     [1 2]], shape=(4, 2), dtype=int64)
+
+    As we can see, the dimensions are swapped and the output is a single Tensor.
+    Note that the number of dimensions of the output does *not* depend on the input
+    shape, it is always two-dimensional.
+
     """
+    if x is None and y is None:
+        interface = _multi_dispatch([condition])
+        return np.where(condition, like=interface)
+
     return np.where(condition, x, y, like=_multi_dispatch([condition, x, y]))
 
 
diff --git a/pennylane/math/single_dispatch.py b/pennylane/math/single_dispatch.py
index f6dc9c5e4..124f8245d 100644
--- a/pennylane/math/single_dispatch.py
+++ b/pennylane/math/single_dispatch.py
@@ -124,11 +124,17 @@ def _scatter_element_add_autograd(tensor, index, value):
     """In-place addition of a multidimensional value over various
     indices of a tensor. Since Autograd doesn't support indexing
     assignment, we have to be clever and use ravel_multi_index."""
+    pnp = _i("qml").numpy
     size = tensor.size
-    flat_index = _i("qml").numpy.ravel_multi_index(index, tensor.shape)
+    flat_index = pnp.ravel_multi_index(index, tensor.shape)
+    if pnp.isscalar(flat_index):
+        flat_index = [flat_index]
+    if pnp.isscalar(value) or len(pnp.shape(value)) == 0:
+        value = [value]
     t = [0] * size
-    t[flat_index] = value
-    return tensor + _i("qml").numpy.array(t).reshape(tensor.shape)
+    for _id, val in zip(flat_index, value):
+        t[_id] = val
+    return tensor + pnp.array(t).reshape(tensor.shape)
 
 
 ar.register_function("autograd", "scatter_element_add", _scatter_element_add_autograd)
@@ -158,6 +164,8 @@ ar.autoray._SUBMODULE_ALIASES["tensorflow", "arctan2"] = "tensorflow.math"
 ar.autoray._SUBMODULE_ALIASES["tensorflow", "diag"] = "tensorflow.linalg"
 ar.autoray._SUBMODULE_ALIASES["tensorflow", "kron"] = "tensorflow.experimental.numpy"
 ar.autoray._SUBMODULE_ALIASES["tensorflow", "moveaxis"] = "tensorflow.experimental.numpy"
+ar.autoray._SUBMODULE_ALIASES["tensorflow", "sinc"] = "tensorflow.experimental.numpy"
+ar.autoray._SUBMODULE_ALIASES["tensorflow", "isclose"] = "tensorflow.experimental.numpy"
 
 
 ar.autoray._FUNC_ALIASES["tensorflow", "arcsin"] = "asin"
@@ -259,6 +267,8 @@ def _scatter_element_add_tf(tensor, index, value):
     indices of a tensor."""
     import tensorflow as tf
 
+    if not isinstance(index[0], int):
+        index = tuple(zip(*index))
     indices = tf.expand_dims(index, 0)
     value = tf.cast(tf.expand_dims(value, 0), tensor.dtype)
     return tf.tensor_scatter_nd_add(tensor, indices, value)
@@ -412,6 +422,16 @@ def _sort_torch(tensor):
 ar.register_function("torch", "sort", _sort_torch)
 
 
+def _tensordot_torch(tensor1, tensor2, axes):
+    torch = _i("torch")
+    if not semantic_version.match(">=1.10.0", torch.__version__) and axes == 0:
+        return torch.outer(tensor1, tensor2)
+    return torch.tensordot(tensor1, tensor2, axes)
+
+
+ar.register_function("torch", "tensordot", _tensordot_torch)
+
+
 # -------------------------------- JAX --------------------------------- #
 
 


Test patch:
diff --git a/tests/fourier/test_reconstruct.py b/tests/fourier/test_reconstruct.py
new file mode 100644
index 000000000..ac2c4843c
--- /dev/null
+++ b/tests/fourier/test_reconstruct.py
@@ -0,0 +1,1075 @@
+# 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.
+"""
+Tests for the Fourier reconstruction transform.
+"""
+import pytest
+from inspect import signature
+from itertools import chain, combinations
+from functools import reduce
+import numpy as np
+import pennylane as qml
+from pennylane import numpy as pnp
+from pennylane.fourier.reconstruct import (
+    _reconstruct_equ,
+    _reconstruct_gen,
+    _prepare_jobs,
+    reconstruct,
+)
+from pennylane.fourier.utils import join_spectra
+
+dev_0 = qml.device("default.qubit", wires=1)
+
+
+class Lambda:
+    def __init__(self, fun):
+        self.fun = fun
+
+    def __call__(self, *args, **kwargs):
+        return self.fun(*args, **kwargs)
+
+
+@qml.qnode(dev_0)
+def dummy_qnode(x):
+    qml.RX(x, wires=0)
+    return qml.expval(qml.PauliZ(0))
+
+
+def get_RX_circuit(scales):
+    """Generate a circuit with Pauli-X rotation gates with ``f*x``
+    as argument where the ``f`` s are stored in ``scales`` ."""
+
+    @qml.qnode(dev_0)
+    def circuit(x):
+        for f in scales:
+            qml.RX(f * x, wires=0)
+        return qml.expval(qml.PauliZ(0))
+
+    return circuit
+
+
+def fun_close(fun1, fun2, zero=None, tol=1e-5, samples=100):
+    X = np.linspace(-np.pi, np.pi, samples)
+    if zero is not None:
+        X = qml.math.convert_like(X, zero)
+
+    for x in X:
+        if not np.isclose(fun1(x), fun2(x), atol=tol, rtol=0):
+            print(f"{fun1(x)} != {fun2(x)} at {x}")
+            return False
+    return True
+
+
+class TestErrors:
+    """Test that errors are raised e.g. for invalid inputs."""
+
+    def test_nums_frequency_and_spectra_missing(self):
+        """Tests that an error is raised if neither information about the number
+        of frequencies nor about the spectrum is given to ``reconstruct``."""
+        with pytest.raises(ValueError, match="Either nums_frequency or spectra must be given."):
+            reconstruct(dummy_qnode)
+
+    @pytest.mark.parametrize("num_frequency", [-3, -9.2, 0.999])
+    def test_num_frequency_invalid(self, num_frequency):
+        """Tests that an error is raised if ``_reconstruct_equ`` receives a
+        negative or non-integer ``num_frequency`` ."""
+        with pytest.raises(ValueError, match="num_frequency must be a non-negative integer"):
+            _reconstruct_equ(dummy_qnode, num_frequency=num_frequency)
+
+    @pytest.mark.parametrize(
+        "spectra, shifts",
+        [
+            ({"x": {(): [0.0, 1.0]}}, {"x": {(): [0.3]}}),
+            ({"x": {(): [0.0, 1.0, 2.0]}}, {"x": {(): list(range(20))}}),
+        ],
+    )
+    def test_wrong_number_of_shifts(self, spectra, shifts):
+        """Tests that an error is raised if the number of provided shifts does not match."""
+        with pytest.raises(ValueError, match="The number of provided shifts"):
+            reconstruct(dummy_qnode, spectra=spectra, shifts=shifts)
+
+
+class TestWarnings:
+    """Test that warnings are raised e.g. for an ill-conditioned
+    Fourier transform during the reconstruction."""
+
+    def test_ill_conditioned(self):
+        """Test that a warning is raised for an ill-conditioned matrix in the Fourier trafo."""
+        shifts = [-np.pi / 2 - 1e-9, -np.pi / 2, 0, np.pi / 2, np.pi / 2 + 1e-9]
+        with pytest.warns(UserWarning, match="condition number of the Fourier"):
+            _reconstruct_gen(dummy_qnode, spectrum=[1.0, 2.0], shifts=shifts)
+
+
+class TestReconstructEqu:
+    """Tests the one-dimensional reconstruction subroutine based on equidistantly
+    shifted evaluations for equidistant frequencies."""
+
+    c_funs = [
+        lambda x: 13.71 * qml.math.sin(x) - qml.math.cos(2 * x) / 30,
+        lambda x: -0.49 * qml.math.sin(3.2 * x),
+        lambda x: 0.1 * qml.math.cos(-2.1 * x) + 2.9 * qml.math.sin(4.2 * x - 1.2),
+        lambda x: qml.math.ones_like(x) * 4.01,
+        lambda x: qml.math.sum(
+            [i ** 2 * 0.1 * qml.math.sin(i * 3.921 * x - 2.7 / i) for i in range(1, 10)]
+        ),
+    ]
+
+    nums_frequency = [2, 1, 2, 0, 9]
+    base_frequencies = [1.0, 3.2, 2.1, 1.0, 3.921]
+    expected_grads = [
+        lambda x: 13.71 * qml.math.cos(x) + 2 * qml.math.sin(2 * x) / 30,
+        lambda x: -0.49 * qml.math.cos(3.2 * x) * 3.2,
+        lambda x: (-2.1) * (-0.1) * qml.math.sin(-2.1 * x)
+        + 4.2 * 2.9 * qml.math.cos(4.2 * x - 1.2),
+        lambda x: 0.0,
+        lambda x: qml.math.sum(
+            [i * 3.921 * i ** 2 * 0.1 * qml.math.cos(i * 3.921 * x - 2.7 / i) for i in range(1, 10)]
+        ),
+    ]
+
+    @pytest.mark.parametrize(
+        "fun, num_frequency, base_f", zip(c_funs, nums_frequency, base_frequencies)
+    )
+    def test_with_classical_fun(self, fun, num_frequency, base_f, mocker):
+        """Test that equidistant-frequency classical functions are
+        reconstructed correctly (via rescaling to integer frequencies)."""
+        Fun = Lambda(fun)
+        spy = mocker.spy(Fun, "fun")
+        # Convert fun to have integer frequencies
+        _fun = lambda x: Fun(x / base_f)
+        _rec = _reconstruct_equ(_fun, num_frequency)
+
+        # Convert reconstruction to have original frequencies
+        rec = lambda x: _rec(base_f * x)
+        assert spy.call_count == num_frequency * 2 + 1
+        assert fun_close(fun, rec)
+
+        # Repeat, using precomputed f0
+        f0 = _fun(0.0)
+        Fun = Lambda(fun)
+        spy = mocker.spy(Fun, "fun")
+        _rec = _reconstruct_equ(_fun, num_frequency, f0=f0)
+        # Convert reconstruction to have original frequencies
+        rec = lambda x: _rec(base_f * x)
+        assert spy.call_count == num_frequency * 2
+        assert fun_close(fun, rec)
+
+    @pytest.mark.parametrize(
+        "fun, num_frequency, base_f, expected_grad",
+        zip(c_funs, nums_frequency, base_frequencies, expected_grads),
+    )
+    def test_differentiability_autograd(self, fun, num_frequency, base_f, expected_grad):
+        """Test that the reconstruction of equidistant-frequency classical
+        functions are differentiable for Autograd input variables."""
+        # Convert fun to have integer frequencies
+        _fun = lambda x: fun(x / base_f)
+        _rec = _reconstruct_equ(_fun, num_frequency, interface="autograd")
+
+        # Convert reconstruction to have original frequencies
+        rec = lambda x: _rec(base_f * x)
+        grad = qml.grad(rec)
+
+        assert fun_close(fun, rec, zero=pnp.array(0.0, requires_grad=True))
+        assert fun_close(expected_grad, grad, zero=pnp.array(0.0, requires_grad=True))
+
+    @pytest.mark.parametrize(
+        "fun, num_frequency, base_f, expected_grad",
+        zip(c_funs, nums_frequency, base_frequencies, expected_grads),
+    )
+    def test_differentiability_jax(self, fun, num_frequency, base_f, expected_grad):
+        """Test that the reconstruction of equidistant-frequency classical
+        functions are differentiable for JAX input variables."""
+        jax = pytest.importorskip("jax")
+        from jax.config import config
+
+        config.update("jax_enable_x64", True)
+        # Convert fun to have integer frequencies
+        _fun = lambda x: fun(x / base_f)
+        _rec = _reconstruct_equ(_fun, num_frequency, interface="jax")
+
+        # Convert reconstruction to have original frequencies
+        rec = lambda x: _rec(base_f * x)
+        grad = jax.grad(rec)
+        assert fun_close(fun, rec, zero=jax.numpy.array(0.0))
+        assert fun_close(expected_grad, grad, zero=jax.numpy.array(0.0))
+
+    @pytest.mark.parametrize(
+        "fun, num_frequency, base_f, expected_grad",
+        zip(c_funs, nums_frequency, base_frequencies, expected_grads),
+    )
+    def test_differentiability_tensorflow(self, fun, num_frequency, base_f, expected_grad):
+        """Test that the reconstruction of equidistant-frequency classical
+        functions are differentiable for TensorFlow input variables."""
+        tf = pytest.importorskip("tensorflow")
+        # Convert fun to have integer frequencies
+        base_f = tf.constant(base_f, dtype=tf.float64)
+        _fun = lambda x: fun(x / base_f)
+        _rec = _reconstruct_equ(_fun, num_frequency, interface="tensorflow")
+
+        # Convert reconstruction to have original frequencies
+        rec = lambda x: _rec(base_f * x)
+
+        def grad(arg):
+            arg = tf.Variable(arg)
+            with tf.GradientTape() as tape:
+                out = rec(arg)
+            return tape.gradient(out, arg)
+
+        assert fun_close(fun, rec, zero=tf.Variable(0.0, dtype=tf.float64))
+        assert fun_close(expected_grad, grad, zero=tf.Variable(0.0, dtype=tf.float64))
+
+    @pytest.mark.parametrize(
+        "fun, num_frequency, base_f, expected_grad",
+        zip(c_funs, nums_frequency, base_frequencies, expected_grads),
+    )
+    def test_differentiability_torch(self, fun, num_frequency, base_f, expected_grad):
+        """Test that the reconstruction of equidistant-frequency classical
+        functions are differentiable for Torch input variables."""
+        torch = pytest.importorskip("torch")
+        # Convert fun to have integer frequencies
+        _fun = lambda x: fun(x / base_f)
+        _rec = _reconstruct_equ(_fun, num_frequency, interface="torch")
+
+        # Convert reconstruction to have original frequencies
+        rec = lambda x: _rec(base_f * x)
+        grad = lambda x: torch.autograd.functional.jacobian(rec, x)
+        assert fun_close(fun, rec, zero=torch.tensor(0.0, requires_grad=True))
+        assert fun_close(expected_grad, grad, zero=torch.tensor(0.0, requires_grad=True))
+
+    @pytest.mark.parametrize(
+        "fun, num_frequency, base_f", zip(c_funs, nums_frequency, base_frequencies)
+    )
+    def test_with_classical_fun_num_freq_too_small(self, fun, num_frequency, base_f, mocker):
+        """Test that equidistant-frequency classical functions are
+        reconstructed wrongly if num_frequency is too small."""
+        if num_frequency == 0:
+            pytest.skip("Can't reduce the number of frequencies below 0.")
+        num_frequency -= 1
+        Fun = Lambda(fun)
+        spy = mocker.spy(Fun, "fun")
+        # Convert fun to have integer frequencies
+        _fun = lambda x: Fun(x / base_f)
+        _rec = _reconstruct_equ(_fun, num_frequency)
+
+        # Convert reconstruction to have original frequencies
+        rec = lambda x: _rec(base_f * x)
+        assert spy.call_count == num_frequency * 2 + 1
+        assert not fun_close(fun, rec)
+
+    @pytest.mark.parametrize(
+        "fun, num_frequency, base_f", zip(c_funs, nums_frequency, base_frequencies)
+    )
+    def test_with_classical_fun_num_freq_too_large(self, fun, num_frequency, base_f, mocker):
+        """Test that equidistant-frequency classical functions are
+        reconstructed correctly if num_frequency is too large."""
+        num_frequency += 1
+        Fun = Lambda(fun)
+        spy = mocker.spy(Fun, "fun")
+        # Convert fun to have integer frequencies
+        _fun = lambda x: Fun(x / base_f)
+        _rec = _reconstruct_equ(_fun, num_frequency)
+
+        # Convert reconstruction to have original frequencies
+        rec = lambda x: _rec(base_f * x)
+        assert spy.call_count == num_frequency * 2 + 1
+        assert fun_close(fun, rec)
+
+    all_scales = [
+        [1],
+        [0],
+        [1, 1],
+        [1, 2],
+        [1, 5],
+        [1, 2, 10],
+    ]
+
+    @pytest.mark.parametrize("scales", all_scales)
+    def test_with_qnode(self, scales, mocker):
+        """Test that integer-frequency qnodes are reconstructed correctly."""
+        circuit = get_RX_circuit(scales)
+        Fun = Lambda(circuit)
+        spy = mocker.spy(Fun, "fun")
+        num_frequency = sum(scales)
+        rec = _reconstruct_equ(Fun, num_frequency)
+        assert spy.call_count == num_frequency * 2 + 1
+        assert fun_close(circuit, rec)
+
+        # Repeat, using precomputed f0
+        f0 = circuit(0.0)
+        Fun = Lambda(circuit)
+        spy = mocker.spy(Fun, "fun")
+        rec = _reconstruct_equ(Fun, num_frequency, f0=f0)
+        assert spy.call_count == num_frequency * 2
+        assert fun_close(circuit, rec)
+
+
+class TestReconstructGen:
+    """Tests the one-dimensional reconstruction subroutine based on arbitrary
+    shifted evaluations for arbitrary frequencies."""
+
+    c_funs = [
+        lambda x: -3.27 * np.sin(0.1712 * x) - np.cos(20.812 * x) / 23,
+        lambda x: -0.49 * np.sin(3.2 * x),
+        lambda x: 0.1 * np.cos(-0.1 * x) + 2.9 * np.sin(0.3 * x - 1.2),
+        lambda x: np.sum([np.sin(i * x) for i in range(1, 10)]),
+        lambda x: np.sum(
+            [i ** 0.9 * 0.2 * np.sin(i ** 1.2 * 3.921 * x - 5.1 / i) for i in np.arange(1, 10)]
+        ),
+    ]
+
+    spectra = [
+        [0.1712, 20.812],
+        [3.2],
+        [-0.3, -0.1, 0.0, 0.1, 0.3],
+        list(np.arange(1, 10)),
+        [3.921 * i ** 1.2 for i in np.arange(1, 10)],
+    ]
+
+    expected_grads = [
+        lambda x: -3.27 * np.cos(0.1712 * x) * 0.1712 + np.sin(20.812 * x) / 23 * 20.812,
+        lambda x: -0.49 * np.cos(3.2 * x) * 3.2,
+        lambda x: (-0.1) ** 2 * np.sin(-0.1 * x) + 0.3 * 2.9 * np.cos(0.3 * x - 1.2),
+        lambda x: np.sum([i * np.cos(i * x) for i in range(1, 10)]),
+        lambda x: np.sum(
+            [i ** 2.1 * 3.921 * 0.2 * np.cos(i ** 1.2 * 3.921 * x - 5.1 / i) for i in range(1, 10)]
+        ),
+    ]
+
+    all_shifts = [
+        [-np.pi / 3, -np.pi / 20, 0.0, np.pi / 20, np.pi / 3],
+        [-0.15, -0.05, 0.05],
+        [-2 * np.pi, -np.pi, -0.1, np.pi, 2 * np.pi],
+        np.arange(-9, 10) * np.pi / 19,
+        np.arange(-9, 10) * np.pi / 19,
+    ]
+
+    @pytest.mark.parametrize("fun, spectrum", zip(c_funs, spectra))
+    def test_with_classical_fun(self, fun, spectrum, mocker):
+        """Test that arbitrary-frequency classical functions are
+        reconstructed correctly."""
+        Fun = Lambda(fun)
+        spy = mocker.spy(Fun, "fun")
+        rec = _reconstruct_gen(Fun, spectrum)
+        assert spy.call_count == len([f for f in spectrum if f > 0.0]) * 2 + 1
+        assert fun_close(fun, rec)
+
+        # Repeat, using precomputed f0
+        f0 = fun(0.0)
+        Fun = Lambda(fun)
+        spy = mocker.spy(Fun, "fun")
+        rec = _reconstruct_gen(Fun, spectrum, f0=f0)
+        assert spy.call_count == len([f for f in spectrum if f > 0.0]) * 2
+        assert fun_close(fun, rec)
+
+    @pytest.mark.parametrize("fun, spectrum, shifts", zip(c_funs, spectra, all_shifts))
+    def test_with_classical_fun_with_shifts(self, fun, spectrum, shifts, mocker, recwarn):
+        """Test that arbitrary-frequency classical functions are
+        reconstructed correctly."""
+        Fun = Lambda(fun)
+        spy = mocker.spy(Fun, "fun")
+        rec = _reconstruct_gen(Fun, spectrum, shifts=shifts)
+        assert spy.call_count == len([f for f in spectrum if f > 0.0]) * 2 + 1
+        assert fun_close(fun, rec)
+
+        # Repeat, using precomputed f0
+        f0 = fun(0.0)
+        Fun = Lambda(fun)
+        spy = mocker.spy(Fun, "fun")
+        rec = _reconstruct_gen(Fun, spectrum, shifts=shifts, f0=f0)
+        if 0.0 in shifts:
+            assert spy.call_count == len([f for f in spectrum if f > 0.0]) * 2
+        else:
+            assert len(recwarn) == 1
+            assert recwarn[0].category == UserWarning
+            assert recwarn[0].message.args[0].startswith("The provided value")
+        assert fun_close(fun, rec)
+
+    @pytest.mark.parametrize(
+        "fun, spectrum, expected_grad",
+        zip(c_funs, spectra, expected_grads),
+    )
+    def test_differentiability_autograd(self, fun, spectrum, expected_grad):
+        """Test that the reconstruction of equidistant-frequency classical
+        functions are differentiable for Autograd input variables."""
+        # Convert fun to have integer frequencies
+        rec = _reconstruct_gen(fun, spectrum, interface="autograd")
+        grad = qml.grad(rec)
+        assert fun_close(fun, rec, zero=pnp.array(0.0, requires_grad=True))
+        assert fun_close(expected_grad, grad, zero=pnp.array(0.0, requires_grad=True))
+
+    @pytest.mark.parametrize(
+        "fun, spectrum, expected_grad",
+        zip(c_funs, spectra, expected_grads),
+    )
+    def test_differentiability_jax(self, fun, spectrum, expected_grad):
+        """Test that the reconstruction of equidistant-frequency classical
+        functions are differentiable for JAX input variables."""
+        jax = pytest.importorskip("jax")
+        from jax.config import config
+
+        config.update("jax_enable_x64", True)
+        # Convert fun to have integer frequencies
+        rec = _reconstruct_gen(fun, spectrum, interface="jax")
+        grad = jax.grad(rec)
+        assert fun_close(fun, rec, zero=jax.numpy.array(0.0))
+        assert fun_close(expected_grad, grad, zero=jax.numpy.array(0.0))
+
+    @pytest.mark.parametrize(
+        "fun, spectrum, expected_grad",
+        zip(c_funs, spectra, expected_grads),
+    )
+    def test_differentiability_tensorflow(self, fun, spectrum, expected_grad):
+        """Test that the reconstruction of equidistant-frequency classical
+        functions are differentiable for TensorFlow input variables."""
+        tf = pytest.importorskip("tensorflow")
+        spectrum = tf.constant(spectrum, dtype=tf.float64)
+        # Convert fun to have integer frequencies
+        rec = _reconstruct_gen(fun, spectrum, interface="tensorflow")
+
+        def grad(arg):
+            arg = tf.Variable(arg)
+            with tf.GradientTape() as tape:
+                out = rec(arg)
+            return tape.gradient(out, arg)
+
+        assert fun_close(fun, rec, zero=tf.Variable(0.0))
+        assert fun_close(expected_grad, grad, zero=tf.Variable(0.0))
+
+    @pytest.mark.parametrize(
+        "fun, spectrum, expected_grad",
+        zip(c_funs, spectra, expected_grads),
+    )
+    def test_differentiability_torch(self, fun, spectrum, expected_grad):
+        """Test that the reconstruction of equidistant-frequency classical
+        functions are differentiable for Torch input variables."""
+        torch = pytest.importorskip("torch")
+        spectrum = torch.tensor(spectrum, dtype=torch.float64)
+        # Convert fun to have integer frequencies
+        rec = _reconstruct_gen(fun, spectrum, interface="torch")
+        grad = lambda x: torch.autograd.functional.jacobian(rec, x)
+        assert fun_close(fun, rec, zero=torch.tensor(np.float64(0.0), requires_grad=True))
+        assert fun_close(
+            expected_grad, grad, zero=torch.tensor(np.float64(0.0), requires_grad=True)
+        )
+
+    @pytest.mark.parametrize("fun, spectrum", zip(c_funs, spectra))
+    def test_with_classical_fun_spectrum_incomplete(self, fun, spectrum, mocker):
+        """Test that arbitrary-frequency classical functions are reconstructed wrongly
+        if spectrum does not contain all frequencies."""
+        if len(spectrum) <= 1:
+            pytest.skip("Can't skip a frequency if len(spectrum)<=1.")
+        spectrum = spectrum[:-1]
+        Fun = Lambda(fun)
+        spy = mocker.spy(Fun, "fun")
+        rec = _reconstruct_gen(Fun, spectrum)
+
+        assert spy.call_count == len([f for f in spectrum if f > 0.0]) * 2 + 1
+        assert not fun_close(fun, rec)
+
+    @pytest.mark.parametrize("fun, spectrum", zip(c_funs, spectra))
+    def test_with_classical_fun_spectrum_overcomplete(self, fun, spectrum, mocker):
+        """Test that arbitrary-frequency classical functions are reconstructed correctly
+        if spectrum contains additional frequencies."""
+        spectrum = spectrum + [0.4812759, 1.2281]
+        Fun = Lambda(fun)
+        spy = mocker.spy(Fun, "fun")
+        # Convert fun to have integer frequencies
+        rec = _reconstruct_gen(Fun, spectrum)
+        assert spy.call_count == len([f for f in spectrum if f > 0.0]) * 2 + 1
+        assert fun_close(fun, rec)
+
+    all_scales = [
+        [1.3],
+        [1.02, 1.59],
+        [0.08, 20.2],
+        [1.2, 5.0001],
+        [1.0, 2.0, 3.0],
+    ]
+
+    @pytest.mark.parametrize("scales", all_scales)
+    def test_with_qnode(self, scales, mocker):
+        """Test that arbitrary-frequency qnodes are reconstructed correctly."""
+        circuit = get_RX_circuit(scales)
+        Fun = Lambda(circuit)
+        spy = mocker.spy(Fun, "fun")
+        if len(scales) == 1:
+            spectrum = sorted({0.0, scales[0]})
+        else:
+            _spectra = [{0.0, s} for s in scales]
+            spectrum = sorted(reduce(join_spectra, _spectra, {0.0}))
+
+        rec = _reconstruct_gen(Fun, spectrum)
+        assert spy.call_count == len([f for f in spectrum if f > 0.0]) * 2 + 1
+        assert fun_close(circuit, rec)
+
+        # Repeat, using precomputed f0
+        f0 = circuit(0.0)
+        Fun = Lambda(circuit)
+        spy = mocker.spy(Fun, "fun")
+        rec = _reconstruct_gen(Fun, spectrum, f0=f0)
+        assert spy.call_count == len([f for f in spectrum if f > 0.0]) * 2
+        assert fun_close(circuit, rec)
+
+
+all_ids = [
+    None,
+    {"x": [0, 1], "y": {1, 5}},
+    {"z": (0, 9)},
+    ["z", "y", "x", "x"],
+    {"x", "z", "y"},
+    ("y",),
+    "x",
+]
+
+all_spectra = [
+    {
+        "x": {0: [0.0], 1: [4.2, 0.0, 0.2]},
+        "y": {3: [0.3, 0.0, 0.2], 1: [0.0, 1.1, 5.2], 5: [0.0, 1.2]},
+        "z": {i: [0.0, i * 8.7] for i in range(20)},
+    },
+]
+
+all_shifts = [
+    {
+        "x": {0: [-1.3], 1: [1.0, -0.4, 4.2, 2.3, -1.5]},
+        "y": {
+            3: [0.3 * i + 0.05 for i in range(-2, 3)],
+            1: [-1, -0.5, -0.1, 0.1, 0.9],
+            5: [-1, -0.5, -0.2],
+        },
+        "z": {i: [-np.pi / 2, 0.0, np.pi / 2] for i in range(20)},
+    },
+]
+
+all_nums_frequency = [
+    {
+        "x": {0: 1, 1: 4},
+        "y": {3: 1, 1: 1, 5: 9},
+        "z": {i: 2 * i for i in range(20)},
+    },
+]
+
+
+class TestPrepareJobs:
+    """Tests the subroutine that determines the 1D reconstruction
+    jobs to be carried out for a call to ``reconstruct`` ."""
+
+    def nested_dict_ids_match(self, ndict, ids):
+        if ids.keys() != ndict.keys():
+            return False
+        for id_, ids_ in ids.items():
+            if list(ids_) != list(ndict[id_].keys()):
+                return False
+        return True
+
+    def ids_match(self, ids_in, ids_out):
+        if type(ids_in) == dict:
+            return ids_in == ids_out
+        else:
+            return all(id_ in ids_in for id_ in ids_out)
+
+    @pytest.mark.parametrize("ids", all_ids)
+    @pytest.mark.parametrize("shifts", all_shifts)
+    def test_missing_spectra_and_nums_frequency(self, ids, shifts, tol):
+        """Test that an error is raised if both, spectra
+        and nums_frequency are missing."""
+        with pytest.raises(ValueError, match="Either nums_frequency or spectra"):
+            _prepare_jobs(ids, nums_frequency=None, spectra=None, shifts=shifts, atol=tol)
+
+    @pytest.mark.parametrize("ids", all_ids)
+    @pytest.mark.parametrize("spectra", all_spectra)
+    def test_with_spectra(self, ids, spectra, tol):
+        """Test the prepared jobs when using spectra and shifts."""
+        ids_, recon_fn, jobs, need_f0 = _prepare_jobs(
+            ids,
+            nums_frequency=None,
+            spectra=spectra,
+            shifts=None,
+            atol=tol,
+        )
+        if ids is None:
+            assert self.nested_dict_ids_match(spectra, ids_)
+        else:
+            assert self.ids_match(ids, ids_)
+
+        # Check function to use for 1D reconstructions
+        assert recon_fn == _reconstruct_gen
+
+        # Check reconstruction jobs to be run
+        assert self.nested_dict_ids_match(jobs, ids_)
+
+        # Check all job details
+        for _id, _jobs in jobs.items():
+            for idx, job in _jobs.items():
+                if len(spectra[_id][idx]) == 1:
+                    assert job is None
+                    continue
+                assert list(job.keys()) == ["shifts", "spectrum"]
+                assert job["shifts"] is None
+                assert job["spectrum"] == spectra[_id][idx]
+        assert need_f0
+
+    @pytest.mark.parametrize("ids", all_ids)
+    @pytest.mark.parametrize("spectra", all_spectra)
+    @pytest.mark.parametrize("shifts", all_shifts)
+    def test_with_spectra_and_shifts(self, ids, spectra, shifts, tol):
+        """Test the prepared jobs when using spectra and shifts."""
+        ids_, recon_fn, jobs, need_f0 = _prepare_jobs(
+            ids,
+            nums_frequency=None,
+            spectra=spectra,
+            shifts=shifts,
+            atol=tol,
+        )
+        if ids is None:
+            assert self.nested_dict_ids_match(spectra, ids_)
+        else:
+            assert self.ids_match(ids, ids_)
+
+        # Check function to use for 1D reconstructions
+        assert recon_fn == _reconstruct_gen
+
+        # Check reconstruction jobs to be run
+        assert self.nested_dict_ids_match(jobs, ids_)
+
+        # Check all job details
+        for _id, _jobs in jobs.items():
+            for idx, job in _jobs.items():
+                if len(spectra[_id][idx]) == 1:
+                    assert job is None
+                    continue
+                assert list(job.keys()) == ["shifts", "spectrum"]
+                assert job["shifts"] == shifts[_id][idx]
+                assert job["spectrum"] == spectra[_id][idx]
+        # sometimes need fun at zero if general reconstruction is performed
+        _all_shifts = chain.from_iterable(
+            [
+                sum(
+                    [
+                        __shifts
+                        for par_idx, __shifts in _shifts.items()
+                        if id_ in ids_ and par_idx in ids_[id_]
+                    ],
+                    [],
+                )
+                for id_, _shifts in shifts.items()
+            ],
+        )
+        assert need_f0 == any(np.isclose(_shift, 0.0, atol=tol, rtol=0) for _shift in _all_shifts)
+
+    @pytest.mark.parametrize("ids", all_ids)
+    @pytest.mark.parametrize("spectra", all_spectra)
+    @pytest.mark.parametrize("nums_frequency", all_nums_frequency)
+    def test_with_nums_frequency(self, ids, spectra, nums_frequency, tol):
+        """Test the prepared jobs when using nums_frequency."""
+        """Test ``_prepare_jobs`` with a large variety of test cases (cheap)."""
+
+        ids_, recon_fn, jobs, need_f0 = _prepare_jobs(
+            ids,
+            nums_frequency,
+            None,
+            None,
+            atol=tol,
+        )
+
+        # Check ids
+        if ids is None:
+            assert self.nested_dict_ids_match(nums_frequency, ids_)
+        else:
+            assert self.ids_match(ids, ids_)
+
+        # Check function to use for 1D reconstructions
+        assert recon_fn == _reconstruct_equ
+
+        # Check reconstruction jobs to be run
+        assert self.nested_dict_ids_match(jobs, ids_)
+
+        for _id, _jobs in jobs.items():
+            for idx, job in _jobs.items():
+                if nums_frequency[_id][idx] == 0:
+                    assert job is None
+                    continue
+                assert list(job.keys()) == ["num_frequency"]
+                assert job["num_frequency"] == nums_frequency[_id][idx]
+        # always need fun at zero if equidistant reconstruction is performed
+        assert need_f0
+
+
+dev_1 = qml.device("default.qubit", wires=2)
+
+
+def qnode_0(x):
+    qml.RX(x, wires=0)
+    return qml.expval(qml.PauliZ(0))
+
+
+def qnode_1(X):
+    qml.RX(X[0], wires=0)
+    qml.RX(X[0], wires=1)
+    qml.RX(X[1], wires=0)
+    qml.RX(X[1], wires=1)
+    return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
+
+
+def qnode_2(X, y):
+    qml.RX(X[0], wires=0)
+    qml.RX(X[2], wires=1)
+    qml.RY(y, wires=0)
+    qml.RX(0.5 * X[0], wires=0)
+    return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
+
+
+def qnode_3(X, Y):
+    for i in range(5):
+        qml.RX(X[i], wires=0)
+        qml.RY(Y[i], wires=0)
+        qml.RX(X[i], wires=0)
+    return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
+
+
+def qnode_4(x):
+    return qml.expval(qml.PauliX(0))
+
+
+def qnode_5(Z, y):
+    qml.Hadamard(wires=0)
+    qml.Hadamard(wires=1)
+    qml.RZ(Z[0, 1], wires=0)
+    qml.RZ(Z[2, 4], wires=1)
+    qml.RY(y, wires=0)
+    return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
+
+
+x = 0.1
+y = 2.3
+
+X = pnp.array([i ** 1.2 - 2.0 / i for i in range(1, 6)])
+Y = pnp.array([i ** 0.9 - 1.0 / i for i in range(1, 6)])
+Z = pnp.array(
+    [
+        [0.3, 9.1, -0.2, 0.6, 1.2],
+        [0.9, -0.1, 1.6, 2.3, -1.5],
+        [0.3, 0.1, -0.9, 0.6, 1.8],
+    ]
+)
+
+test_cases_qnodes = [
+    (qnode_0, (x,), "x", None, {"x": {(): [0.0, 1.0]}}, None, 3),
+    (
+        qnode_0,
+        (x,),
+        {"x": [()]},
+        None,
+        {"x": {(): [0.0, 1.0]}},
+        {"x": {(): [-np.pi / 3, 0.0, np.pi / 3]}},
+        3,
+    ),
+    (qnode_0, (x,), "x", {"x": {(): 1}}, None, None, 3),
+    (qnode_1, (X,), {"X"}, None, {"X": {(0,): [0.0, 1.0, 2.0], (1,): [0.0, 2.0]}}, None, 7),
+    (
+        qnode_1,
+        (X,),
+        "X",
+        None,
+        {"X": {(0,): [0.0, 2.0]}},
+        {"X": {(0,): [-np.pi / 2, -0.1, np.pi / 5]}},
+        3,
+    ),
+    (qnode_1, (X,), ["X"], {"X": {(0,): 2, (1,): 2}}, None, None, 9),
+    (
+        qnode_2,
+        (X, y),
+        ["X", "y"],
+        None,
+        {"X": {(0,): [0.0, 0.5, 1.0, 1.5], (2,): [0.0, 1.0]}, "y": {(): [0.0, 1.0]}},
+        None,
+        11,
+    ),
+    (
+        qnode_3,
+        (X, Y),
+        {"X": [(0,), (3,)], "Y": ((4,), (1,))},
+        {"X": {(i,): 2 for i in range(5)}, "Y": {(i,): 1 for i in range(5)}},
+        None,
+        None,
+        13,
+    ),
+    (qnode_4, (x,), ["x"], {"x": {(): 0}}, None, None, 1),
+    (qnode_5, (Z, y), ["Z"], {"Z": {(0, 1): 1, (2, 4): 1, (1, 3): 0}}, None, None, 5),
+]
+
+
+class TestReconstruct:
+    """Tests the integration of ``_reconstruct_equ`` and ``_reconstruct_gen`` via
+    the full ``reconstruct`` function as well as the differentiability of the
+    reconstructed function with respect to their single scalar argument."""
+
+    @pytest.mark.parametrize(
+        "qnode, params, ids, nums_frequency, spectra, shifts, exp_calls",
+        test_cases_qnodes,
+    )
+    def test_with_qnode(
+        self, qnode, params, ids, nums_frequency, spectra, shifts, exp_calls, mocker
+    ):
+        """Run a full reconstruction on a QNode."""
+        qnode = qml.QNode(qnode, dev_1)
+
+        with qml.Tracker(qnode.device) as tracker:
+            recons = reconstruct(qnode, ids, nums_frequency, spectra, shifts)(*params)
+        assert tracker.totals["executions"] == exp_calls
+        arg_names = list(signature(qnode.func).parameters.keys())
+        for outer_key in recons:
+            outer_key_num = arg_names.index(outer_key)
+            for inner_key, rec in recons[outer_key].items():
+                x0 = params[outer_key_num]
+                if not pnp.isscalar(x0):
+                    x0 = x0[inner_key]
+                    shift_vec = qml.math.zeros_like(params[outer_key_num])
+                    shift_vec[inner_key] = 1.0
+                shift_vec = 1.0 if pnp.isscalar(params[outer_key_num]) else shift_vec
+                mask = (
+                    0.0
+                    if pnp.isscalar(params[outer_key_num])
+                    else pnp.ones(qml.math.shape(params[outer_key_num])) - shift_vec
+                )
+                univariate = lambda x: qnode(
+                    *params[:outer_key_num],
+                    params[outer_key_num] * mask + x * shift_vec,
+                    *params[outer_key_num + 1 :],
+                )
+                assert np.isclose(rec(x0), qnode(*params))
+                assert np.isclose(rec(x0 + 0.1), univariate(x0 + 0.1))
+                assert fun_close(rec, univariate, 10)
+
+    @pytest.mark.parametrize(
+        "qnode, params, ids, nums_frequency, spectra, shifts, exp_calls",
+        test_cases_qnodes,
+    )
+    def test_differentiability_autograd(
+        self, qnode, params, ids, nums_frequency, spectra, shifts, exp_calls, mocker
+    ):
+        """Tests the reconstruction and differentiability with autograd."""
+        qnode = qml.QNode(qnode, dev_1, interface="autograd")
+        with qml.Tracker(qnode.device) as tracker:
+            recons = reconstruct(qnode, ids, nums_frequency, spectra, shifts)(*params)
+        assert tracker.totals["executions"] == exp_calls
+        arg_names = list(signature(qnode.func).parameters.keys())
+        for outer_key in recons:
+            outer_key_num = arg_names.index(outer_key)
+            for inner_key, rec in recons[outer_key].items():
+                x0 = params[outer_key_num]
+                if not pnp.isscalar(x0):
+                    x0 = x0[inner_key]
+                    shift_vec = qml.math.zeros_like(params[outer_key_num])
+                    shift_vec[inner_key] = 1.0
+                shift_vec = 1.0 if pnp.isscalar(params[outer_key_num]) else shift_vec
+                mask = (
+                    0.0
+                    if pnp.isscalar(params[outer_key_num])
+                    else pnp.ones(qml.math.shape(params[outer_key_num])) - shift_vec
+                )
+                univariate = lambda x: qnode(
+                    *params[:outer_key_num],
+                    params[outer_key_num] * mask + x * shift_vec,
+                    *params[outer_key_num + 1 :],
+                )
+                exp_qnode_grad = qml.grad(qnode, argnum=outer_key_num)
+                exp_grad = qml.grad(univariate)
+                grad = qml.grad(rec)
+                if nums_frequency is None:
+                    # Gradient evaluation at reconstruction point not supported for
+                    # Dirichlet reconstruction
+                    assert np.isclose(grad(x0), exp_qnode_grad(*params)[inner_key])
+                assert np.isclose(grad(x0 + 0.1), exp_grad(x0 + 0.1))
+                assert fun_close(grad, exp_grad, 10)
+
+    @pytest.mark.parametrize(
+        "qnode, params, ids, nums_frequency, spectra, shifts, exp_calls",
+        test_cases_qnodes,
+    )
+    def test_differentiability_jax(
+        self, qnode, params, ids, nums_frequency, spectra, shifts, exp_calls, mocker
+    ):
+        """Tests the reconstruction and differentiability with JAX."""
+        jax = pytest.importorskip("jax")
+        from jax.config import config
+
+        config.update("jax_enable_x64", True)
+        params = tuple(jax.numpy.array(par) for par in params)
+        qnode = qml.QNode(qnode, dev_1, interface="jax")
+        with qml.Tracker(qnode.device) as tracker:
+            recons = reconstruct(qnode, ids, nums_frequency, spectra, shifts)(*params)
+        assert tracker.totals["executions"] == exp_calls
+        arg_names = list(signature(qnode.func).parameters.keys())
+        for outer_key in recons:
+            outer_key_num = arg_names.index(outer_key)
+            for inner_key, rec in recons[outer_key].items():
+                x0 = params[outer_key_num]
+                if not pnp.isscalar(x0):
+                    x0 = x0[inner_key]
+                    shift_vec = qml.math.zeros_like(params[outer_key_num])
+                    shift_vec = qml.math.scatter_element_add(shift_vec, inner_key, 1.0)
+                shift_vec = 1.0 if pnp.isscalar(params[outer_key_num]) else shift_vec
+                mask = (
+                    0.0
+                    if pnp.isscalar(params[outer_key_num])
+                    else pnp.ones(qml.math.shape(params[outer_key_num])) - shift_vec
+                )
+                univariate = lambda x: qnode(
+                    *params[:outer_key_num],
+                    params[outer_key_num] * mask + x * shift_vec,
+                    *params[outer_key_num + 1 :],
+                )
+                exp_qnode_grad = jax.grad(qnode, argnums=outer_key_num)
+                exp_grad = jax.grad(univariate)
+                grad = jax.grad(rec)
+                assert np.isclose(grad(x0), exp_qnode_grad(*params)[inner_key])
+                assert np.isclose(grad(x0 + 0.1), exp_grad(x0 + 0.1))
+                assert fun_close(grad, exp_grad, 10)
+
+    @pytest.mark.parametrize(
+        "qnode, params, ids, nums_frequency, spectra, shifts, exp_calls",
+        test_cases_qnodes,
+    )
+    def test_differentiability_tensorflow(
+        self, qnode, params, ids, nums_frequency, spectra, shifts, exp_calls, mocker
+    ):
+        """Tests the reconstruction and differentiability with TensorFlow."""
+        if qnode == qnode_4:
+            pytest.skip("Gradients are empty in TensorFlow for independent functions.")
+        tf = pytest.importorskip("tensorflow")
+        qnode = qml.QNode(qnode, dev_1, interface="tf")
+        params = tuple(tf.Variable(par, dtype=tf.float64) for par in params)
+        if spectra is not None:
+            spectra = {
+                outer_key: {
+                    inner_key: tf.constant(val, dtype=tf.float64)
+                    for inner_key, val in outer_val.items()
+                }
+                for outer_key, outer_val in spectra.items()
+            }
+        if shifts is not None:
+            shifts = {
+                outer_key: {
+                    inner_key: tf.constant(val, dtype=tf.float64)
+                    for inner_key, val in outer_val.items()
+                }
+                for outer_key, outer_val in shifts.items()
+            }
+        with qml.Tracker(qnode.device) as tracker:
+            recons = reconstruct(qnode, ids, nums_frequency, spectra, shifts)(*params)
+        assert tracker.totals["executions"] == exp_calls
+        arg_names = list(signature(qnode.func).parameters.keys())
+        for outer_key in recons:
+            outer_key_num = arg_names.index(outer_key)
+            for inner_key, rec in recons[outer_key].items():
+                if outer_key == "Z" and inner_key == (1, 3):
+                    # This is a constant function dependence, which can
+                    # not be properly resolved by this test.
+                    continue
+                x0 = params[outer_key_num]
+                if not len(qml.math.shape(x0)) == 0:
+                    x0 = x0[inner_key]
+                    shift_vec = qml.math.zeros_like(params[outer_key_num])
+                    shift_vec = qml.math.scatter_element_add(shift_vec, inner_key, 1.0)
+                    mask = pnp.ones(qml.math.shape(params[outer_key_num])) - shift_vec
+                else:
+                    shift_vec = 1.0
+                    mask = 0.0
+                univariate = lambda x: qnode(
+                    *params[:outer_key_num],
+                    params[outer_key_num] * mask + x * shift_vec,
+                    *params[outer_key_num + 1 :],
+                )
+                with tf.GradientTape() as tape:
+                    out = qnode(*params)
+                exp_qnode_grad = tape.gradient(out, params[outer_key_num])
+
+                def exp_grad(x):
+                    x = tf.Variable(x, dtype=tf.float64)
+                    with tf.GradientTape() as tape:
+                        out = univariate(x)
+                    return tape.gradient(out, x)
+
+                def grad(x):
+                    x = tf.Variable(x, dtype=tf.float64)
+                    with tf.GradientTape() as tape:
+                        out = rec(x)
+                    return tape.gradient(out, x)
+
+                if nums_frequency is None:
+                    # Gradient evaluation at reconstruction point not supported for
+                    # Dirichlet reconstruction
+                    assert np.isclose(grad(x0), exp_qnode_grad[inner_key])
+                assert np.isclose(grad(x0 + 0.1), exp_grad(x0 + 0.1))
+                assert fun_close(grad, exp_grad, 10)
+
+    @pytest.mark.parametrize(
+        "qnode, params, ids, nums_frequency, spectra, shifts, exp_calls",
+        test_cases_qnodes,
+    )
+    def test_differentiability_torch(
+        self, qnode, params, ids, nums_frequency, spectra, shifts, exp_calls, mocker
+    ):
+        """Tests the reconstruction and differentiability with Torch."""
+        torch = pytest.importorskip("torch")
+        qnode = qml.QNode(qnode, dev_1, interface="torch")
+        params = tuple(torch.tensor(par, requires_grad=True, dtype=torch.float64) for par in params)
+        if spectra is not None:
+            spectra = {
+                outer_key: {
+                    inner_key: torch.tensor(val, dtype=torch.float64)
+                    for inner_key, val in outer_val.items()
+                }
+                for outer_key, outer_val in spectra.items()
+            }
+        if shifts is not None:
+            shifts = {
+                outer_key: {
+                    inner_key: torch.tensor(val, dtype=torch.float64)
+                    for inner_key, val in outer_val.items()
+                }
+                for outer_key, outer_val in shifts.items()
+            }
+        with qml.Tracker(qnode.device) as tracker:
+            recons = reconstruct(qnode, ids, nums_frequency, spectra, shifts)(*params)
+        assert tracker.totals["executions"] == exp_calls
+        arg_names = list(signature(qnode.func).parameters.keys())
+        for outer_key in recons:
+            outer_key_num = arg_names.index(outer_key)
+            for inner_key, rec in recons[outer_key].items():
+                x0 = params[outer_key_num]
+                if not len(qml.math.shape(x0)) == 0:
+                    x0 = x0[inner_key]
+                    shift_vec = qml.math.zeros_like(params[outer_key_num])
+                    shift_vec = qml.math.scatter_element_add(shift_vec, inner_key, 1.0)
+                    mask = torch.ones(qml.math.shape(params[outer_key_num])) - shift_vec
+                else:
+                    shift_vec = 1.0
+                    mask = 0.0
+                univariate = lambda x: qnode(
+                    *params[:outer_key_num],
+                    params[outer_key_num] * mask + x * shift_vec,
+                    *params[outer_key_num + 1 :],
+                )
+                exp_qnode_grad = torch.autograd.functional.jacobian(qnode, params)[outer_key_num]
+
+                exp_grad = lambda x: torch.autograd.functional.jacobian(univariate, x)
+                grad = lambda x: torch.autograd.functional.jacobian(rec, x)
+
+                assert np.isclose(grad(x0), exp_qnode_grad[inner_key])
+                assert np.isclose(grad(x0 + 0.1), exp_grad(x0 + 0.1))
+                assert fun_close(
+                    grad, exp_grad, zero=torch.tensor(0.0, requires_grad=True), samples=10
+                )
diff --git a/tests/math/test_functions.py b/tests/math/test_functions.py
index be69f0b7c..da00166aa 100644
--- a/tests/math/test_functions.py
+++ b/tests/math/test_functions.py
@@ -22,6 +22,7 @@ from pennylane import numpy as np
 from pennylane import math as fn
 from autograd.numpy.numpy_boxes import ArrayBox
 
+import semantic_version
 
 tf = pytest.importorskip("tensorflow", minversion="2.1")
 torch = pytest.importorskip("torch")
@@ -534,6 +535,223 @@ class TestDot:
         assert fn.allequal(res, expected)
 
 
+class TestTensordot:
+    """Tests for the dot product function."""
+
+    v1 = torch.tensor([0.1, 0.5, -0.9, 1.0, -4.2, 0.1], dtype=torch.float64)
+    v2 = torch.tensor([4.3, -1.2, 8.2, 0.6, -4.2, -11.0], dtype=torch.float64)
+    _arange = np.arange(0, 54).reshape((9, 6)).astype(np.float64)
+    _shuffled_arange = np.array(
+        [
+            [42.0, 43.0, 44.0, 45.0, 46.0, 47.0],
+            [6.0, 7.0, 8.0, 9.0, 10.0, 11.0],
+            [30.0, 31.0, 32.0, 33.0, 34.0, 35.0],
+            [0.0, 1.0, 2.0, 3.0, 4.0, 5.0],
+            [48.0, 49.0, 50.0, 51.0, 52.0, 53.0],
+            [12.0, 13.0, 14.0, 15.0, 16.0, 17.0],
+            [24.0, 25.0, 26.0, 27.0, 28.0, 29.0],
+            [18.0, 19.0, 20.0, 21.0, 22.0, 23.0],
+            [36.0, 37.0, 38.0, 39.0, 40.0, 41.0],
+        ],
+        dtype=np.float64,
+    )
+    M1 = torch.tensor(_arange)
+    M2 = torch.tensor(_shuffled_arange)
+    T1 = np.arange(0, 3 * 6 * 9 * 2).reshape((3, 6, 9, 2)).astype(np.float64)
+    T1 = np.array([T1[1], T1[0], T1[2]])
+
+    v1_dot_v2 = 9.59
+    v1_outer_v2 = np.array(
+        [
+            [0.43, -0.12, 0.82, 0.06, -0.42, -1.1],
+            [2.15, -0.6, 4.1, 0.3, -2.1, -5.5],
+            [-3.87, 1.08, -7.38, -0.54, 3.78, 9.9],
+            [4.3, -1.2, 8.2, 0.6, -4.2, -11.0],
+            [-18.06, 5.04, -34.44, -2.52, 17.64, 46.2],
+            [0.43, -0.12, 0.82, 0.06, -0.42, -1.1],
+        ],
+        dtype=np.float64,
+    )
+
+    M1_dot_v1 = torch.tensor(
+        [-14.6, -35.0, -55.4, -75.8, -96.2, -116.6, -137.0, -157.4, -177.8], dtype=torch.float64
+    )
+    M1_dot_v2 = torch.tensor(
+        [-54.8, -74.6, -94.4, -114.2, -134.0, -153.8, -173.6, -193.4, -213.2], dtype=torch.float64
+    )
+    M2_dot_v1 = torch.tensor(
+        [-157.4, -35.0, -116.6, -14.6, -177.8, -55.4, -96.2, -75.8, -137.0], dtype=torch.float64
+    )
+    M2_dot_v2 = torch.tensor(
+        [-193.4, -74.6, -153.8, -54.8, -213.2, -94.4, -134.0, -114.2, -173.6], dtype=torch.float64
+    )
+    M1_dot_M2T = torch.tensor(
+        [
+            [685, 145, 505, 55, 775, 235, 415, 325, 595],
+            [2287, 451, 1675, 145, 2593, 757, 1369, 1063, 1981],
+            [3889, 757, 2845, 235, 4411, 1279, 2323, 1801, 3367],
+            [5491, 1063, 4015, 325, 6229, 1801, 3277, 2539, 4753],
+            [7093, 1369, 5185, 415, 8047, 2323, 4231, 3277, 6139],
+            [8695, 1675, 6355, 505, 9865, 2845, 5185, 4015, 7525],
+            [10297, 1981, 7525, 595, 11683, 3367, 6139, 4753, 8911],
+            [11899, 2287, 8695, 685, 13501, 3889, 7093, 5491, 10297],
+            [13501, 2593, 9865, 775, 15319, 4411, 8047, 6229, 11683],
+        ],
+        dtype=torch.float64,
+    )
+    M1T_dot_M2 = torch.tensor(
+        [
+            [5256, 5472, 5688, 5904, 6120, 6336],
+            [5472, 5697, 5922, 6147, 6372, 6597],
+            [5688, 5922, 6156, 6390, 6624, 6858],
+            [5904, 6147, 6390, 6633, 6876, 7119],
+            [6120, 6372, 6624, 6876, 7128, 7380],
+            [6336, 6597, 6858, 7119, 7380, 7641],
+        ],
+        dtype=torch.float64,
+    )
+
+    T1_dot_v1 = torch.tensor(
+        [
+            [
+                [-630.0, -633.4],
+                [-636.8, -640.2],
+                [-643.6, -647.0],
+                [-650.4, -653.8],
+                [-657.2, -660.6],
+                [-664.0, -667.4],
+                [-670.8, -674.2],
+                [-677.6, -681.0],
+                [-684.4, -687.8],
+            ],
+            [
+                [-262.8, -266.2],
+                [-269.6, -273.0],
+                [-276.4, -279.8],
+                [-283.2, -286.6],
+                [-290.0, -293.4],
+                [-296.8, -300.2],
+                [-303.6, -307.0],
+                [-310.4, -313.8],
+                [-317.2, -320.6],
+            ],
+            [
+                [-997.2, -1000.6],
+                [-1004.0, -1007.4],
+                [-1010.8, -1014.2],
+                [-1017.6, -1021.0],
+                [-1024.4, -1027.8],
+                [-1031.2, -1034.6],
+                [-1038.0, -1041.4],
+                [-1044.8, -1048.2],
+                [-1051.6, -1055.0],
+            ],
+        ],
+        dtype=torch.float64,
+    )
+
+    T1_dot_v2 = torch.tensor(
+        [
+            [
+                [-1342.8, -1346.1],
+                [-1349.4, -1352.7],
+                [-1356.0, -1359.3],
+                [-1362.6, -1365.9],
+                [-1369.2, -1372.5],
+                [-1375.8, -1379.1],
+                [-1382.4, -1385.7],
+                [-1389.0, -1392.3],
+                [-1395.6, -1398.9],
+            ],
+            [
+                [-986.4, -989.7],
+                [-993.0, -996.3],
+                [-999.6, -1002.9],
+                [-1006.2, -1009.5],
+                [-1012.8, -1016.1],
+                [-1019.4, -1022.7],
+                [-1026.0, -1029.3],
+                [-1032.6, -1035.9],
+                [-1039.2, -1042.5],
+            ],
+            [
+                [-1699.2, -1702.5],
+                [-1705.8, -1709.1],
+                [-1712.4, -1715.7],
+                [-1719.0, -1722.3],
+                [-1725.6, -1728.9],
+                [-1732.2, -1735.5],
+                [-1738.8, -1742.1],
+                [-1745.4, -1748.7],
+                [-1752.0, -1755.3],
+            ],
+        ],
+        dtype=torch.float64,
+    )
+
+    T1_dot_M1 = torch.tensor(
+        [
+            [237546.0, 238977.0],
+            [82998.0, 84429.0],
+            [392094.0, 393525.0],
+        ],
+        dtype=torch.float64,
+    )
+
+    T1_dot_M2 = torch.tensor(
+        [
+            [233370.0, 234801.0],
+            [78822.0, 80253.0],
+            [387918.0, 389349.0],
+        ],
+        dtype=torch.float64,
+    )
+
+    @pytest.mark.parametrize("axes", [[[0], [0]], [[-1], [0]], [[0], [-1]], [[-1], [-1]]])
+    def test_tensordot_torch_vector_vector(self, axes):
+        assert fn.allclose(fn.tensordot(self.v1, self.v2, axes=axes), self.v1_dot_v2)
+
+    def test_tensordot_torch_outer(self):
+        assert fn.allclose(fn.tensordot(self.v1, self.v2, axes=0), self.v1_outer_v2)
+        assert fn.allclose(fn.tensordot(self.v2, self.v1, axes=0), qml.math.T(self.v1_outer_v2))
+
+    def test_tensordot_torch_outer_with_old_version(self, monkeypatch):
+        with monkeypatch.context() as m:
+            m.setattr("torch.__version__", "1.9.0")
+            assert fn.allclose(fn.tensordot(self.v1, self.v2, axes=0), self.v1_outer_v2)
+            assert fn.allclose(fn.tensordot(self.v2, self.v1, axes=0), qml.math.T(self.v1_outer_v2))
+
+    @pytest.mark.parametrize(
+        "M, v, expected",
+        [(M1, v1, M1_dot_v1), (M1, v2, M1_dot_v2), (M2, v1, M2_dot_v1), (M2, v2, M2_dot_v2)],
+    )
+    @pytest.mark.parametrize("axes", [[[1], [0]], [[-1], [0]], [[1], [-1]], [[-1], [-1]]])
+    def test_tensordot_torch_matrix_vector(self, M, v, expected, axes):
+        assert fn.allclose(fn.tensordot(M, v, axes=axes), expected)
+
+    @pytest.mark.parametrize("axes", [[[1], [0]], [[-1], [0]], [[1], [-2]], [[-1], [-2]]])
+    def test_tensordot_torch_matrix_matrix(self, axes):
+        assert fn.allclose(fn.tensordot(self.M1, qml.math.T(self.M2), axes=axes), self.M1_dot_M2T)
+        assert fn.allclose(
+            fn.tensordot(self.M2, qml.math.T(self.M1), axes=axes), qml.math.T(self.M1_dot_M2T)
+        )
+        assert fn.allclose(fn.tensordot(qml.math.T(self.M1), self.M2, axes=axes), self.M1T_dot_M2)
+        assert fn.allclose(
+            fn.tensordot(qml.math.T(self.M2), self.M1, axes=axes), qml.math.T(self.M1T_dot_M2)
+        )
+
+    @pytest.mark.parametrize("axes", [[[1], [0]], [[-3], [0]], [[1], [-1]], [[-3], [-1]]])
+    @pytest.mark.parametrize("v, expected", [(v1, T1_dot_v1), (v2, T1_dot_v2)])
+    def test_tensordot_torch_tensor_vector(self, v, expected, axes):
+        assert fn.allclose(fn.tensordot(self.T1, v, axes=axes), expected)
+
+    @pytest.mark.parametrize("axes1", [[1, 2], [-3, -2], [1, -2], [-3, 2]])
+    @pytest.mark.parametrize("axes2", [[1, 0], [-1, -2], [1, -2]])
+    @pytest.mark.parametrize("M, expected", [(M1, T1_dot_M1), (M2, T1_dot_M2)])
+    def test_tensordot_torch_tensor_matrix(self, M, expected, axes1, axes2):
+        assert fn.allclose(fn.tensordot(self.T1, M, axes=[axes1, axes2]), expected)
+
+
 # the following test data is of the form
 # [original shape, axis to expand, new shape]
 expand_dims_test_data = [
@@ -1059,22 +1277,34 @@ class TestTake:
 
 
 where_data = [
-    np.array([[[1, 2], [3, 4], [-1, 1]], [[5, 6], [0, -1], [2, 1]]]),
-    torch.tensor([[[1, 2], [3, 4], [-1, 1]], [[5, 6], [0, -1], [2, 1]]]),
-    onp.array([[[1, 2], [3, 4], [-1, 1]], [[5, 6], [0, -1], [2, 1]]]),
-    tf.constant([[[1, 2], [3, 4], [-1, 1]], [[5, 6], [0, -1], [2, 1]]]),
-    tf.Variable([[[1, 2], [3, 4], [-1, 1]], [[5, 6], [0, -1], [2, 1]]]),
-    jnp.array([[[1, 2], [3, 4], [-1, 1]], [[5, 6], [0, -1], [2, 1]]]),
+    ("autograd", np.array([[[1, 2], [3, 4], [-1, 1]], [[5, 6], [0, -1], [2, 1]]])),
+    ("torch", torch.tensor([[[1, 2], [3, 4], [-1, 1]], [[5, 6], [0, -1], [2, 1]]])),
+    ("numpy", onp.array([[[1, 2], [3, 4], [-1, 1]], [[5, 6], [0, -1], [2, 1]]])),
+    ("tf", tf.constant([[[1, 2], [3, 4], [-1, 1]], [[5, 6], [0, -1], [2, 1]]])),
+    ("tf", tf.Variable([[[1, 2], [3, 4], [-1, 1]], [[5, 6], [0, -1], [2, 1]]])),
+    ("jax", jnp.array([[[1, 2], [3, 4], [-1, 1]], [[5, 6], [0, -1], [2, 1]]])),
 ]
 
 
-@pytest.mark.parametrize("t", where_data)
-def test_where(t):
+@pytest.mark.parametrize("interface, t", where_data)
+def test_where(interface, t):
     """Test that the where function works as expected"""
+    # With output values
     res = fn.where(t < 0, 100 * fn.ones_like(t), t)
     expected = np.array([[[1, 2], [3, 4], [100, 1]], [[5, 6], [0, 100], [2, 1]]])
     assert fn.allclose(res, expected)
 
+    # Without output values
+    res = fn.where(t > 0)
+    expected = (
+        [0, 0, 0, 0, 0, 1, 1, 1, 1],
+        [0, 0, 1, 1, 2, 0, 0, 2, 2],
+        [0, 1, 0, 1, 1, 0, 1, 0, 1],
+    )
+    if interface == "tf":
+        expected = qml.math.T(expected)
+    assert all(fn.allclose(_res, _exp) for _res, _exp in zip(res, expected))
+
 
 squeeze_data = [
     np.ones((1, 2, 3, 1, 5, 1)),
@@ -1095,102 +1325,213 @@ def test_squeeze(t):
 class TestScatterElementAdd:
     """Tests for the scatter_element_add function"""
 
+    x = onp.ones((2, 3), dtype=np.float64)
+    y = onp.array(0.56)
+    index = [1, 2]
+    expected_val = onp.array([[1.0, 1.0, 1.0], [1.0, 1.0, 1.3136]])
+    expected_grad_x = onp.array([[0, 0, 0], [0, 0, 1.0]])
+    expected_grad_y = 2 * y
+    expected_jac_x = onp.eye(6).reshape((2, 3, 2, 3))
+    expected_jac_y = onp.array([[0, 0, 0], [0, 0, 2 * y]])
+
     def test_array(self):
         """Test that a NumPy array is differentiable when using scatter addition"""
-        x = np.array([[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]], requires_grad=True)
-        y = np.array(0.56, requires_grad=True)
+        x = np.array(self.x, requires_grad=True)
+        y = np.array(self.y, requires_grad=True)
 
         def cost(weights):
-            return fn.scatter_element_add(weights[0], [1, 2], weights[1] ** 2)
+            return fn.scatter_element_add(weights[0], self.index, weights[1] ** 2)
 
         res = cost([x, y])
         assert isinstance(res, np.ndarray)
-        assert fn.allclose(res, onp.array([[1.0, 1.0, 1.0], [1.0, 1.0, 1.3136]]))
+        assert fn.allclose(res, self.expected_val)
 
-        grad = qml.grad(lambda weights: cost(weights)[1, 2])([x, y])
-        assert fn.allclose(grad[0], onp.array([[0, 0, 0], [0, 0, 1.0]]))
-        assert fn.allclose(grad[1], 2 * y)
+        grad = qml.grad(lambda weights: cost(weights)[self.index[0], self.index[1]])([x, y])
+        assert fn.allclose(grad[0], self.expected_grad_x)
+        assert fn.allclose(grad[1], self.expected_grad_y)
 
     def test_array_multi(self):
         """Test that a NumPy array and the addend are differentiable when using
         scatter addition (multi dispatch)."""
-        x = np.array([[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]], requires_grad=True)
-        y = np.array(0.56, requires_grad=True)
+        x = np.array(self.x, requires_grad=True)
+        y = np.array(self.y, requires_grad=True)
 
         def cost_multi(weight_0, weight_1):
-            return fn.scatter_element_add(weight_0, [1, 2], weight_1 ** 2)
+            return fn.scatter_element_add(weight_0, self.index, weight_1 ** 2)
 
         res = cost_multi(x, y)
         assert isinstance(res, np.ndarray)
-        assert fn.allclose(res, onp.array([[1.0, 1.0, 1.0], [1.0, 1.0, 1.3136]]))
+        assert fn.allclose(res, self.expected_val)
 
         jac = qml.jacobian(lambda *weights: cost_multi(*weights))(x, y)
-        assert fn.allclose(jac[0], onp.eye(6).reshape((2, 3, 2, 3)))
-        assert fn.allclose(jac[1], onp.array([[0, 0, 0], [0, 0, 2 * y]]))
+        assert fn.allclose(jac[0], self.expected_jac_x)
+        assert fn.allclose(jac[1], self.expected_jac_y)
 
     def test_tensorflow(self):
         """Test that a TF tensor is differentiable when using scatter addition"""
-        x = tf.Variable([[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]])
-        y = tf.Variable(0.56)
+        x = tf.Variable(self.x)
+        y = tf.Variable(self.y)
 
         with tf.GradientTape() as tape:
-            res = fn.scatter_element_add(x, [1, 2], y ** 2)
-            loss = res[1, 2]
+            res = fn.scatter_element_add(x, self.index, y ** 2)
+            loss = res[self.index[0], self.index[1]]
 
         assert isinstance(res, tf.Tensor)
-        assert fn.allclose(res, onp.array([[1.0, 1.0, 1.0], [1.0, 1.0, 1.3136]]))
+        assert fn.allclose(res, self.expected_val)
 
         grad = tape.gradient(loss, [x, y])
-        assert fn.allclose(grad[0], onp.array([[0, 0, 0], [0, 0, 1.0]]))
-        assert fn.allclose(grad[1], 2 * y)
+        assert fn.allclose(grad[0], self.expected_grad_x)
+        assert fn.allclose(grad[1], self.expected_grad_y)
 
     def test_torch(self):
         """Test that a torch tensor is differentiable when using scatter addition"""
-        x = torch.tensor([[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]], requires_grad=True)
-        y = torch.tensor(0.56, requires_grad=True)
+        x = torch.tensor(self.x, requires_grad=True)
+        y = torch.tensor(self.y, requires_grad=True)
 
-        res = fn.scatter_element_add(x, [1, 2], y ** 2)
-        loss = res[1, 2]
+        res = fn.scatter_element_add(x, self.index, y ** 2)
+        loss = res[self.index[0], self.index[1]]
 
         assert isinstance(res, torch.Tensor)
-        assert fn.allclose(res.detach(), onp.array([[1.0, 1.0, 1.0], [1.0, 1.0, 1.3136]]))
+        assert fn.allclose(res.detach(), self.expected_val)
 
         loss.backward()
-        assert fn.allclose(x.grad, onp.array([[0, 0, 0], [0, 0, 1.0]]))
-        assert fn.allclose(y.grad, 2 * y)
+        assert fn.allclose(x.grad, self.expected_grad_x)
+        assert fn.allclose(y.grad, self.expected_grad_y)
 
     def test_jax(self):
         """Test that a JAX array is differentiable when using scatter addition"""
-        x = jnp.array([[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]])
-        y = jnp.array(0.56)
+        x = jnp.array(self.x)
+        y = jnp.array(self.y)
 
         def cost(weights):
-            return fn.scatter_element_add(weights[0], [1, 2], weights[1] ** 2)
+            return fn.scatter_element_add(weights[0], self.index, weights[1] ** 2)
 
         res = cost([x, y])
         assert isinstance(res, jax.interpreters.xla.DeviceArray)
-        assert fn.allclose(res, onp.array([[1.0, 1.0, 1.0], [1.0, 1.0, 1.3136]]))
+        assert fn.allclose(res, self.expected_val)
 
-        grad = jax.grad(lambda weights: cost(weights)[1, 2])([x, y])
-        assert fn.allclose(grad[0], onp.array([[0, 0, 0], [0, 0, 1.0]]))
-        assert fn.allclose(grad[1], 2 * y)
+        grad = jax.grad(lambda weights: cost(weights)[self.index[0], self.index[1]])([x, y])
+        assert fn.allclose(grad[0], self.expected_grad_x)
+        assert fn.allclose(grad[1], self.expected_grad_y)
 
     def test_jax_multi(self):
         """Test that a NumPy array and the addend are differentiable when using
         scatter addition (multi dispatch)."""
-        x = jnp.array([[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]])
-        y = jnp.array(0.56)
+        x = jnp.array(self.x)
+        y = jnp.array(self.y)
 
         def cost_multi(weight_0, weight_1):
-            return fn.scatter_element_add(weight_0, [1, 2], weight_1 ** 2)
+            return fn.scatter_element_add(weight_0, self.index, weight_1 ** 2)
 
         res = cost_multi(x, y)
         assert isinstance(res, jax.interpreters.xla.DeviceArray)
-        assert fn.allclose(res, onp.array([[1.0, 1.0, 1.0], [1.0, 1.0, 1.3136]]))
+        assert fn.allclose(res, self.expected_val)
 
         jac = jax.jacobian(lambda *weights: cost_multi(*weights), argnums=[0, 1])(x, y)
-        assert fn.allclose(jac[0], onp.eye(6).reshape((2, 3, 2, 3)))
-        assert fn.allclose(jac[1], onp.array([[0, 0, 0], [0, 0, 2 * y]]))
+        assert fn.allclose(jac[0], self.expected_jac_x)
+        assert fn.allclose(jac[1], self.expected_jac_y)
+
+
+class TestScatterElementAddMultiValue:
+    """Tests for the scatter_element_add function when adding
+    multiple values at multiple positions."""
+
+    x = onp.ones((2, 3), dtype=np.float64)
+    y = onp.array(0.56)
+    indices = [[1, 0], [2, 1]]
+    expected_val = onp.array([[1.0, 1.3136, 1.0], [1.0, 1.0, 1.27636]])
+    expected_grad_x = onp.array([[0, 1.0, 0], [0, 0, 1.0]])
+    expected_grad_y = 2 * y + onp.cos(y / 2) / 2
+
+    def test_array(self):
+        """Test that a NumPy array is differentiable when using scatter addition
+        with multiple values."""
+        x = np.array(self.x, requires_grad=True)
+        y = np.array(self.y, requires_grad=True)
+
+        def cost(weights):
+            return fn.scatter_element_add(
+                weights[0], self.indices, [fn.sin(weights[1] / 2), weights[1] ** 2]
+            )
+
+        res = cost([x, y])
+        assert isinstance(res, np.ndarray)
+        assert fn.allclose(res, self.expected_val)
+
+        scalar_cost = (
+            lambda weights: cost(weights)[self.indices[0][0], self.indices[1][0]]
+            + cost(weights)[self.indices[0][1], self.indices[1][1]]
+        )
+        grad = qml.grad(scalar_cost)([x, y])
+        assert fn.allclose(grad[0], self.expected_grad_x)
+        assert fn.allclose(grad[1], self.expected_grad_y)
+
+    def test_tensorflow(self):
+        """Test that a TF tensor is differentiable when using scatter addition
+        with multiple values."""
+        x = tf.Variable(self.x)
+        y = tf.Variable(self.y)
+
+        with tf.GradientTape() as tape:
+            res = fn.scatter_element_add(x, self.indices, [tf.sin(y / 2), y ** 2])
+            loss = (
+                res[self.indices[0][0], self.indices[1][0]]
+                + res[self.indices[0][1], self.indices[1][1]]
+            )
+
+        assert isinstance(res, tf.Tensor)
+        assert fn.allclose(res, self.expected_val)
+
+        grad = tape.gradient(loss, [x, y])
+        assert fn.allclose(grad[0], self.expected_grad_x)
+        assert fn.allclose(grad[1], self.expected_grad_y)
+
+    def test_torch(self):
+        """Test that a torch tensor is differentiable when using scatter addition
+        with multiple values."""
+        x = torch.tensor(self.x, requires_grad=True)
+        y = torch.tensor(self.y, requires_grad=True)
+
+        values = torch.zeros(2)
+        values[0] += torch.sin(y / 2)
+        values[1] += y ** 2
+        res = fn.scatter_element_add(x, self.indices, values)
+        loss = (
+            res[self.indices[0][0], self.indices[1][0]]
+            + res[self.indices[0][1], self.indices[1][1]]
+        )
+
+        assert isinstance(res, torch.Tensor)
+        assert fn.allclose(res.detach(), self.expected_val)
+
+        loss.backward()
+        assert fn.allclose(x.grad, self.expected_grad_x)
+        assert fn.allclose(y.grad, self.expected_grad_y)
+
+    def test_jax(self):
+        """Test that a JAX array is differentiable when using scatter addition
+        with multiple values."""
+        x = jnp.array(self.x)
+        y = jnp.array(self.y)
+
+        def cost(weights):
+            return fn.scatter_element_add(
+                weights[0], self.indices, [fn.sin(weights[1] / 2), weights[1] ** 2]
+            )
+
+        res = cost([x, y])
+        assert isinstance(res, jax.interpreters.xla.DeviceArray)
+        print(res)
+        print(self.expected_val)
+        assert fn.allclose(res, self.expected_val)
+
+        scalar_cost = (
+            lambda weights: cost(weights)[self.indices[0][0], self.indices[1][0]]
+            + cost(weights)[self.indices[0][1], self.indices[1][1]]
+        )
+        grad = jax.grad(scalar_cost)([x, y])
+        assert fn.allclose(grad[0], self.expected_grad_x)
+        assert fn.allclose(grad[1], self.expected_grad_y)
 
 
 class TestDiag: