Repository: PennyLaneAI/pennylane
Pull request #3373 (2022-11-16 19:38:53)

Implement QPE for operators
**Context:**

Currently the [QPE](https://pennylane.readthedocs.io/en/stable/code/api/pennylane.QuantumPhaseEstimation.html) template only supports operator pass-through as a matrix. It would be useful to be able to pass a quantum function and not have to think in arrays.

**Description of the Change:**

QPE is updated to take an `Operator` as the input unitary to perform phase estimation with.

**Benefits:**

Now we can perform QPE on a quantum operator

**Possible Drawbacks:**

This is a new feature added, which does not impact existing functionalities. However, it is worth noting that there are two potential approaches to performing phase estimation with a quantum function, specifically when constructing powers of the unitary. The first approach is to find the matrix representation of the circuit and calculate the powers numerically. 

This PR takes the second approach, which is to simply apply the quantum function multiple times. Admittedly, using this approach, the number of gate operations would scale exponentially and not linearly with the number of estimation bits, but I would argue that the second approach is closer to the intended implementation of QPE. 

When we analyze the theoretical complexity of QPE, we always treat the unitary as an oracle and discuss the complexity in terms of calls to the controlled unitary. I believe that the implementation here should reflect that, as opposed to taking the shortcut of calculating the matrix powers numerically.

**Related GitHub Issues:**
https://github.com/PennyLaneAI/pennylane/issues/2935
https://github.com/PennyLaneAI/pennylane/issues/3407



albi3ro: Hi @astralcai  and thanks for the pull request 🎉 

@KetpuntoG and I have clarified the problem description and the desired behaviour in the Issue.  You can see my comment on the problem description at: https://github.com/PennyLaneAI/pennylane/issues/2935#issuecomment-1318736193
astralcai: Hi @albi3ro,

I'm glad that we agree on using operations and not matrices when applying the controlled unitary.

This new call signature would definitely be a lot easier to implement, but there is one thing in my current implementation that I would like to keep. I imagine a list of operators can be defined on any set of wire labels, and once it is defined, it can be reused in many different places. Let's say someone decides to define a list of operations on wires 0 and 1, maybe we don't want to constrain the user to have to choose 0 and 1 as target wires in their QPE circuit (maybe their QPE is part of a larger circuit, and wires 0 and 1 are not even part of the QPE circuit). That is why I have this wire map in my code, let me know what you think of this.

Best,

Astral
albi3ro: Our new `qml.map_wires` function (which I see you are using 😄 ) acts on individual Operators, Measurement Processes, quantum scripts, quantum functions, or QNodes.  So we can still map wires on individual operators.

We can still define one `Operator`, and use it in multiple places with mapped wires through `qml.map_wires`.

I'm also wary of implicitly mapping wires under the hood.   I could see this leading to hard-to-diagnose unexpected behaviour. If the target wires are somehow described by the object provided by `unitary`, they should *only* be provided through `unitary`.

One thing we also need to start popularizing and advertising more, we can turn a list of operations into a `Prod` Operator.  For example,  the target operator could be:

```
qml.RX(x, 0) @ qml.RY(y, 1)
```
which would be equivalent to:
```
def qfunc():
    qml.RY(y, 1)
    qml.RY(x, 0)
```

As a final note, since the information on `target_wires` would be stored in the `unitary` operation, we would also need to override the `Operator.map_wires` method to update the wires on the base operation as well.
astralcai: Hi, how about if target wires are not provided, we use the wires defined in the unitary, but if target wires are explicitly provided, we do the wire map to the target wires? We can clearly explain this behaviour in the docstring of the QPE method. 
astralcai: Hi @albi3ro,

I just made the update. Now the QPE can take a single operator or a list of operators as the unitary with no target_wires specified. However, I kept the logic where if target wires are specified in addition to the operators, I would still map the operators to the target wires. 

Best,

Astral
codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3373?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#3373](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3373?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (c56c09c) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/a6049f21338ad1a61f0d33cca9f0bc449f764ea7?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (a6049f2) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

> :exclamation: Current head c56c09c differs from pull request most recent head 72c8e71. Consider uploading reports for the commit 72c8e71 to get more accurate results

```diff
@@           Coverage Diff           @@
##           master    #3373   +/-   ##
=======================================
  Coverage   99.70%   99.70%           
=======================================
  Files         298      298           
  Lines       26159    26195   +36     
=======================================
+ Hits        26081    26117   +36     
  Misses         78       78           
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3373?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/templates/subroutines/qpe.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3373/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RlbXBsYXRlcy9zdWJyb3V0aW5lcy9xcGUucHk=) | `100.00% <100.00%> (ø)` | |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)

albi3ro: ### Switching the order of arguments in the call signature:
Switching `target_wires` and `estimation wires` in the call signature would be a breaking change that could lead to a lot of hard-to-diagnose problems.  Instead, we can go with the simpler solution: giving both `target_wires` and `estimation_wires` default values of `None`.  While `estimation_wires` would then need to be provided as a keyword argument, that is perfectly reasonable restriction.
```
QPE(op, estimation_wires=[1,2,3])
```
This leads to the call signature:
```
QPE(unitary: Union[ndarray, Operator], target_wires=None, estimation_wires=None, do_queue=True, id=None)
```

### Mapping the wires internally

Deciding whether or not to map the wires internally hinges on the answers to two questions:

1) How hard is it to let the user control mapping the wires themselves?
2) What is the probability of bugs and unexpected behavior if we map the wires internally? How hard would these problems be to debug?

With the release of `qml.map_wires`, users can now easily map the wires themselves. 

Even if we document the behavior of the internal wire mapping, not everyone will not read the docstring closely enough. If wires are mistakenly mapped, this could be tricky to diagnose. I'd rather preemptively avoid those situations.

We should not map the wires internally. Instead, if `unitary` is an `Operator` and `target_wires` is specified, we should raise an error.

### Accepting a list of Operators

We will already be overloading the `unitary` keyword with both `array` and `Operator`. Including `list[Operator]` will increase the complexity of the interface and behavior with only minimal improvements to functionality. As we start expanding the use of `qml.prod` and working with the product of operators, the use cases for `list[Operator]` will decrease even more.

We should not accept a list of operators.  That adds too much complexity to the interface.

### Map wires method

You can check out the source code for [`Controlled.map_wires`](https://github.com/PennyLaneAI/pennylane/blob/ffa6e46aae3ab1c94c8c9c8d2761a900ca91abd0/pennylane/ops/op_math/controlled_class.py#L276) for an example implementation.  This instance method should create a  new operation that is a copy except for mapped wires. `Controlled` is another class that has a target operation and additional "control"/ "estimation" wires.

Though we need to update the method to touching the private `_wires` attribute.  Just recursively call `op.map_wires(map)` isntead.
astralcai: Hi @albi3ro,

I am trying to implement QPE using the ```qml.pow``` operator as you suggested, but have encountered an issue. If the unitary is specified as a ```Prod```, such as ```qml.RX(param, wires=[0]) @ qml.CNOT(wires=[0, 1])```, then during decomposition of ```(qml.RX(param, wires=[0]) @ qml.CNOT(wires=[0, 1])) ** 1```, we call ```Pow.decomposition()```, which then calls ```Prod.pow``` (line 283 of ```ops/op_math/pow.py```, where ```self.base``` is a ```Prod``` as defined earlier). However, since ```Prod.pow``` is not defined, it defaults to ```Operator.pow```, which returns ```[copy.copy(self)]``` (line 1041 of ```operations.py```), but this is not queued. 

As a result, in ```Operator.expand()```, when we first create a tape (line 1094 of ```operations.py```), and call ```self.decomposition()``` (line 1100), the result of that decomposition is not added to the tape, which leads to the final tape simply missing the single power of the unitary.

I imagine we can override ```Prod.pow``` with the same logic as ```Operator.pow```, except that we return ```[qml.apply(self)]``` instead of ```[copy.copy(self)]``` if ```QueuingManager.recording()```, similar to the logic in ```Pow.decomposition``` (line 286 of ```ops/op_math/pow.py```). What do you suggest?

Best,

Astral

albi3ro: @astralcai  This appears to be a bug.  Thank you for finding it. Would you mind opening a bug report issue on this? That way we can track what bugs exist and were fixed.

This problem extends far beyond `Prod`, and can be cleanly described by this:

```pycon
>>> qml.pow(qml.PauliX(0), 1).expand().circuit
[]
>>> qml.pow(qml.PauliX(0), 1).decomposition()
[PauliX(wires=[0])]
```

Basically, `Operator.pow` needs to queue all its constituents. As you stated, this can be solved by using `apply` instead of `copy` in the default behaviour.
astralcai: Hi @albi3ro,

If I simply update ```Operator.pow``` to use ```apply``` instead of ```copy```, the issue would be fixed. Do I still need to open a bug report in this case?

Best,

Astral
albi3ro: Yes please.  It helps us keep track of what things have required bug fixes. A bug fix would also require dedicated tests and a changelog entry.  

I've also realized we will need to override the `queue` method as well, see [`SymbolicOp.queue`](https://github.com/PennyLaneAI/pennylane/blob/6881ca8076983f4f324492cfeb8edb10ec06a5be/pennylane/ops/op_math/symbolicop.py#L112) for an example.  `SymbolicOp` is the base class for "operator that contains another operator", which is what `QPE` will now be.  Basically, if the target operation is constructed inside the quantum function, we don't want it to end up in the final circuit, so we need to make the target operation as "owned".
astralcai: > Yes please. It helps us keep track of what things have required bug fixes. A bug fix would also require dedicated tests and a changelog entry.
> 
> I've also realized we will need to override the `queue` method as well, see [`SymbolicOp.queue`](https://github.com/PennyLaneAI/pennylane/blob/6881ca8076983f4f324492cfeb8edb10ec06a5be/pennylane/ops/op_math/symbolicop.py#L112) for an example. `SymbolicOp` is the base class for "operator that contains another operator", which is what `QPE` will now be. Basically, if the target operation is constructed inside the quantum function, we don't want it to end up in the final circuit, so we need to make the target operation as "owned".

Makes sense, I will add the `queue` method override.
astralcai: @albi3ro 

I made the update. One test failed, but the failing test case has nothing to do with QPE, and I can't even replicate the fail on my own machine.
astralcai: > @astralcai I updated with master, and we'll see if that fixes the CI problem. If not, I'll debug it myself this afternoon.
> 
> Thanks for your work on this PR!

@albi3ro seems that the issue is gone now. I just merged from master again. I'm actually going out this afternoon and won't be able to work on this for the rest of the day. Hopefully you can have it merged. Thanks for all the support :)

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 61da9fd86..aed8a4a2c 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -368,6 +368,9 @@
 * Remove private `_wires` setter from the `Controlled.map_wires` method.
   [(#3405)](https://github.com/PennyLaneAI/pennylane/pull/3405)
 
+* QPE can now accept a target operator instead of a matrix and target wires pair.
+  [#3373](https://github.com/PennyLaneAI/pennylane/pull/3373)
+
 * `QuantumTape._process_queue` has been moved to `qml.queuing.process_queue` to disentangle
   its functionality from the `QuantumTape` class.
   [(#3401)](https://github.com/PennyLaneAI/pennylane/pull/3401)
@@ -815,7 +818,11 @@ Deprecations cycles are tracked at [doc/developement/deprecations.rst](https://d
 * The `pad_with` argument in the `qml.AmplitudeEmbedding` template is now compatible with all interfaces
   [(#3392)](https://github.com/PennyLaneAI/pennylane/pull/3392)
 
-* Fixed a bug where a QNode returning `qml.sample` would produce incorrect results when run on a device defined with a shot vector.
+* `Operator.pow` now queues its constituents by default.
+  [(#3373)](https://github.com/PennyLaneAI/pennylane/pull/3373)
+
+* Fixed a bug where a QNode returning `qml.sample` would produce incorrect results when
+  run on a device defined with a shot vector.
   [(#3422)](https://github.com/PennyLaneAI/pennylane/pull/3422)
 
 * The `qml.data` module now works as expected on Windows.
diff --git a/pennylane/operation.py b/pennylane/operation.py
index 61ee568a5..05b0d4391 100644
--- a/pennylane/operation.py
+++ b/pennylane/operation.py
@@ -1190,6 +1190,8 @@ class Operator(abc.ABC):
         if z == 0:
             return []
         if z == 1:
+            if QueuingManager.recording():
+                return [qml.apply(self)]
             return [copy.copy(self)]
         raise PowUndefinedError
 
diff --git a/pennylane/ops/op_math/symbolicop.py b/pennylane/ops/op_math/symbolicop.py
index 8f4fb41d9..e1b6565ea 100644
--- a/pennylane/ops/op_math/symbolicop.py
+++ b/pennylane/ops/op_math/symbolicop.py
@@ -33,8 +33,8 @@ class SymbolicOp(Operator):
     This *developer-facing* class can serve as a parent to single base symbolic operators, such as
     :class:`~.ops.op_math.Adjoint` and :class:`~.ops.op_math.Pow`.
 
-    New symbolic operators can inherit from this class to recieve some common default behavior, such
-    as deferring properties to the the base class, copying the base class during a shallow copy, and
+    New symbolic operators can inherit from this class to receive some common default behavior, such
+    as deferring properties to the base class, copying the base class during a shallow copy, and
     updating the metadata of the base operator during queueing.
 
     The child symbolic operator should define the `_name` property during initialization and define
diff --git a/pennylane/templates/subroutines/qpe.py b/pennylane/templates/subroutines/qpe.py
index 8a7dc77d9..748af272a 100644
--- a/pennylane/templates/subroutines/qpe.py
+++ b/pennylane/templates/subroutines/qpe.py
@@ -16,8 +16,8 @@ Contains the QuantumPhaseEstimation template.
 """
 # pylint: disable=too-many-arguments,arguments-differ
 import pennylane as qml
-from pennylane.operation import AnyWires, Operation
-from pennylane.ops import Hadamard, ControlledQubitUnitary
+from pennylane.queuing import QueuingManager
+from pennylane.operation import AnyWires, Operation, Operator
 
 
 class QuantumPhaseEstimation(Operation):
@@ -53,14 +53,17 @@ class QuantumPhaseEstimation(Operation):
     usage details below give an example of this case.
 
     Args:
-        unitary (array): the phase estimation unitary, specified as a matrix
-        target_wires (Union[Wires, Sequence[int], or int]): the target wires to apply the unitary
+        unitary (array or Operator): the phase estimation unitary, specified as a matrix or an
+            :class:`~.Operator`
+        target_wires (Union[Wires, Sequence[int], or int]): the target wires to apply the unitary.
+            If the unitary is specified as an operator, the target wires should already have been
+            defined as part of the operator. In this case, target_wires should not be specified.
         estimation_wires (Union[Wires, Sequence[int], or int]): the wires to be used for phase
             estimation
 
     Raises:
         QuantumFunctionError: if the ``target_wires`` and ``estimation_wires`` share a common
-            element
+            element, or if ``target_wires`` are specified for an operator unitary.
 
     .. details::
         :title: Usage Details
@@ -104,11 +107,62 @@ class QuantumPhaseEstimation(Operation):
 
             # Need to rescale phase due to convention of RX gate
             phase_estimated = 4 * np.pi * (1 - phase_estimated)
+
+        We can also perform phase estimation on an operator. Note that since operators are defined
+        with target wires, the target wires should not be provided for the QPE.
+
+        .. code-block:: python
+
+
+            # use the product to specify compound operators
+            unitary = qml.RX(np.pi / 2, wires=[0]) @ qml.CNOT(wires=[0, 1])
+            eigenvector = np.array([-1/2, -1/2, 1/2, 1/2])
+
+            n_estimation_wires = 5
+            estimation_wires = range(2, n_estimation_wires + 2)
+            target_wires = [0, 1]
+
+            dev = qml.device("default.qubit", wires=n_estimation_wires + 2)
+
+            @qml.qnode(dev)
+            def circuit():
+                qml.QubitStateVector(eigenvector, wires=target_wires)
+                QuantumPhaseEstimation(
+                    unitary,
+                    estimation_wires=estimation_wires,
+                )
+                return qml.probs(estimation_wires)
+
+            phase_estimated = np.argmax(circuit()) / 2 ** n_estimation_wires
+
     """
     num_wires = AnyWires
     grad_method = None
 
-    def __init__(self, unitary, target_wires, estimation_wires, do_queue=True, id=None):
+    def __init__(self, unitary, target_wires=None, estimation_wires=None, do_queue=True, id=None):
+
+        if isinstance(unitary, Operator):
+
+            # If the unitary is expressed in terms of operators, do not provide target wires
+            if target_wires is not None:
+                raise qml.QuantumFunctionError(
+                    "The unitary is expressed as an operator, which already has target wires "
+                    "defined, do not additionally specify target wires."
+                )
+            target_wires = unitary.wires
+
+        elif target_wires is None:
+            raise qml.QuantumFunctionError(
+                "Target wires must be specified if the unitary is expressed as a matrix."
+            )
+
+        else:
+            unitary = qml.QubitUnitary(unitary, wires=target_wires)
+
+        # Estimation wires are required, but kept as an optional argument so that it can be
+        # placed after target_wires for backwards compatibility.
+        if estimation_wires is None:
+            raise qml.QuantumFunctionError("No estimation wires specified.")
 
         target_wires = list(target_wires)
         estimation_wires = list(estimation_wires)
@@ -116,57 +170,72 @@ class QuantumPhaseEstimation(Operation):
 
         if any(wire in target_wires for wire in estimation_wires):
             raise qml.QuantumFunctionError(
-                "The target wires and estimation wires must be different"
+                "The target wires and estimation wires must not overlap."
             )
 
         self._hyperparameters = {
+            "unitary": unitary,
             "target_wires": target_wires,
             "estimation_wires": estimation_wires,
         }
 
-        super().__init__(unitary, wires=wires, do_queue=do_queue, id=id)
+        super().__init__(wires=wires, do_queue=do_queue, id=id)
 
     @property
-    def num_params(self):
-        return 1
+    def target_wires(self):
+        """The target wires of the QPE"""
+        return self._hyperparameters["target_wires"]
+
+    @property
+    def estimation_wires(self):
+        """The estimation wires of the QPE"""
+        return self._hyperparameters["estimation_wires"]
+
+    # pylint: disable=protected-access
+    def map_wires(self, wire_map: dict):
+        new_op = super().map_wires(wire_map)
+        new_op._hyperparameters["unitary"] = qml.map_wires(
+            new_op._hyperparameters["unitary"], wire_map
+        )
+
+        new_op._hyperparameters["estimation_wires"] = [
+            wire_map.get(wire, wire) for wire in self.estimation_wires
+        ]
+        new_op._hyperparameters["target_wires"] = [
+            wire_map.get(wire, wire) for wire in self.target_wires
+        ]
+
+        return new_op
+
+    def queue(self, context=QueuingManager):
+        context.update_info(self._hyperparameters["unitary"], owner=self)
+        context.append(self, owns=self._hyperparameters["unitary"])
+        return self
 
     @staticmethod
     def compute_decomposition(
-        unitary, wires, target_wires, estimation_wires
+        wires, unitary, target_wires, estimation_wires
     ):  # pylint: disable=arguments-differ,unused-argument
-        r"""Representation of the operator as a product of other operators.
+        r"""Representation of the QPE circuit as a product of other operators.
 
         .. math:: O = O_1 O_2 \dots O_n.
 
 
-
         .. seealso:: :meth:`~.QuantumPhaseEstimation.decomposition`.
 
         Args:
-            unitary (array): the phase estimation unitary, specified as a matrix
-            wires (Any or Iterable[Any]): wires that the operator acts on
+            wires (Any or Iterable[Any]): wires that the QPE circuit acts on
+            unitary (Operator): the phase estimation unitary, specified as an operator
             target_wires (Any or Iterable[Any]): the target wires to apply the unitary
-            estimation_wires (Any or Iterable[Any]): the wires to be used for phase
-                estimation
+            estimation_wires (Any or Iterable[Any]): the wires to be used for phase estimation
 
         Returns:
             list[.Operator]: decomposition of the operator
         """
 
-        unitary_powers = [unitary]
-
-        for _ in range(len(estimation_wires) - 1):
-            new_power = unitary_powers[-1] @ unitary_powers[-1]
-            unitary_powers.append(new_power)
-
-        op_list = []
-
-        for wire in estimation_wires:
-            op_list.append(Hadamard(wire))
-            op_list.append(
-                ControlledQubitUnitary(unitary_powers.pop(), control_wires=wire, wires=target_wires)
-            )
-
+        op_list = [qml.Hadamard(w) for w in estimation_wires]
+        pow_ops = (qml.pow(unitary, 2**i) for i in range(len(estimation_wires) - 1, -1, -1))
+        op_list.extend(qml.ctrl(op, w) for op, w in zip(pow_ops, estimation_wires))
         op_list.append(qml.adjoint(qml.templates.QFT(wires=estimation_wires)))
 
         return op_list


Test patch:
diff --git a/tests/templates/test_subroutines/test_qpe.py b/tests/templates/test_subroutines/test_qpe.py
index e32e17ddb..addb415d7 100644
--- a/tests/templates/test_subroutines/test_qpe.py
+++ b/tests/templates/test_subroutines/test_qpe.py
@@ -28,19 +28,27 @@ class TestDecomposition:
         op = qml.QuantumPhaseEstimation(m, target_wires=[0], estimation_wires=[1, 2])
         qscript = op.expand()
 
+        unitary = qml.QubitUnitary(m, wires=[0])
         with qml.queuing.AnnotatedQueue() as q:
-            qml.Hadamard(1),
-            qml.ControlledQubitUnitary(m @ m, control_wires=[1], wires=[0]),
-            qml.Hadamard(2),
-            qml.ControlledQubitUnitary(m, control_wires=[2], wires=[0]),
+            qml.Hadamard(1)
+            qml.Hadamard(2)
+            qml.ctrl(qml.pow(unitary, 2), control=[1])
+            qml.ctrl(qml.pow(unitary, 1), control=[2])
             qml.adjoint(qml.QFT(wires=[1, 2]))
         qscript2 = qml.tape.QuantumScript.from_queue(q)
-
         assert len(qscript) == len(qscript2)
-        # qml.equal doesn't work for Adjoint op yet, so we stop before we get to it.
-        for op1, op2 in zip(qscript[:-1], qscript2[:-1]):
+        # qml.equal doesn't work for Adjoint or Pow op yet, so we stop before we get to it.
+        for op1, op2 in zip(qscript[:2], qscript2[:2]):
             assert qml.equal(op1, op2)
 
+        assert qml.equal(qscript[2].base.base, qscript2[2].base.base)
+        assert qscript[2].base.z, qscript2[2].base.z
+        assert qscript[2].control_wires == qscript2[2].control_wires
+
+        assert qml.equal(qscript[3].base.base, qscript2[3].base.base)
+        assert qscript[3].base.z == qscript2[3].base.z
+        assert qscript[3].control_wires == qscript2[3].control_wires
+
         assert isinstance(qscript[-1], qml.ops.op_math.Adjoint)
         assert qml.equal(qscript[-1].base, qml.QFT(wires=(1, 2)))
 
@@ -138,6 +146,133 @@ class TestDecomposition:
         # lower
         assert np.allclose(estimates[-1], phase, rtol=1e-2)
 
+    @pytest.mark.parametrize("param", np.linspace(0, 2 * np.pi, 4))
+    def test_phase_estimated_single_ops(self, param):
+        """Tests that the QPE works correctly for a single operator"""
+
+        unitary = qml.RX(param, wires=[0])
+
+        # Analytical eigenvectors and phase of the unitary
+        eig_vec = np.array([-1 / np.sqrt(2), 1 / np.sqrt(2)])
+        phase = param / (4 * np.pi)
+
+        estimates = []
+        wire_range = range(3, 11)
+
+        for wires in wire_range:
+            dev = qml.device("default.qubit", wires=wires)
+
+            estimation_wires = range(1, wires - 1)
+            target_wires = [0]
+
+            tape = qml.tape.QuantumScript(
+                [qml.QuantumPhaseEstimation(unitary, estimation_wires=estimation_wires)],
+                [qml.probs(estimation_wires)],
+                prep=[qml.QubitStateVector(eig_vec, wires=target_wires)],
+            )
+
+            tape = tape.expand(depth=2, stop_at=lambda obj: obj.name in dev.operations)
+            res = dev.execute(tape).flatten()
+
+            estimate = np.argmax(res) / 2 ** (wires - 2)
+            estimates.append(estimate)
+
+        # Check that the error is monotonically decreasing
+        for i in range(len(estimates) - 1):
+            err1 = np.abs(estimates[i] - phase)
+            err2 = np.abs(estimates[i + 1] - phase)
+            assert err1 >= err2
+
+        # This is a large error, but we'd need to push the qubit number up more to get it lower
+        assert np.allclose(estimates[-1], phase, rtol=1e-2)
+
+    @pytest.mark.parametrize("param", np.linspace(0, 2 * np.pi, 4))
+    def test_phase_estimated_ops(self, param):
+        """Tests that the QPE works correctly for compound operators"""
+
+        unitary = qml.RX(param, wires=[0]) @ qml.CNOT(wires=[0, 1])
+
+        # Analytical eigenvectors and phase of the unitary
+        eig_vec = np.array([-1 / 2, -1 / 2, 1 / 2, 1 / 2])
+        phase = param / (4 * np.pi)
+
+        estimates = []
+        wire_range = range(3, 11)
+
+        for wires in wire_range:
+
+            dev = qml.device("default.qubit", wires=wires)
+
+            # Offset the index of target wires to test the wire maÏp
+            estimation_wires = range(2, wires)
+            target_wires = [0, 1]
+
+            tape = qml.tape.QuantumScript(
+                [qml.QuantumPhaseEstimation(unitary, estimation_wires=estimation_wires)],
+                [qml.probs(estimation_wires)],
+                prep=[qml.QubitStateVector(eig_vec, wires=target_wires)],
+            )
+
+            tape = tape.expand(depth=2, stop_at=lambda obj: obj.name in dev.operations)
+            res = dev.execute(tape).flatten()
+
+            estimate = np.argmax(res) / 2 ** (wires - 2)
+            estimates.append(estimate)
+
+        # Check that the error is monotonically decreasing
+        for i in range(len(estimates) - 1):
+            err1 = np.abs(estimates[i] - phase)
+            err2 = np.abs(estimates[i + 1] - phase)
+            assert err1 >= err2
+
+        # This is a large error, but we'd need to push the qubit number up more to get it lower
+        assert np.allclose(estimates[-1], phase, rtol=1e-2)
+
+    def test_wires_specified(self):
+        """Tests errors with specifying target_wires and estimation_wires"""
+
+        unitary = unitary_group.rvs(4, random_state=1967)
+
+        with pytest.raises(
+            qml.QuantumFunctionError,
+            match="Target wires must be specified if the unitary is expressed as a matrix.",
+        ):
+            qml.QuantumPhaseEstimation(unitary, estimation_wires=[2, 3])
+
+        unitary = qml.RX(3, wires=[0])
+        with pytest.raises(
+            qml.QuantumFunctionError,
+            match="The unitary is expressed as an operator, which already has target wires "
+            "defined, do not additionally specify target wires.",
+        ):
+            qml.QuantumPhaseEstimation(unitary, target_wires=[1], estimation_wires=[2, 3])
+
+        with pytest.raises(
+            qml.QuantumFunctionError,
+            match="No estimation wires specified.",
+        ):
+            qml.QuantumPhaseEstimation(unitary)
+
+    def test_map_wires(self):
+        """Tests that QPE behaves correctly in a wire map"""
+
+        unitary = qml.RX(np.pi / 4, wires=[0]) @ qml.CNOT(wires=[0, 1])
+        qpe = qml.QuantumPhaseEstimation(unitary, estimation_wires=[2, 3])
+        new_qpe = qml.map_wires(
+            qpe,
+            wire_map={
+                0: 2,
+                1: 3,
+                2: 4,
+                3: 5,
+            },
+        )
+
+        assert list(new_qpe.wires) == [2, 3, 4, 5]
+        assert list(new_qpe._hyperparameters["target_wires"]) == [2, 3]
+        assert list(new_qpe._hyperparameters["estimation_wires"]) == [4, 5]
+        assert list(new_qpe._hyperparameters["unitary"].wires) == [2, 3]
+
     def test_adjoint(self):
         """Test that the QPE adjoint works."""
         dev = qml.device("default.qubit", wires=3)
@@ -174,11 +309,11 @@ class TestInputs:
         common element"""
 
         with pytest.raises(qml.QuantumFunctionError, match="The target wires and estimation wires"):
-            qml.QuantumPhaseEstimation(np.eye(2), target_wires=[0, 1], estimation_wires=[1, 2])
+            qml.QuantumPhaseEstimation(np.eye(4), target_wires=[0, 1], estimation_wires=[1, 2])
 
     def test_id(self):
         """Tests that the id attribute can be set."""
         template = qml.QuantumPhaseEstimation(
-            np.eye(2), target_wires=[0, 1], estimation_wires=[2, 3], id="a"
+            np.eye(4), target_wires=[0, 1], estimation_wires=[2, 3], id="a"
         )
         assert template.id == "a"