Repository: PennyLaneAI/pennylane
Pull request #1378 (2021-05-31 18:09:12)

Add spectrum function to Fourier module
**Context:**

Recently, PL got a new Fourier module that can be used to investigate the Fourier representation of quantum circuits. Due to unforeseen difficulties we did not merge the function that would analytically derive the Fourier spectrum from a circuit. This is a new attempt with a much simplified scope.

**Description of the Change:**

Adds `spectrum` qnode transform.

**Benefits:**

The new design lets users mark the gates that are considered as input-encoding gates via their `id` attribute. It will not try to track parameters, but simply see if the marked gates define generators, and thus support the assumptions of how to compute the Fourier spectrum directly.


github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [.github/CHANGELOG.md](/PennyLaneAI/pennylane/blob/master/.github/CHANGELOG.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.
mariaschuld: @glassnotes this is still WIP, but maybe you want to check if the idea is sound?
codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1378?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#1378](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1378?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (5c246c8) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/a4d7abdc047869a2d378e54c7b1cc0b036ae7625?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (a4d7abd) will **decrease** coverage by `0.01%`.
> The diff coverage is `96.22%`.

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

```diff
@@            Coverage Diff             @@
##           master    #1378      +/-   ##
==========================================
- Coverage   98.16%   98.14%   -0.02%     
==========================================
  Files         154      155       +1     
  Lines       11548    11609      +61     
==========================================
+ Hits        11336    11394      +58     
- Misses        212      215       +3     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1378?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/fourier/spectrum.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1378/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ZvdXJpZXIvc3BlY3RydW0ucHk=) | `95.74% <95.74%> (ø)` | |
| [pennylane/fourier/\_\_init\_\_.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1378/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%> (ø)` | |
| [pennylane/operation.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1378/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wZXJhdGlvbi5weQ==) | `95.67% <100.00%> (-0.16%)` | :arrow_down: |
| [pennylane/transforms/control.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1378/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvY29udHJvbC5weQ==) | `100.00% <0.00%> (ø)` | |

------

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

mariaschuld: @glassnotes I hope I addressed all comments!
mariaschuld: @co9olguy I wonder if we can get this in before feature freeze - let us know if we must pick another reviewer!
co9olguy: Sorry, this didn't show up in my `review-requested` filter for some reason! :slightly_frowning_face: 
I'll try to have a look today
mariaschuld: Thanks @co9olguy, I hope I can convince you to not add too many more changes :) 

Patch:
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
index c85ce07b0..2b7cda011 100644
--- a/.github/CHANGELOG.md
+++ b/.github/CHANGELOG.md
@@ -78,7 +78,6 @@
   >>> qnode2 = qml.QNode(circuit, dev, diff_method="parameter-shift", argnum=[0])
   >>> print(qml.grad(qnode2)(x,y))
   (array(0.31434679), array(0.))
->>>>>>> master
   ```
 
 * The `quantum_monte_carlo` transform has been added, allowing an input circuit to be transformed
@@ -156,8 +155,9 @@
   that help with investigating the Fourier representation of functions
   implemented by quantum circuits.
   [(#1160)](https://github.com/PennyLaneAI/pennylane/pull/1160)
+  [(#1378)](https://github.com/PennyLaneAI/pennylane/pull/1378)
 
-  For example, one can plot distributions over Fourier series coefficients like
+  For example, one can plot distributions over Fourier series coefficients like 
   this one:
 
   <img src="https://pennylane.readthedocs.io/en/latest/_static/fourier.png" width=70%/>
diff --git a/doc/code/qml_fourier.rst b/doc/code/qml_fourier.rst
index 73669c128..b0cbf8d61 100644
--- a/doc/code/qml_fourier.rst
+++ b/doc/code/qml_fourier.rst
@@ -87,10 +87,70 @@ available to a quantum model, the larger the class of functions that model can
 represent, potentially leading to greater utility for quantum machine learning
 applications.
 
+Calculating the frequencies supported by a circuit
+--------------------------------------------------
+
+For certain circuits, information on the frequency spectra :math:`\Omega_j`
+can be derived solely from the structure of the gates
+that encode the corresponding inputs :math:`x_j` (see for example `Schuld,
+Sweke, and Meyer (2020) <https://arxiv.org/abs/2008.08605>`__). More precisely, if
+all input-encoding gates are of the form :math:`e^{-ix_j G}`, where :math:`G` is
+a Hermitian operator that "generates" the operation, we can deduce a maximum set of frequencies
+that can theoretically appear in :math:`\Omega_j`.  Depending on the non-input-encoding
+gates in the circuit, some of these theoretically supported frequencies may end up
+having vanishing Fourier coefficients, and :math:`\Omega_j` effectively turns out to be smaller.
+However, estimates based on the input-encoding strategy can still be useful to understand
+the potential expressivity of a type of ansatz.
+
+The theoretically supported frequencies can be computed
+using the :func:`~.pennylane.fourier.spectrum` function. To mark which gates encode
+inputs (and, for example, which ones are only used for trainable parameters), we
+have to give input-encoding gates an ``id``:
+
+.. code::
+
+    import pennylane as qml
+
+    dev = qml.device('default.qubit', wires=2)
+
+    @qml.qnode(dev)
+    def simple_circuit_marked(x):
+        qml.RX(x[0], wires=0, id="x")
+        qml.RY(x[0], wires=1, id="x")
+        qml.CNOT(wires=[1, 0])
+        return qml.expval(qml.PauliZ(0))
+
+We can then compute the frequencies supported by the input-encoding gates as:
+
+.. code::
+
+   >>> from pennylane.fourier import spectrum
+   >>> freqs = spectrum(simple_circuit_marked)([0.1])
+   >>> for k, v in freqs.items():
+   >>>     print(k, ":", v)
+
+   x : [-2.0, -1.0, 0.0, 1.0, 2.0]
+
+.. note::
+
+    Some encoding-gate types may give rise to non-integer-valued frequencies. In this case,
+    the :func:`~.pennylane.fourier.spectrum` function computes the frequency sets :math:`\Omega_j`
+    of the *Fourier sum* of the form
+
+    .. math::
+
+        f(x) = \sum \limits_{\omega_1\in \Omega_1} \dots \sum \limits_{\omega_N \in \Omega_N}
+        c_{\omega_1,\dots, \omega_N} e^{-i x_1 \omega_1} \dots e^{-i x_N \omega_N},
+
+    with :math:`\omega \in \mathbb{R}`. Just like any other function, such a sum can be turned into a
+    proper Fourier series, but the mapping from non-integer-valued to integer-valued frequencies
+    is not always trivial.
+
 Calculating the Fourier coefficients
 ------------------------------------
 
-The Fourier coefficients can be numerically calculated with the
+To get a more accurate picture of the Fourier series representation of a quantum circuit,
+we have to compute the Fourier coefficients :math:`c_{n_1,\ldots,n_N}`. This is done using numerical methods in the
 :func:`~.pennylane.fourier.coefficients` function:
 
 .. code::
@@ -120,7 +180,7 @@ For more details and examples of coefficient calculation, please see the
 documentation for :func:`~.pennylane.fourier.coefficients`.
 
 Fourier coefficient visualization
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 A key application of the Fourier module is to analyze the *expressivity* of
 classes of quantum circuit families. The set of frequencies in the Fourier
diff --git a/pennylane/fourier/__init__.py b/pennylane/fourier/__init__.py
index 215f1c211..b650c6b84 100644
--- a/pennylane/fourier/__init__.py
+++ b/pennylane/fourier/__init__.py
@@ -14,3 +14,4 @@
 """This module contains functions to analyze the Fourier representation
 of quantum circuits."""
 from .coefficients import coefficients
+from .spectrum import spectrum
diff --git a/pennylane/fourier/spectrum.py b/pennylane/fourier/spectrum.py
new file mode 100644
index 000000000..c7809f427
--- /dev/null
+++ b/pennylane/fourier/spectrum.py
@@ -0,0 +1,285 @@
+# 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 transform that computes the frequency spectrum of a quantum
+circuit."""
+from itertools import chain, combinations
+from functools import wraps
+import numpy as np
+import pennylane as qml
+
+
+def _get_spectrum(op):
+    r"""Extract the frequencies contributed by an input-encoding gate to the
+    overall Fourier representation of a quantum circuit.
+
+    If :math:`G` is the generator of the input-encoding gate :math:`\exp(-i x G)`,
+    the frequencies are the differences between any two of :math:`G`'s eigenvalues.
+
+    Args:
+        op (~pennylane.operation.Operation): an instance of the `Operation` class
+
+    Returns:
+        list: frequencies contributed by this input-encoding gate
+    """
+    no_generator = False
+    if hasattr(op, "generator"):
+        g, coeff = op.generator
+
+        if isinstance(g, np.ndarray):
+            matrix = g
+        elif hasattr(g, "matrix"):
+            matrix = g.matrix
+        else:
+            no_generator = True
+    else:
+        no_generator = True
+
+    if no_generator:
+        raise ValueError(f"Generator of operation {op} is not defined.")
+
+    matrix = coeff * matrix
+    # eigenvalues of hermitian ops are guaranteed to be real
+    # todo: use qml.math.linalg once it is tested properly
+    evals = qml.math.real(np.linalg.eigvals(matrix))
+
+    # compute all differences of eigenvalues
+    unique_frequencies = set(
+        chain.from_iterable(
+            np.round((x[1] - x[0], x[0] - x[1]), decimals=8) for x in combinations(evals, 2)
+        )
+    )
+    unique_frequencies = unique_frequencies.union({0})
+    return sorted(unique_frequencies)
+
+
+def _join_spectra(spec1, spec2):
+    r"""Join two sets of frequencies that belong to the same input.
+
+    Since :math:`\exp(i a x)\exp(i b x) = \exp(i (a+b) x)`, frequency sets of two gates
+    encoding the same :math:`x` are joined by computing the set of sums of their elements.
+
+    Args:
+        spec1 (list[float]): first spectrum
+        spec2 (list[float]): second spectrum
+    Returns:
+        list[float]: joined spectrum
+    """
+    if spec1 == []:
+        return sorted(set(spec2))
+    if spec2 == []:
+        return sorted(set(spec1))
+
+    sums = [s1 + s2 for s1 in spec1 for s2 in spec2]
+    return sorted(set(sums))
+
+
+def spectrum(qnode, encoding_gates=None):
+    r"""Compute the frequency spectrum of the Fourier representation of simple quantum circuits.
+
+    The circuit must only use simple single-parameter gates of the form :math:`e^{-i x_j G}` as
+    input-encoding gates, which allows the computation of the spectrum by inspecting the gates'
+    generators :math:`G`. The most important example of such gates are Pauli rotations.
+
+    .. note::
+
+        More precisely, the spectrum function relies on the gate to define a ``generator``, and will
+        fail if gates marked as inputs do not have this attribute.
+
+    Gates are marked as input-encoding gates in the quantum function by giving them an ``id``.
+    If two gates have the same ``id``, they are considered
+    to be used to encode the same input :math:`x_j`. The ``encoding_gates`` argument can be used
+    to indicate that only gates with a specific ``id`` should be interpreted as input-encoding gates.
+    Otherwise, all gates with an explicit ``id`` are considered to be input-encoding gates.
+
+    .. note::
+        If no input-encoding gates are found, an empty dictionary is returned.
+
+    Args:
+        qnode (pennylane.QNode): a quantum node representing a circuit in which
+            input-encoding gates are marked by their ``id`` attribute
+        encoding_gates (list[str]): list of input-encoding gate ``id`` strings
+            for which to compute the frequency spectra
+
+    Returns:
+        (dict[str, list[float]]): Dictionary with the input-encoding gate ``id`` as keys and
+            their frequency spectra as values.
+
+
+    **Details**
+
+    A circuit that returns an expectation value which depends on
+    :math:`N` scalar inputs :math:`x_j` can be interpreted as a function
+    :math:`f: \mathbb{R}^N \rightarrow \mathbb{R}`. This function can always be
+    expressed by a Fourier-type sum
+
+    .. math::
+
+        \sum \limits_{\omega_1\in \Omega_1} \dots \sum \limits_{\omega_N \in \Omega_N}
+        c_{\omega_1,\dots, \omega_N} e^{-i x_1 \omega_1} \dots e^{-i x_N \omega_N}
+
+    over the *frequency spectra* :math:`\Omega_j \subseteq \mathbb{R},`
+    :math:`j=1,\dots,N`. Each spectrum has the property that
+    :math:`0 \in \Omega_j`, and the spectrum is
+    symmetric (for every :math:`\omega \in \Omega_j` we have that :math:`-\omega \in
+    \Omega_j`). If all frequencies are integer-valued, the Fourier sum becomes a
+    *Fourier series*.
+
+    As shown in `Vidal and Theis (2019)
+    <https://arxiv.org/abs/1901.11434>`_ and `Schuld, Sweke and Meyer (2020)
+    <https://arxiv.org/abs/2008.08605>`_, if an input :math:`x_j, j = 1 \dots N`,
+    only enters into single-parameter gates of the form :math:`e^{-i x_j G}` (where :math:`G` is a Hermitian generator),
+    the frequency spectrum :math:`\Omega_j` is fully determined by the eigenvalues
+    of :math:`G`. In many situations, the spectra are limited
+    to a few frequencies only, which in turn limits the function class that the circuit
+    can express.
+
+    The ``spectrum`` function computes all frequencies that will potentially appear in the
+    sets :math:`\Omega_1` to :math:`\Omega_N`.
+
+    **Example**
+
+    Consider the following example, which uses non-trainable inputs ``x`` and
+    trainable parameters ``w`` as arguments to the qnode.
+
+    .. code-block:: python
+
+        import pennylane as qml
+        import numpy as np
+        from pennylane.fourier import spectrum
+
+        n_layers = 2
+        n_qubits = 3
+        dev = qml.device("default.qubit", wires=n_qubits)
+
+        @qml.qnode(dev)
+        def circuit(x, w):
+            for l in range(n_layers):
+                for i in range(n_qubits):
+                    qml.RX(x[i], wires=0, id="x"+str(i))
+                    qml.Rot(w[l,i,0], w[l,i,1], w[l,i,2], wires=0)
+            qml.RZ(x[0], wires=0, id="x0")
+            return qml.expval(qml.PauliZ(wires=0))
+
+        x = np.array([1, 2, 3])
+        w = np.random.random((n_layers, n_qubits, 3))
+        res = spectrum(circuit)(x, w)
+
+    >>> print(qml.draw(circuit)(x, w))
+    0: ──RX(1)──Rot(0.863, 0.611, 0.281)───RX(1)──Rot(0.47, 0.158, 0.648)───RZ(1)──┤ ⟨Z⟩
+    1: ──RX(2)──Rot(0.0781, 0.971, 0.457)──RX(2)──Rot(0.896, 0.224, 0.731)─────────┤
+    2: ──RX(3)──Rot(0.462, 0.286, 0.929)───RX(3)──Rot(0.879, 0.399, 0.215)─────────┤
+
+    >>> for inp, freqs in res.items():
+    >>>     print(f"{inp}: {freqs}")
+    'x0': [-3.0, -2.0, -1.0, 0.0, 1.0, 2.0, 3.0]
+    'x1': [-2.0, -1.0, 0.0, 1.0, 2.0]
+    'x2': [-2.0, -1.0, 0.0, 1.0, 2.0]
+
+    .. note::
+        While the Fourier spectrum usually does not depend
+        on trainable circuit parameters or the actual values of the inputs,
+        it may still change based on inputs to the QNode that alter the architecture
+        of the circuit.
+
+    The input-encoding gates to consider can also be explicitly selected by using the
+    ``encoding_gates`` keyword argument:
+
+    .. code-block:: python
+
+        dev = qml.device("default.qubit", wires=1)
+
+        @qml.qnode(dev)
+        def circuit(x):
+            qml.RX(x[0], wires=0, id="x0")
+            qml.PhaseShift(x[0], wires=0, id="x0")
+            qml.RX(x[1], wires=0, id="x1")
+            return qml.expval(qml.PauliZ(wires=0))
+
+        x = np.array([1, 2])
+        res = spectrum(circuit, encoding_gates=["x0"])(x)
+
+    >>> for inp, freqs in res.items():
+    >>>     print(f"{inp}: {freqs}")
+    'x0': [-2.0, -1.0, 0.0, 1.0, 2.0]
+
+    .. note::
+        The ``spectrum`` function does not check if the result of the
+        circuit is an expectation, or if gates with the same ``id``
+        take the same value in a given call of the function.
+
+    The ``spectrum`` function works in all interfaces:
+
+    .. code-block:: python
+
+        import tensorflow as tf
+
+        dev = qml.device("default.qubit", wires=1)
+
+        @qml.qnode(dev, interface='tf')
+        def circuit(x):
+            qml.RX(x[0], wires=0, id="x0")
+            qml.PhaseShift(x[1], wires=0, id="x1")
+            return qml.expval(qml.PauliZ(wires=0))
+
+        x = tf.constant([1, 2])
+        res = spectrum(circuit)(x)
+
+    >>> for inp, freqs in res.items():
+    >>>     print(f"{inp}: {freqs}")
+    'x0': [-1.0, 0.0, 1.0]
+    'x1': [-1.0, 0.0, 1.0]
+
+    """
+
+    @wraps(qnode)
+    def wrapper(*args, **kwargs):
+        qnode.construct(args, kwargs)
+        tape = qnode.qtape
+
+        if encoding_gates is None:
+            freqs = {}
+        else:
+            freqs = {input_id: [] for input_id in encoding_gates}
+
+        for op in tape.operations:
+            id = op.id
+
+            # if the operator has no specific ID, move to the next
+            if id is None:
+                continue
+
+            # if user has not specified encoding_gate id's,
+            # consider any id
+            is_encoding_gate = encoding_gates is None or id in encoding_gates
+
+            if is_encoding_gate:
+
+                if len(op.parameters) != 1:
+                    raise ValueError(
+                        "Can only consider one-parameter gates as data-encoding gates; "
+                        f"got {op.name}."
+                    )
+
+                spec = _get_spectrum(op)
+
+                # if id has been seen before,
+                # join this spectrum to another one
+                if id in freqs:
+                    spec = _join_spectra(freqs[id], spec)
+
+                freqs[id] = spec
+
+        return freqs
+
+    return wrapper


Test patch:
diff --git a/tests/fourier/test_spectrum.py b/tests/fourier/test_spectrum.py
new file mode 100644
index 000000000..d5e6df9c7
--- /dev/null
+++ b/tests/fourier/test_spectrum.py
@@ -0,0 +1,240 @@
+# 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 spectrum transform.
+"""
+import pytest
+import numpy as np
+import pennylane as qml
+from pennylane import numpy as pnp
+from pennylane.fourier.spectrum import spectrum, _join_spectra, _get_spectrum
+
+
+class TestHelpers:
+    @pytest.mark.parametrize(
+        "spectrum1, spectrum2, expected",
+        [
+            ([-1, 0, 1], [-1, 0, 1], [-2, -1, 0, 1, 2]),
+            ([-3, 0, 3], [-5, 0, 5], [-8, -5, -3, -2, 0, 2, 3, 5, 8]),
+            ([-2, -1, 0, 1, 2], [-1, 0, 1], [-3, -2, -1, 0, 1, 2, 3]),
+            ([-0.5, 0, 0.5], [-1, 0, 1], [-1.5, -1, -0.5, 0, 0.5, 1.0, 1.5]),
+            ([-0.5, 0, 0.5], [], [-0.5, 0, 0.5]),
+            ([], [-0.5, 0, 0.5], [-0.5, 0, 0.5]),
+        ],
+    )
+    def test_join_spectra(self, spectrum1, spectrum2, expected, tol):
+        """Test that spectra are joined correctly."""
+        joined = _join_spectra(spectrum1, spectrum2)
+        assert np.allclose(joined, expected, atol=tol, rtol=0)
+
+    @pytest.mark.parametrize(
+        "op, expected",
+        [
+            (qml.RX(0.1, wires=0), [-1, 0, 1]),  # generator is a class
+            (qml.RY(0.1, wires=0), [-1, 0, 1]),  # generator is a class
+            (qml.RZ(0.1, wires=0), [-1, 0, 1]),  # generator is a class
+            (qml.PhaseShift(0.5, wires=0), [-1, 0, 1]),  # generator is an array
+            (qml.ControlledPhaseShift(0.5, wires=[0, 1]), [-1, 0, 1]),  # generator is an array
+        ],
+    )
+    def test_get_spectrum(self, op, expected, tol):
+        """Test that the spectrum is correctly extracted from an operator."""
+        spec = _get_spectrum(op)
+        assert np.allclose(spec, expected, atol=tol, rtol=0)
+
+    def test_get_spectrum_complains_no_generator(self):
+        """Test that an error is raised if the operator has no generator defined."""
+
+        # Observables have no generator attribute
+        with pytest.raises(ValueError, match="Generator of operation"):
+            _get_spectrum(qml.P(wires=0))
+
+        # CNOT is an operation where generator is an abstract property
+        with pytest.raises(ValueError, match="Generator of operation"):
+            _get_spectrum(qml.CNOT(wires=[0, 1]))
+
+
+class TestCircuits:
+    """Tests that the spectrum is returned as expected."""
+
+    @pytest.mark.parametrize("n_layers, n_qubits", [(1, 1), (2, 3), (4, 1)])
+    def test_spectrum_grows_with_gates(self, n_layers, n_qubits):
+        """Test that the spectrum grows linearly with the number of
+        encoding gates if we use Pauli rotation encoding."""
+
+        dev = qml.device("default.qubit", wires=n_qubits)
+
+        @qml.qnode(dev)
+        def circuit(x):
+            for l in range(n_layers):
+                for i in range(n_qubits):
+                    qml.RX(x, wires=i, id="x")
+                    qml.RY(0.4, wires=i)
+            return qml.expval(qml.PauliZ(wires=0))
+
+        res = spectrum(circuit)(0.1)
+        expected_degree = n_qubits * n_layers
+        assert np.allclose(res["x"], range(-expected_degree, expected_degree + 1))
+
+    def test_encoding_gates(self):
+        """Test that the spectrum contains the ids provided in encoding_gates, or
+        all ids if encoding_gates is None."""
+
+        dev = qml.device("default.qubit", wires=1)
+
+        @qml.qnode(dev)
+        def circuit(x):
+            qml.RX(x, wires=0, id="x")
+            qml.RY(0.4, wires=0, id="other")
+            return qml.expval(qml.PauliZ(wires=0))
+
+        res = spectrum(circuit, encoding_gates=["x"])(0.1)
+        assert res == {"x": [-1.0, 0.0, 1.0]}
+
+        res = spectrum(circuit, encoding_gates=["x", "other"])(0.1)
+        assert res == {"x": [-1.0, 0.0, 1.0], "other": [-1.0, 0.0, 1.0]}
+
+        res = spectrum(circuit)(0.1)
+        assert res == {"x": [-1.0, 0.0, 1.0], "other": [-1.0, 0.0, 1.0]}
+
+        res = spectrum(circuit, encoding_gates=["a"])(0.1)
+        assert res == {"a": []}
+
+    def test_spectrum_changes_with_qnode_args(self):
+        """Test that the spectrum changes per call if a qnode argument changes the
+        circuit architecture."""
+
+        dev = qml.device("default.qubit", wires=3)
+
+        @qml.qnode(dev)
+        def circuit(last_gate):
+            qml.RX(0.1, wires=0, id="x")
+            qml.RX(0.2, wires=1, id="x")
+            if last_gate:
+                qml.RX(0.3, wires=2, id="x")
+            return qml.expval(qml.PauliZ(wires=0))
+
+        res_true = spectrum(circuit)(True)
+        assert np.allclose(res_true["x"], range(-3, 4))
+
+        res_false = spectrum(circuit)(False)
+        assert np.allclose(res_false["x"], range(-2, 3))
+
+    def test_input_gates_not_of_correct_form(self):
+        """Test that an error is thrown if gates marked as encoding gates
+        are not single-parameter gates."""
+
+        dev = qml.device("default.qubit", wires=3)
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.RX(0.1, wires=0, id="x")
+            qml.Rot(0.2, 0.3, 0.4, wires=1, id="x")
+            return qml.expval(qml.PauliZ(wires=0))
+
+        with pytest.raises(ValueError, match="Can only consider one-parameter gates"):
+            spectrum(circuit)()
+
+
+def circuit(x, w):
+    """Test circuit"""
+    for l in range(2):
+        for i in range(3):
+            qml.RX(x[i], wires=0, id="x" + str(i))
+            qml.RY(w[l][i], wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.CNOT(wires=[1, 2])
+    qml.RZ(x[0], wires=0, id="x0")
+    return qml.expval(qml.PauliZ(wires=0))
+
+
+expected_result = {
+    "x0": [-3.0, -2.0, -1.0, 0.0, 1.0, 2.0, 3.0],
+    "x1": [-2.0, -1.0, 0.0, 1.0, 2.0],
+    "x2": [-2.0, -1.0, 0.0, 1.0, 2.0],
+}
+
+
+class TestInterfaces:
+    """Test that inputs are correctly identified and spectra computed in
+    all interfaces."""
+
+    def test_integration_autograd(self):
+        """Test that the spectra of a circuit is calculated correctly
+        in the autograd interface."""
+
+        x = pnp.array([1.0, 2.0, 3.0], requires_grad=False)
+        w = pnp.array([[-1, -2, -3], [-4, -5, -6]], requires_grad=True)
+
+        dev = qml.device("default.qubit", wires=3)
+        qnode = qml.QNode(circuit, dev, interface="autograd")
+
+        res = spectrum(qnode)(x, w)
+        for (k1, v1), (k2, v2) in zip(res.items(), expected_result.items()):
+            assert k1 == k2
+            assert v1 == v2
+
+    def test_integration_torch(self):
+        """Test that the spectra of a circuit is calculated correctly
+        in the torch interface."""
+
+        torch = pytest.importorskip("torch")
+        x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
+        w = torch.tensor([[-1, -2, -3], [-4, -5, -6]], requires_grad=False)
+
+        dev = qml.device("default.qubit", wires=3)
+        qnode = qml.QNode(circuit, dev, interface="torch")
+
+        res = spectrum(qnode)(x, w)
+        assert res
+        for (k1, v1), (k2, v2) in zip(res.items(), expected_result.items()):
+            assert k1 == k2
+            assert v1 == v2
+
+    def test_integration_tf(self):
+        """Test that the spectra of a circuit is calculated correctly
+        in the tf interface."""
+        tf = pytest.importorskip("tensorflow")
+
+        dev = qml.device("default.qubit", wires=3)
+        qnode = qml.QNode(circuit, dev, interface="tf")
+
+        x = tf.Variable([1.0, 2.0, 3.0])
+        w = tf.constant([[-1, -2, -3], [-4, -5, -6]])
+        res = spectrum(qnode)(x, w)
+
+        assert res
+        for (k1, v1), (k2, v2) in zip(res.items(), expected_result.items()):
+            assert k1 == k2
+            assert v1 == v2
+
+    def test_integration_jax(self):
+        """Test that the spectra of a circuit is calculated correctly
+        in the jax interface."""
+
+        jax = pytest.importorskip("jax")
+        from jax import numpy as jnp
+
+        x = jnp.array([1.0, 2.0, 3.0])
+        w = [[-1, -2, -3], [-4, -5, -6]]
+
+        dev = qml.device("default.qubit", wires=3)
+        qnode = qml.QNode(circuit, dev, interface="jax")
+
+        res = spectrum(qnode)(x, w)
+
+        assert res
+        for (k1, v1), (k2, v2) in zip(res.items(), expected_result.items()):
+            assert k1 == k2
+            assert v1 == v2