Repository: PennyLaneAI/pennylane
Pull request #1609 (2021-08-30 12:56:47)

Add a transform get_unitary_matrix to obtain matrix representations of circuits
**Context:**

Create a new function to construct the unitary matrix corresponding to any quantum circuit.

**Description of the Change:**

The new transform `get_unitary_matrix` works similarly to  `_apply_unitary()` in `devices/default_qubit.py.`
**Benefits:**

Can process all operations, including controlled gates on non-adjacent wires or wires in reverse order.

**Possible Drawbacks:**

**Related GitHub Issues:**
 #1462

codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1609?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#1609](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1609?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (e7fcc2a) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/90baa1d0cc509ce2e0936ff0ccd2e8303c35f457?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (90baa1d) will **decrease** coverage by `0.31%`.
> The diff coverage is `41.77%`.

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

```diff
@@            Coverage Diff             @@
##           master    #1609      +/-   ##
==========================================
- Coverage   99.13%   98.82%   -0.32%     
==========================================
  Files         196      197       +1     
  Lines       14162    14277     +115     
==========================================
+ Hits        14040    14109      +69     
- Misses        122      168      +46     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1609?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/transforms/get\_unitary\_matrix.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1609/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvZ2V0X3VuaXRhcnlfbWF0cml4LnB5) | `41.77% <41.77%> (ø)` | |
| [pennylane/qnode.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1609/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3Fub2RlLnB5) | `98.67% <0.00%> (+0.18%)` | :arrow_up: |

------

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

ingstra: Thanks @rmoyard! I have made adjustments now. What is left is I need to edit the changelog, and I suppose it also should be added [here](https://pennylane.readthedocs.io/en/stable/code/qml_transforms.html). Under which section?  (I see I also need to add more documentation to the function itself)
rmoyard: @ingstra Thank you for the great work, yes you can add it in the `CHANGELOG.md` file under new features 💯 The documentation should be added automatically due to `doc/code/qml_transforms.rst`. You can check it by running locally `make docs`.
ingstra: About documentation: I did not get it to generate automatically just running `make docs`.. But by looking and testing I managed after some time to get it generated after adding a file `pennylane.transforms.get_unitary_matrix.rst` to the directory `doc/code/api/` and adding `from .get_unitary_matrix import get_unitary_matrix` in `transforms/__init__.py`. To get it to show up I'm also fairly certain it needs to be added to the docstring in `transforms/__init__.py`?
josh146: Hey @ingstra --- to get it autogenerated, all you should need to do is:

- Make sure it is imported into the [`pennylane/transforms/__init__.py`](https://github.com/PennyLaneAI/pennylane/blob/master/pennylane/transforms/__init__.py) file, and
- Make sure it is listed in the `pennylane/transforms/__init__.py` docstring

Once those two changes are made, it should automatically show up on the page https://pennylane--1609.org.readthedocs.build/en/1609/code/qml_transforms.html, no need to modify anything in the `doc/` folder 🙂 

Feel free to have a look at the other transform functions to see the docstring conventions we use. A guide to writing docstrings can also be found on our [developer hub](https://pennylane.readthedocs.io/en/latest/development/guide/documentation.html#docstrings).
ingstra: I see! Thanks. I discovered it backwards. I started looking in the `doc/` folder and tried adding the .rst file, but then it complained that there was no such function in the transform module. After that, I noticed the `__init__.py` file. 
ingstra: Aww it does not play nice with tensorflow. Trying to get the matrix of the operator (here RZ) throws an error. Not sure how to deal with that. 

```
theta = tf.Variable([0.1])
op = qml.RZ(theta[0], wires=0)
print(op)
print(op.matrix)
```
gives
```
RZ(<tf.Tensor: shape=(), dtype=float32, numpy=0.1>, wires=[0])
Traceback (most recent call last):
  File "/home/ingrid/pennylane/ingstra/pennylane/pennylane/transforms/test_again.py", line 15, in <module>
    print(op.matrix)
  File "/home/ingrid/pennylane/ingstra/pennylane/pennylane/operation.py", line 754, in matrix
    op_matrix = self._matrix(*self.parameters)
  File "/home/ingrid/pennylane/ingstra/pennylane/pennylane/ops/qubit/parametric_ops.py", line 160, in _matrix
    p = cmath.exp(-0.5j * theta)
  File "/home/ingrid/.local/lib/python3.9/site-packages/tensorflow/python/ops/math_ops.py", line 1266, in r_binary_op_wrapper
    y, x = maybe_promote_tensors(y, x)
  File "/home/ingrid/.local/lib/python3.9/site-packages/tensorflow/python/ops/math_ops.py", line 1202, in maybe_promote_tensors
    ops.convert_to_tensor(tensor, dtype, name="x"))
  File "/home/ingrid/.local/lib/python3.9/site-packages/tensorflow/python/profiler/trace.py", line 163, in wrapped
    return func(*args, **kwargs)
  File "/home/ingrid/.local/lib/python3.9/site-packages/tensorflow/python/framework/ops.py", line 1566, in convert_to_tensor
    ret = conversion_func(value, dtype=dtype, name=name, as_ref=as_ref)
  File "/home/ingrid/.local/lib/python3.9/site-packages/tensorflow/python/framework/constant_op.py", line 339, in _constant_tensor_conversion_function
    return constant(v, dtype=dtype, name=name)
  File "/home/ingrid/.local/lib/python3.9/site-packages/tensorflow/python/framework/constant_op.py", line 264, in constant
    return _constant_impl(value, dtype, shape, name, verify_shape=False,
  File "/home/ingrid/.local/lib/python3.9/site-packages/tensorflow/python/framework/constant_op.py", line 276, in _constant_impl
    return _constant_eager_impl(ctx, value, dtype, shape, verify_shape)
  File "/home/ingrid/.local/lib/python3.9/site-packages/tensorflow/python/framework/constant_op.py", line 301, in _constant_eager_impl
    t = convert_to_eager_tensor(value, ctx, dtype)
  File "/home/ingrid/.local/lib/python3.9/site-packages/tensorflow/python/framework/constant_op.py", line 98, in convert_to_eager_tensor
    return ops.EagerTensor(value, ctx.device_name, dtype)
TypeError: Cannot convert (-0-0.5j) to EagerTensor of dtype float

```
glassnotes: > Aww it does not play nice with tensorflow. Trying to get the matrix of the operator (here RZ) throws an error. Not sure how to deal with that.

Ahh yeah, Tensorflow is a tricky one, to multiply a real number by a complex one it needs to be casted. I ran into this recently in a separate PR and put together a rough solution; you can see the implementation [here](
https://github.com/PennyLaneAI/pennylane/blob/46a4fbdd56c5566a70f1de51bbcf18f5c65cbbda/tests/transforms/test_optimization/utils.py), feel free to adapt it here. I know `RZ` and `Rot` need to be handled separately, you might find there are others needed as well. 

ingstra: Ok so now it works with tensorflow (the other interfaces gave no problems), but pylint doesn't think it's good (too many branches). What do you guys think? Please also comment on the default `wire_order` issue mentioned in an unresolved comment above.
rmoyard: Hi @ingstra, thank you for tackling those problems 👍 I will give a review and answer your questions as soon as possible!
ingstra: Ugh I notice pylint is also complaining about "anomalous backslash in string". I see in other documentation there is only one backslash before a greek letter, but for some reason that wouldn't render properly on my computer when I ran `make docs` so I had to put two. Can remove later
ingstra: Oh what is going on now... I merged upstream, and now `tests/transforms/test_decompositions.py` is failing? It doesn't even have anything to do with`get_unitary_matrix`?? The first pytest error is

```
 @pytest.mark.parametrize("U,expected_gate,expected_params", single_qubit_decomps)
    def test_zyz_decomposition(self, U, expected_gate, expected_params):
        """Test that a one-qubit matrix in isolation is correctly decomposed."""
        obtained_gates = zyz_decomposition(U, Wires("a"))
    
        assert len(obtained_gates) == 1
    
        assert isinstance(obtained_gates[0], expected_gate)
        assert obtained_gates[0].wires == Wires("a")
>       assert qml.math.allclose(obtained_gates[0].parameters, expected_params)

E       AssertionError: assert False
```

Thoroughly confused. Don't know what else to do so I wanted to look at `zyz_decomposition.py` but it doesn't even exist in my directory `pennylane/transforms/decompositions`/. But git saysI'm up to date with both the main repo and my fork. How can this happen...
ingstra: I'm very confused as to why `tests/transforms/test_decompositions.py` is failing 

```
def test_zyz_decomposition(self, U, expected_gate, expected_params):
        """Test that a one-qubit matrix in isolation is correctly decomposed."""
        obtained_gates = zyz_decomposition(U, Wires("a"))
    
        assert len(obtained_gates) == 1
    
        assert isinstance(obtained_gates[0], expected_gate)
        assert obtained_gates[0].wires == Wires("a")
        print("OARAMS!", obtained_gates[0].parameters,expected_params)
>       assert qml.math.allclose(obtained_gates[0].parameters, expected_params)
E       AssertionError: assert False
```
I printed the variables, and I get` obtained_gates[0].parameters= [0.0, -3.141592653589793, -3.141592653589793]` and `expected_params=[0.0, 3.141592653589793, 3.141592653589793]`. I fail to see how anything I've touched interacts with this function?
josh146: @ingstra, just yesterday, we merged in @glassnotes's PR #1552. This likely got pulled in when you merged master into your branch. I think the issue might be related to this!

Just tagging @glassnotes in case she has any ideas here
glassnotes: > @ingstra, just yesterday, we merged in @glassnotes's PR #1552. This likely got pulled in when you merged master into your branch. I think the issue might be related to this!
> 
> Just tagging @glassnotes in case she has any ideas here

Yes there was a bug fix in the `zyz_decomposition` that changed the expected parameters. However I also updated the tests, and this change should have propagated through to your branch. Let me check it out and then check it out :+1: 
glassnotes: Ahhh, I checked the error log, and I think I know the problem. The tests I merged in yesterday in `test_decompositions` and `test_unitary_to_rot` all still use the `compute_matrix_from_ops`... functions, which are no longer present in your branch, so there is an error in the test collection.  Changing them to use the transform should smooth things out :crossed_fingers: 
ingstra: @glassnotes I can't find it! `grep -r compute_matrix_from_ops *` in the base directory gives me nothing?
glassnotes: > @glassnotes I can't find it! `grep -r compute_matrix_from_ops *` in the base directory gives me nothing?

@ingstra if you don't mind me pushing to your branch, I'm happy to just quickly update these tests to use the new transform! 
ingstra: @glassnotes sure go ahead :) 
glassnotes: @ingstra everything has worked smoothly except the tensorflow cases, am investigating that now :weary: 
glassnotes: @ingstra No problem :) Tensorflow is tricky to deal with. The first fix I pushed modified the actual matrix construction of the operations, but then @josh146 realized we could make use of a new feature called `qml.tape.Unwrap`. This avoids having to manually check if the interface is tensorflow, and then the parameters are "unwrapped" into numpy arrays automatically when looping through the operations on the tape.
ingstra: Why is Sphinx suddenly failing? The error message is not very forthcoming:

[sphinx-action] Build failed with 0 warnings
Traceback (most recent call last):
  File "/entrypoint.py", line 22, in <module>
    action.build_all_docs(github_env, [os.environ.get("INPUT_DOCS-FOLDER")])
  File "/sphinx_action/action.py", line 167, in build_all_docs
    raise RuntimeError("Build failed")
RuntimeError: Build failed
ingstra: Thanks @glassnotes! I accepted your suggestions and added a test with JAX. 

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index d1a837fe6..50a7bd7af 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -4,6 +4,10 @@
 
 <h3>New features since last release</h3>
 
+* The unitary matrix corresponding to a quantum circuit can now be created using the new
+  `get_unitary_matrix()` transform.
+  [(#1609)](https://github.com/PennyLaneAI/pennylane/pull/1609)
+
 * Arbitrary two-qubit unitaries can now be decomposed into elementary gates. This
   functionality has been incorporated into the `qml.transforms.unitary_to_rot` transform, and is
   available separately as `qml.transforms.two_qubit_decomposition`.
@@ -76,9 +80,9 @@
   For more usage details, please see the
   [classical Jacobian docstring](https://pennylane.readthedocs.io/en/latest/code/api/pennylane.transforms.classical_jacobian.html).
 
-* Added a new operation `OrbitalRotation`, which implements the spin-adapted spatial orbital rotation gate. 
+* Added a new operation `OrbitalRotation`, which implements the spin-adapted spatial orbital rotation gate.
   [(#1665)](https://github.com/PennyLaneAI/pennylane/pull/1665)
-  
+
   An example circuit that uses `OrbitalRotation` operation is:
 
   ```python
@@ -92,7 +96,7 @@
 
   If we run this circuit, we will get the following output
 
-  ```pycon 
+  ```pycon
   >>> circuit(0.1)
   array([ 0.        +0.j,  0.        +0.j,  0.        +0.j,
           0.00249792+0.j,  0.        +0.j,  0.        +0.j,
@@ -224,4 +228,4 @@
 
 This release contains contributions from (in alphabetical order):
 
-Utkarsh Azad, Olivia Di Matteo, Andrew Gardhouse, Josh Izaac, Christina Lee, David Wierichs.
+Utkarsh Azad, Olivia Di Matteo, Andrew Gardhouse, Josh Izaac, Christina Lee, Ingrid Strandberg, David Wierichs.
diff --git a/pennylane/transforms/__init__.py b/pennylane/transforms/__init__.py
index 39138b52a..5a04a41d7 100644
--- a/pennylane/transforms/__init__.py
+++ b/pennylane/transforms/__init__.py
@@ -31,6 +31,7 @@ that compute the desired quantity.
 
     ~transforms.classical_jacobian
     ~draw
+    ~transforms.get_unitary_matrix
     ~metric_tensor
     ~specs
 
@@ -121,3 +122,4 @@ from .optimization import (
 from .specs import specs
 from .qmc import apply_controlled_Q, quantum_monte_carlo
 from .unitary_to_rot import unitary_to_rot
+from .get_unitary_matrix import get_unitary_matrix
diff --git a/pennylane/transforms/get_unitary_matrix.py b/pennylane/transforms/get_unitary_matrix.py
new file mode 100644
index 000000000..4acfef478
--- /dev/null
+++ b/pennylane/transforms/get_unitary_matrix.py
@@ -0,0 +1,133 @@
+# 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.
+"""
+A transform to obtain the matrix representation of a quantum circuit.
+"""
+from functools import wraps
+import numpy as np
+from pennylane.wires import Wires
+import pennylane as qml
+
+
+def get_unitary_matrix(circuit, wire_order=None):
+    r"""Construct the matrix representation of a quantum circuit.
+
+    Args:
+        circuit (pennylane.QNode, .QuantumTape, or Callable): A quantum node, tape,
+            or function that applies quantum operations.
+        wire_order (Sequence[Any], optional): Order of the wires in the quantum circuit.
+            Defaults to the order in which the wires appear in the quantum function.
+
+    Returns:
+         function: Function which accepts the same arguments as the QNode or quantum function.
+         When called, this function will return the unitary matrix as a numpy array.
+
+    **Example**
+
+    Consider the following function (the same applies for a QNode or tape):
+
+    .. code-block:: python3
+
+        def circuit(theta):
+            qml.RX(theta, wires=1)
+            qml.PauliZ(wires=0)
+
+
+    We can use ``get_unitary_matrix`` to generate a new function
+    that returns the unitary matrix corresponding to the function ``circuit``:
+
+
+    >>> get_matrix = get_unitary_matrix(circuit)
+    >>> theta = np.pi/4
+    >>> get_matrix(theta)
+    array([[ 0.92387953+0.j,  0.+0.j ,  0.-0.38268343j,  0.+0.j],
+       [ 0.+0.j,  -0.92387953+0.j,  0.+0.j,  0. +0.38268343j],
+       [ 0. -0.38268343j,  0.+0.j,  0.92387953+0.j,  0.+0.j],
+       [ 0.+0.j,  0.+0.38268343j,  0.+0.j,  -0.92387953+0.j]])
+
+
+    Note that since ``wire_order`` was not specified, the default order ``[1, 0]``
+    for ``circuit`` was used, and the unitary matrix corresponds to the operation
+    :math:`Z\otimes R_X(\theta)`. To obtain the matrix for :math:`R_X(\theta)\otimes Z`,
+    specify ``wire_order=[0, 1]`` in the function call:
+
+    >>> get_matrix = get_unitary_matrix(circuit, wire_order=[0, 1])
+
+    You can also get the unitary matrix for operations on a subspace of a
+    larger Hilbert space. For example, with the same function ``circuit`` and
+    ``wire_order=["a", 0, "b", 1]`` you obtain the :math:`16\times 16` matrix for
+    the operation :math:`I\otimes Z\otimes I\otimes  R_X(\theta)`.
+    """
+
+    wires = wire_order
+
+    @wraps(circuit)
+    def wrapper(*args, **kwargs):
+
+        if isinstance(circuit, qml.QNode):
+            # user passed a QNode, get the tape
+            circuit.construct(args, kwargs)
+            tape = circuit.qtape
+
+        elif isinstance(circuit, qml.tape.QuantumTape):
+            # user passed a tape
+            tape = circuit
+
+        elif callable(circuit):
+            # user passed something that is callable but not a tape or qnode.
+            tape = qml.transforms.make_tape(circuit)(*args, **kwargs)
+            # raise exception if it is not a quantum function
+            if len(tape.operations) == 0:
+                raise ValueError("Function contains no quantum operation")
+
+        else:
+            raise ValueError("Input is not a tape, QNode, or quantum function")
+
+        # if no wire ordering is specified, take wire list from tape
+        wire_order = tape.wires if wires is None else Wires(wires)
+
+        n_wires = len(wire_order)
+
+        # check that all wire labels in the circuit are contained in wire_order
+        if not set(tape.wires).issubset(wire_order):
+            raise ValueError("Wires in circuit are inconsistent with those in wire_order")
+
+        # initialize the unitary matrix
+        unitary_matrix = np.eye(2 ** n_wires)
+
+        with qml.tape.Unwrap(tape):
+            for op in tape.operations:
+                # operator wire position relative to wire ordering
+                op_wire_pos = wire_order.indices(op.wires)
+
+                I = np.reshape(np.eye(2 ** n_wires), [2] * n_wires * 2)
+                axes = (np.arange(len(op.wires), 2 * len(op.wires)), op_wire_pos)
+
+                # reshape op.matrix
+                U_op_reshaped = np.reshape(op.matrix, [2] * len(op.wires) * 2)
+                U_tensordot = np.tensordot(U_op_reshaped, I, axes=axes)
+
+                unused_idxs = [idx for idx in range(n_wires) if idx not in op_wire_pos]
+                # permute matrix axes to match wire ordering
+                perm = op_wire_pos + unused_idxs
+                U = np.moveaxis(U_tensordot, wire_order.indices(wire_order), perm)
+
+                U = np.reshape(U, ((2 ** n_wires, 2 ** n_wires)))
+
+                # add to total matrix if there are multiple ops
+                unitary_matrix = np.dot(U, unitary_matrix)
+
+        return unitary_matrix
+
+    return wrapper


Test patch:
diff --git a/tests/transforms/test_decompositions.py b/tests/transforms/test_decompositions.py
index 09b84371c..d871943f6 100644
--- a/tests/transforms/test_decompositions.py
+++ b/tests/transforms/test_decompositions.py
@@ -21,6 +21,7 @@ import pennylane as qml
 from pennylane import numpy as np
 
 from pennylane.wires import Wires
+from pennylane.transforms.get_unitary_matrix import get_unitary_matrix
 
 from pennylane.transforms.decompositions import zyz_decomposition
 from pennylane.transforms.decompositions import two_qubit_decomposition
@@ -30,7 +31,7 @@ from pennylane.transforms.decompositions.two_qubit_unitary import (
     _compute_num_cnots,
 )
 
-from test_optimization.utils import compute_matrix_from_ops_two_qubit, check_matrix_equivalence
+from test_optimization.utils import check_matrix_equivalence
 from gate_data import I, Z, S, T, H, X, CNOT, SWAP
 
 single_qubit_decomps = [
@@ -661,7 +662,7 @@ class TestTwoQubitUnitaryDecomposition:
     @pytest.mark.parametrize("wires", [[0, 1], ["a", "b"], [3, 2], ["c", 0]])
     @pytest.mark.parametrize("U", samples_3_cnots)
     def test_two_qubit_decomposition_3_cnots(self, U, wires):
-        """Test that a two-qubit matrix in isolation is correctly decomposed."""
+        """Test that a two-qubit matrix using 3 CNOTs is correctly decomposed."""
         U = _convert_to_su4(np.array(U))
 
         assert _compute_num_cnots(U) == 3
@@ -669,9 +670,11 @@ class TestTwoQubitUnitaryDecomposition:
         obtained_decomposition = two_qubit_decomposition(U, wires=wires)
         assert len(obtained_decomposition) == 10
 
-        obtained_matrix = compute_matrix_from_ops_two_qubit(
-            obtained_decomposition, wire_order=wires
-        )
+        with qml.tape.QuantumTape() as tape:
+            for op in obtained_decomposition:
+                qml.apply(op)
+
+        obtained_matrix = get_unitary_matrix(tape, wire_order=wires)()
 
         # We check with a slightly great tolerance threshold here simply because the
         # test matrices were copied in here with reduced precision.
@@ -679,8 +682,8 @@ class TestTwoQubitUnitaryDecomposition:
 
     @pytest.mark.parametrize("wires", [[0, 1], ["a", "b"], [3, 2], ["c", 0]])
     @pytest.mark.parametrize("U", samples_2_cnots)
-    def test_two_qubit_decomposition_1_cnot(self, U, wires):
-        """Test that a two-qubit matrix in isolation is correctly decomposed."""
+    def test_two_qubit_decomposition_2_cnots(self, U, wires):
+        """Test that a two-qubit matrix using 2 CNOTs isolation is correctly decomposed."""
         U = _convert_to_su4(np.array(U))
 
         assert _compute_num_cnots(U) == 2
@@ -688,16 +691,18 @@ class TestTwoQubitUnitaryDecomposition:
         obtained_decomposition = two_qubit_decomposition(U, wires=wires)
         assert len(obtained_decomposition) == 8
 
-        obtained_matrix = compute_matrix_from_ops_two_qubit(
-            obtained_decomposition, wire_order=wires
-        )
+        with qml.tape.QuantumTape() as tape:
+            for op in obtained_decomposition:
+                qml.apply(op)
+
+        obtained_matrix = get_unitary_matrix(tape, wire_order=wires)()
 
         assert check_matrix_equivalence(U, obtained_matrix, atol=1e-7)
 
     @pytest.mark.parametrize("wires", [[0, 1], ["a", "b"], [3, 2], ["c", 0]])
     @pytest.mark.parametrize("U", samples_1_cnot)
     def test_two_qubit_decomposition_1_cnot(self, U, wires):
-        """Test that a two-qubit matrix in isolation is correctly decomposed."""
+        """Test that a two-qubit matrix using one CNOT is correctly decomposed."""
         U = _convert_to_su4(np.array(U))
 
         assert _compute_num_cnots(U) == 1
@@ -705,16 +710,18 @@ class TestTwoQubitUnitaryDecomposition:
         obtained_decomposition = two_qubit_decomposition(U, wires=wires)
         assert len(obtained_decomposition) == 5
 
-        obtained_matrix = compute_matrix_from_ops_two_qubit(
-            obtained_decomposition, wire_order=wires
-        )
+        with qml.tape.QuantumTape() as tape:
+            for op in obtained_decomposition:
+                qml.apply(op)
+
+        obtained_matrix = get_unitary_matrix(tape, wire_order=wires)()
 
         assert check_matrix_equivalence(U, obtained_matrix, atol=1e-7)
 
     @pytest.mark.parametrize("wires", [[0, 1], ["a", "b"], [3, 2], ["c", 0]])
     @pytest.mark.parametrize("U_pair", samples_su2_su2)
     def test_two_qubit_decomposition_tensor_products(self, U_pair, wires):
-        """Test that a one-qubit matrix in isolation is correctly decomposed."""
+        """Test that a two-qubit tensor product matrix is correctly decomposed."""
         U = _convert_to_su4(qml.math.kron(np.array(U_pair[0]), np.array(U_pair[1])))
 
         assert _compute_num_cnots(U) == 0
@@ -722,7 +729,11 @@ class TestTwoQubitUnitaryDecomposition:
         obtained_decomposition = two_qubit_decomposition(U, wires=wires)
         assert len(obtained_decomposition) == 2
 
-        obtained_matrix = compute_matrix_from_ops_two_qubit(obtained_decomposition, wires)
+        with qml.tape.QuantumTape() as tape:
+            for op in obtained_decomposition:
+                qml.apply(op)
+
+        obtained_matrix = get_unitary_matrix(tape, wire_order=wires)()
 
         assert check_matrix_equivalence(U, obtained_matrix, atol=1e-7)
 
@@ -739,9 +750,12 @@ class TestTwoQubitUnitaryDecompositionInterfaces:
         U = torch.tensor(U, dtype=torch.complex128)
 
         obtained_decomposition = two_qubit_decomposition(U, wires=wires)
-        obtained_matrix = compute_matrix_from_ops_two_qubit(
-            obtained_decomposition, wire_order=wires
-        )
+
+        with qml.tape.QuantumTape() as tape:
+            for op in obtained_decomposition:
+                qml.apply(op)
+
+        obtained_matrix = get_unitary_matrix(tape, wire_order=wires)()
 
         assert check_matrix_equivalence(U, obtained_matrix, atol=1e-7)
 
@@ -757,7 +771,11 @@ class TestTwoQubitUnitaryDecompositionInterfaces:
 
         obtained_decomposition = two_qubit_decomposition(U, wires=wires)
 
-        obtained_matrix = compute_matrix_from_ops_two_qubit(obtained_decomposition, wires)
+        with qml.tape.QuantumTape() as tape:
+            for op in obtained_decomposition:
+                qml.apply(op)
+
+        obtained_matrix = get_unitary_matrix(tape, wire_order=wires)()
 
         assert check_matrix_equivalence(U, obtained_matrix, atol=1e-7)
 
@@ -771,9 +789,11 @@ class TestTwoQubitUnitaryDecompositionInterfaces:
 
         obtained_decomposition = two_qubit_decomposition(U, wires=wires)
 
-        obtained_matrix = compute_matrix_from_ops_two_qubit(
-            obtained_decomposition, wire_order=wires
-        )
+        with qml.tape.QuantumTape() as tape:
+            for op in obtained_decomposition:
+                qml.apply(op)
+
+        obtained_matrix = get_unitary_matrix(tape, wire_order=wires)()
 
         assert check_matrix_equivalence(U, obtained_matrix, atol=1e-7)
 
@@ -789,8 +809,11 @@ class TestTwoQubitUnitaryDecompositionInterfaces:
 
         obtained_decomposition = two_qubit_decomposition(U, wires=wires)
 
-        print(obtained_decomposition)
-        obtained_matrix = compute_matrix_from_ops_two_qubit(obtained_decomposition, wires)
+        with qml.tape.QuantumTape() as tape:
+            for op in obtained_decomposition:
+                qml.apply(op)
+
+        obtained_matrix = get_unitary_matrix(tape, wire_order=wires)()
 
         assert check_matrix_equivalence(U, obtained_matrix, atol=1e-7)
 
@@ -808,9 +831,12 @@ class TestTwoQubitUnitaryDecompositionInterfaces:
         U = jax.numpy.array(U, dtype=jax.numpy.complex128)
 
         obtained_decomposition = two_qubit_decomposition(U, wires=wires)
-        obtained_matrix = compute_matrix_from_ops_two_qubit(
-            obtained_decomposition, wire_order=wires
-        )
+
+        with qml.tape.QuantumTape() as tape:
+            for op in obtained_decomposition:
+                qml.apply(op)
+
+        obtained_matrix = get_unitary_matrix(tape, wire_order=wires)()
 
         assert check_matrix_equivalence(U, obtained_matrix, atol=1e-7)
 
@@ -831,6 +857,10 @@ class TestTwoQubitUnitaryDecompositionInterfaces:
 
         obtained_decomposition = two_qubit_decomposition(U, wires=wires)
 
-        obtained_matrix = compute_matrix_from_ops_two_qubit(obtained_decomposition, wires)
+        with qml.tape.QuantumTape() as tape:
+            for op in obtained_decomposition:
+                qml.apply(op)
+
+        obtained_matrix = get_unitary_matrix(tape, wire_order=wires)()
 
         assert check_matrix_equivalence(U, obtained_matrix, atol=1e-7)
diff --git a/tests/transforms/test_get_unitary_matrix.py b/tests/transforms/test_get_unitary_matrix.py
new file mode 100644
index 000000000..c961c3c2c
--- /dev/null
+++ b/tests/transforms/test_get_unitary_matrix.py
@@ -0,0 +1,467 @@
+# 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.
+"""
+Unit tests for the get_unitary_matrix transform
+"""
+from functools import reduce
+import pytest
+from pennylane import numpy as np
+from gate_data import I, X, Y, Z, H, S, CNOT
+import pennylane as qml
+
+from pennylane.transforms.get_unitary_matrix import get_unitary_matrix
+
+# test non-parametric single qubit gates
+nonparam_1qubit_op_list = [qml.PauliX, qml.PauliY, qml.PauliZ, qml.Hadamard, qml.S, qml.T, qml.SX]
+
+
+@pytest.mark.parametrize("op", nonparam_1qubit_op_list)
+@pytest.mark.parametrize("wire", [0, 1, 2])
+def test_get_unitary_matrix_nonparam_1qubit_ops(op, wire):
+    """Check the matrices for different nonparametrized single-qubit gates, which are acting on different qubits in a space of three qubits."""
+    wires = [0, 1, 2]
+
+    def testcircuit(wire):
+        op(wires=wire)
+
+    get_matrix = get_unitary_matrix(testcircuit, wires)
+    matrix = get_matrix(wire)
+
+    if wire == 0:
+        expected_matrix = np.kron(op(wires=wire).matrix, np.eye(4))
+    if wire == 1:
+        expected_matrix = np.kron(np.eye(2), np.kron(op(wires=wire).matrix, np.eye(2)))
+    if wire == 2:
+        expected_matrix = np.kron(np.eye(4), op(wires=wire).matrix)
+
+    assert np.allclose(matrix, expected_matrix)
+
+
+# Test a circuit containing multiple gates
+def test_get_unitary_matrix_multiple_ops():
+    """Check the total matrix for a circuit containing multiple gates. Also
+    checks that non-integer wires work"""
+    wires = ["a", "b", "c"]
+
+    def testcircuit():
+        qml.PauliX(wires="a")
+        qml.S(wires="b")
+        qml.Hadamard(wires="c")
+        qml.CNOT(wires=["b", "c"])
+
+    get_matrix = get_unitary_matrix(testcircuit, wires)
+    matrix = get_matrix()
+
+    expected_matrix = np.kron(I, CNOT) @ np.kron(X, np.kron(S, H))
+
+    assert np.allclose(matrix, expected_matrix)
+
+
+@pytest.mark.parametrize("target_wire", [0, 2, 3, 4])
+def test_get_unitary_matrix_CNOT(target_wire):
+    """Test CNOT: 2-qubit gate with different target wires, some non-adjacent."""
+    wires = [0, 1, 2, 3, 4]
+
+    def testcircuit():
+        qml.CNOT(wires=[1, target_wire])
+
+    get_matrix = get_unitary_matrix(testcircuit, wires)
+    matrix = get_matrix()
+
+    # test the matrix operation on a state
+    state0 = [1, 0]
+    state1 = [0, 1]
+    teststate = reduce(np.kron, [state1, state1, state1, state1, state1])
+
+    if target_wire == 0:
+        expected_state = reduce(np.kron, [state0, state1, state1, state1, state1])
+    elif target_wire == 2:
+        expected_state = reduce(np.kron, [state1, state1, state0, state1, state1])
+    elif target_wire == 3:
+        expected_state = reduce(np.kron, [state1, state1, state1, state0, state1])
+    elif target_wire == 4:
+        expected_state = reduce(np.kron, [state1, state1, state1, state1, state0])
+
+    obtained_state = matrix @ teststate
+
+    assert np.allclose(obtained_state, expected_state)
+
+
+def test_get_unitary_matrix_CRX():
+    """Test controlled rotation with non-adjacent control and target wires"""
+    testangle = np.pi / 4
+
+    wires = [0, 1, 2]
+
+    def testcircuit():
+        qml.CRX(testangle, wires=[2, 0])
+
+    # test applying to state
+    state0 = [1, 0]
+    state1 = [0, 1]
+
+    # perform controlled rotation
+    teststate1 = reduce(np.kron, [state1, state1, state1])
+    # do not perform controlled rotation
+    teststate0 = reduce(np.kron, [state1, state1, state0])
+
+    expected_state1 = reduce(np.kron, [qml.RX(testangle, wires=1).matrix @ state1, state1, state1])
+    expected_state0 = teststate0
+
+    get_matrix = get_unitary_matrix(testcircuit, wires)
+    matrix = get_matrix()
+
+    obtained_state1 = matrix @ teststate1
+    obtained_state0 = matrix @ teststate0
+
+    assert np.allclose(obtained_state1, expected_state1)
+    assert np.allclose(obtained_state0, expected_state0)
+
+
+def test_get_unitary_matrix_Toffoli():
+    """Check the Toffoli matrix by its action on states"""
+    wires = [0, "a", 2, "c", 4]
+
+    def testcircuit():
+        qml.Toffoli(wires=[0, 4, "a"])
+
+    # test applying to state
+    state0 = [1, 0]
+    state1 = [0, 1]
+
+    teststate1 = reduce(np.kron, [state1, state1, state1, state1, state1])
+    teststate2 = reduce(np.kron, [state0, state0, state1, state1, state0])
+
+    expected_state1 = reduce(np.kron, [state1, state0, state1, state1, state1])
+    expected_state2 = teststate2
+
+    get_matrix = get_unitary_matrix(testcircuit, wires)
+    matrix = get_matrix()
+
+    obtained_state1 = matrix @ teststate1
+    obtained_state2 = matrix @ teststate2
+
+    assert np.allclose(obtained_state1, expected_state1)
+    assert np.allclose(obtained_state2, expected_state2)
+
+
+def test_get_unitary_matrix_MultiControlledX():
+    """Test with many control wires"""
+    wires = [0, 1, 2, 3, 4, 5]
+
+    def testcircuit():
+        qml.MultiControlledX(control_wires=[0, 2, 4, 5], wires=3)
+
+    state0 = [1, 0]
+    state1 = [0, 1]
+
+    teststate1 = reduce(np.kron, [state1, state1, state1, state1, state1, state1])
+    teststate2 = reduce(np.kron, [state0, state1, state0, state0, state1, state0])
+
+    expected_state1 = reduce(np.kron, [state1, state1, state1, state0, state1, state1])
+    expected_state2 = teststate2
+
+    get_matrix = get_unitary_matrix(testcircuit, wires)
+    matrix = get_matrix()
+
+    obtained_state1 = matrix @ teststate1
+    obtained_state2 = matrix @ teststate2
+
+    assert np.allclose(obtained_state1, expected_state1)
+    assert np.allclose(obtained_state2, expected_state2)
+
+
+def test_get_unitary_matrix_default_wireorder():
+    """Test without specified wire order"""
+
+    def testcircuit():
+        qml.PauliX(wires=0)
+        qml.PauliY(wires=1)
+        qml.PauliZ(wires=2)
+
+    get_matrix = get_unitary_matrix(testcircuit)
+
+    matrix = get_matrix()
+    expected_matrix = np.kron(X, np.kron(Y, Z))
+
+    assert np.allclose(matrix, expected_matrix)
+
+
+def test_get_unitary_matrix_input_tape():
+    """Test with quantum tape as input"""
+    with qml.tape.QuantumTape() as tape:
+        qml.RX(0.432, wires=0)
+        qml.RY(0.543, wires=0)
+        qml.CNOT(wires=[0, 1])
+        qml.RX(0.133, wires=1)
+
+    get_matrix = get_unitary_matrix(tape)
+
+    matrix = get_matrix()
+
+    part_expected_matrix = np.kron(qml.RY(0.543, wires=0).matrix @ qml.RX(0.432, wires=0).matrix, I)
+
+    expected_matrix = np.kron(I, qml.RX(0.133, wires=1).matrix) @ CNOT @ part_expected_matrix
+
+    assert np.allclose(matrix, expected_matrix)
+
+
+def test_get_unitary_matrix_input_tape_wireorder():
+    """Test with quantum tape as input, and nonstandard wire ordering"""
+    with qml.tape.QuantumTape() as tape:
+        qml.RX(0.432, wires=0)
+        qml.RY(0.543, wires=0)
+        qml.CNOT(wires=[0, 1])
+        qml.RX(0.133, wires=1)
+
+    get_matrix = get_unitary_matrix(tape, wire_order=[1, 0])
+    matrix = get_matrix()
+
+    # CNOT where the second wire is the control wire, as opposed to qml.CNOT.matrix
+    CNOT10 = np.array([[1, 0, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0], [0, 1, 0, 0]])
+
+    part_expected_matrix = np.kron(I, qml.RY(0.543, wires=0).matrix @ qml.RX(0.432, wires=0).matrix)
+
+    expected_matrix = np.kron(qml.RX(0.133, wires=1).matrix, I) @ CNOT10 @ part_expected_matrix
+
+    assert np.allclose(matrix, expected_matrix)
+
+
+def test_get_unitary_matrix_input_QNode():
+    """Test with QNode as input"""
+    dev = qml.device("default.qubit", wires=5)
+
+    @qml.qnode(dev)
+    def my_quantum_function():
+        qml.PauliZ(wires=0)
+        qml.CNOT(wires=[0, 1])
+        qml.PauliY(wires=1)
+        qml.CRZ(0.2, wires=[2, 3])
+        qml.PauliX(wires=4)
+        return qml.expval(qml.PauliZ(1))
+
+    get_matrix = get_unitary_matrix(my_quantum_function)  # default wire_order = [0, 1, 2, 3, 4]
+    matrix = get_matrix()
+
+    expected_matrix = (
+        reduce(np.kron, [I, I, I, I, X])
+        @ reduce(np.kron, [I, I, qml.CRZ(0.2, wires=[2, 3]).matrix, I])
+        @ reduce(np.kron, [I, Y, I, I, I])
+        @ reduce(np.kron, [CNOT, I, I, I])
+        @ reduce(np.kron, [Z, I, I, I, I])
+    )
+
+    assert np.allclose(matrix, expected_matrix)
+
+
+def test_get_unitary_matrix_input_QNode_wireorder():
+    """Test with QNode as input, and nonstandard wire ordering"""
+    dev = qml.device("default.qubit", wires=5)
+
+    @qml.qnode(dev)
+    def my_quantum_function():
+        qml.PauliZ(wires=0)
+        qml.CNOT(wires=[0, 1])
+        qml.PauliY(wires=1)
+        qml.CRZ(0.2, wires=[2, 3])
+        qml.PauliX(wires=4)
+        return qml.expval(qml.PauliZ(1))
+
+    get_matrix = get_unitary_matrix(my_quantum_function, wire_order=[1, 0, 4, 2, 3])
+    matrix = get_matrix()
+
+    CNOT10 = np.array([[1, 0, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0], [0, 1, 0, 0]])
+
+    expected_matrix = (
+        reduce(np.kron, [I, I, X, I, I])
+        @ reduce(np.kron, [I, I, I, qml.CRZ(0.2, wires=[2, 3]).matrix])
+        @ reduce(np.kron, [Y, I, I, I, I])
+        @ reduce(np.kron, [CNOT10, I, I, I])
+        @ reduce(np.kron, [I, Z, I, I, I])
+    )
+
+    assert np.allclose(matrix, expected_matrix)
+
+
+def test_get_unitary_matrix_invalid_argument():
+    """Assert error raised when input is neither a tape, QNode, nor quantum function"""
+
+    get_matrix = get_unitary_matrix(qml.PauliZ(0))
+
+    with pytest.raises(ValueError, match="Input is not a tape, QNode, or quantum function"):
+        matrix = get_matrix()
+
+
+def test_get_unitary_matrix_wrong_function():
+    """Assert error raised when input function is not a quantum function"""
+
+    def testfunction(x):
+        return x
+
+    get_matrix = get_unitary_matrix(testfunction, [0])
+
+    with pytest.raises(ValueError, match="Function contains no quantum operation"):
+        matrix = get_matrix(1)
+
+
+def test_get_unitary_matrix_interface_tf():
+    """Test with tensorflow interface"""
+
+    tf = pytest.importorskip("tensorflow")
+
+    dev = qml.device("default.qubit", wires=3)
+
+    def circuit(beta, theta):
+        qml.RZ(beta, wires=0)
+        qml.RZ(theta[0], wires=1)
+        qml.CRY(theta[1], wires=[1, 2])
+        return qml.expval(qml.PauliZ(1))
+
+    # set qnode interface
+    qnode_tensorflow = qml.QNode(circuit, dev, interface="tf")
+
+    get_matrix = get_unitary_matrix(qnode_tensorflow)
+
+    beta = 0.1
+    # input tensorflow parameters
+    theta = tf.Variable([0.2, 0.3])
+
+    matrix = get_matrix(beta, theta)
+
+    # expected matrix
+    theta_np = theta.numpy()
+    matrix1 = np.kron(qml.RZ(beta, wires=0).matrix, np.kron(qml.RZ(theta_np[0], wires=1).matrix, I))
+    matrix2 = np.kron(I, qml.CRY(theta_np[1], wires=[1, 2]).matrix)
+    expected_matrix = matrix2 @ matrix1
+
+    assert np.allclose(matrix, expected_matrix)
+
+
+def test_get_unitary_matrix_interface_torch():
+    """Test with torch interface"""
+
+    torch = pytest.importorskip("torch", minversion="1.8")
+
+    dev = qml.device("default.qubit", wires=3)
+
+    def circuit(theta):
+        qml.RZ(theta[0], wires=0)
+        qml.RZ(theta[1], wires=1)
+        qml.CRY(theta[2], wires=[1, 2])
+        return qml.expval(qml.PauliZ(1))
+
+    # set qnode interface
+    qnode_torch = qml.QNode(circuit, dev, interface="torch")
+
+    get_matrix = get_unitary_matrix(qnode_torch)
+
+    # input torch parameters
+    theta = torch.tensor([0.1, 0.2, 0.3])
+
+    matrix = get_matrix(theta)
+
+    # expected matrix
+    matrix1 = np.kron(
+        qml.RZ(theta[0], wires=0).matrix, np.kron(qml.RZ(theta[1], wires=1).matrix, I)
+    )
+    matrix2 = np.kron(I, qml.CRY(theta[2], wires=[1, 2]).matrix)
+    expected_matrix = matrix2 @ matrix1
+
+    assert np.allclose(matrix, expected_matrix)
+
+
+def test_get_unitary_matrix_interface_autograd():
+    """Test with autograd interface"""
+
+    dev = qml.device("default.qubit", wires=3)
+
+    def circuit(theta):
+        qml.RZ(theta[0], wires=0)
+        qml.RZ(theta[1], wires=1)
+        qml.CRY(theta[2], wires=[1, 2])
+        return qml.expval(qml.PauliZ(1))
+
+    # set qnode interface
+    qnode = qml.QNode(circuit, dev, interface="autograd")
+
+    get_matrix = get_unitary_matrix(qnode)
+
+    # set input parameters
+    theta = np.array([0.1, 0.2, 0.3], requires_grad=True)
+
+    matrix = get_matrix(theta)
+
+    # expected matrix
+    matrix1 = np.kron(
+        qml.RZ(theta[0], wires=0).matrix, np.kron(qml.RZ(theta[1], wires=1).matrix, I)
+    )
+    matrix2 = np.kron(I, qml.CRY(theta[2], wires=[1, 2]).matrix)
+    expected_matrix = matrix2 @ matrix1
+
+    assert np.allclose(matrix, expected_matrix)
+
+
+def test_get_unitary_matrix_interface_jax():
+    """Test with JAX interface"""
+
+    jax = pytest.importorskip("jax")
+    from jax import numpy as jnp
+    from jax.config import config
+
+    remember = config.read("jax_enable_x64")
+    config.update("jax_enable_x64", True)
+
+    dev = qml.device("default.qubit", wires=3)
+
+    def circuit(theta):
+        qml.RZ(theta[0], wires=0)
+        qml.RZ(theta[1], wires=1)
+        qml.CRY(theta[2], wires=[1, 2])
+        return qml.expval(qml.PauliZ(1))
+
+    # set qnode interface
+    qnode = qml.QNode(circuit, dev, interface="jax")
+
+    get_matrix = get_unitary_matrix(qnode)
+
+    # input jax parameters
+    theta = jnp.array([0.1, 0.2, 0.3], dtype=jnp.float64)
+
+    matrix = get_matrix(theta)
+
+    # expected matrix
+    matrix1 = np.kron(
+        qml.RZ(theta[0], wires=0).matrix, np.kron(qml.RZ(theta[1], wires=1).matrix, I)
+    )
+    matrix2 = np.kron(I, qml.CRY(theta[2], wires=[1, 2]).matrix)
+    expected_matrix = matrix2 @ matrix1
+
+    assert np.allclose(matrix, expected_matrix)
+
+
+def test_get_unitary_matrix_wronglabel():
+    """Assert error raised when wire labels in wire_order and circuit are inconsistent"""
+
+    def circuit():
+        qml.PauliX(wires=1)
+        qml.PauliZ(wires=0)
+
+    wires = [0, "b"]
+
+    get_matrix = get_unitary_matrix(circuit, wires)
+
+    with pytest.raises(
+        ValueError, match="Wires in circuit are inconsistent with those in wire_order"
+    ):
+        matrix = get_matrix()
diff --git a/tests/transforms/test_optimization/test_commute_controlled.py b/tests/transforms/test_optimization/test_commute_controlled.py
index f02a2306c..0c7b2eb94 100644
--- a/tests/transforms/test_optimization/test_commute_controlled.py
+++ b/tests/transforms/test_optimization/test_commute_controlled.py
@@ -18,9 +18,9 @@ import pennylane as qml
 from pennylane.wires import Wires
 
 from pennylane.transforms.optimization import commute_controlled
+from pennylane.transforms.get_unitary_matrix import get_unitary_matrix
 from utils import (
     compare_operation_lists,
-    compute_matrix_from_ops_two_qubit,
     check_matrix_equivalence,
 )
 
@@ -296,8 +296,12 @@ class TestCommuteControlled:
         assert len(original_ops) == len(transformed_ops)
 
         # Compare matrices
-        matrix_expected = compute_matrix_from_ops_two_qubit(original_ops, wire_order=[0, 1])
-        matrix_obtained = compute_matrix_from_ops_two_qubit(transformed_ops, wire_order=[0, 1])
+        compute_matrix = get_unitary_matrix(qfunc, [0, 1])
+        matrix_expected = compute_matrix()
+
+        compute_transformed_matrix = get_unitary_matrix(transformed_qfunc, [0, 1])
+        matrix_obtained = compute_transformed_matrix()
+
         assert check_matrix_equivalence(matrix_expected, matrix_obtained)
 
 
diff --git a/tests/transforms/test_optimization/test_optimization_utils.py b/tests/transforms/test_optimization/test_optimization_utils.py
index 49d04be68..bf4ff6416 100644
--- a/tests/transforms/test_optimization/test_optimization_utils.py
+++ b/tests/transforms/test_optimization/test_optimization_utils.py
@@ -21,7 +21,8 @@ from pennylane.transforms.optimization.optimization_utils import (
     fuse_rot_angles,
 )
 
-from utils import check_matrix_equivalence, compute_matrix_from_ops_one_qubit
+from utils import check_matrix_equivalence
+from pennylane.transforms.get_unitary_matrix import get_unitary_matrix
 
 
 sample_op_list = [
@@ -60,16 +61,23 @@ class TestRotGateFusion:
         """Test that a set of rotations of the form YZY is correctly converted
         to a sequence of the form ZYZ."""
 
-        original_ops = [
+        def original_ops():
             qml.RY(angles[0], wires=0),
             qml.RZ(angles[1], wires=0),
             qml.RY(angles[2], wires=0),
-        ]
-        product_yzy = compute_matrix_from_ops_one_qubit(original_ops)
+
+        compute_matrix = get_unitary_matrix(original_ops, [0])
+        product_yzy = compute_matrix()
 
         z1, y, z2 = _yzy_to_zyz(angles)
-        transformed_ops = [qml.RZ(z1, wires=0), qml.RY(y, wires=0), qml.RZ(z2, wires=0)]
-        product_zyz = compute_matrix_from_ops_one_qubit(transformed_ops)
+
+        def transformed_ops():
+            qml.RZ(z1, wires=0)
+            qml.RY(y, wires=0)
+            qml.RZ(z2, wires=0)
+
+        compute_transformed_matrix = get_unitary_matrix(transformed_ops, [0])
+        product_zyz = compute_transformed_matrix()
 
         assert check_matrix_equivalence(product_yzy, product_zyz)
 
@@ -95,8 +103,12 @@ class TestRotGateFusion:
         """Test that the fusion of two Rot gates has the same effect as
         applying the Rots sequentially."""
 
-        original_ops = [qml.Rot(*angles_1, wires=0), qml.Rot(*angles_2, wires=0)]
-        matrix_expected = compute_matrix_from_ops_one_qubit(original_ops)
+        def original_ops():
+            qml.Rot(*angles_1, wires=0)
+            qml.Rot(*angles_2, wires=0)
+
+        compute_matrix = get_unitary_matrix(original_ops, [0])
+        matrix_expected = compute_matrix()
 
         fused_angles = fuse_rot_angles(angles_1, angles_2)
         matrix_obtained = qml.Rot(*fused_angles, wires=0).matrix
diff --git a/tests/transforms/test_optimization/test_single_qubit_fusion.py b/tests/transforms/test_optimization/test_single_qubit_fusion.py
index 5a718d31b..868871412 100644
--- a/tests/transforms/test_optimization/test_single_qubit_fusion.py
+++ b/tests/transforms/test_optimization/test_single_qubit_fusion.py
@@ -18,6 +18,7 @@ from pennylane import numpy as np
 import pennylane as qml
 from pennylane.wires import Wires
 from pennylane.transforms.optimization import single_qubit_fusion
+from pennylane.transforms.get_unitary_matrix import get_unitary_matrix
 
 from utils import *
 
@@ -39,14 +40,12 @@ class TestSingleQubitFusion:
 
         transformed_qfunc = single_qubit_fusion()(qfunc)
 
-        original_ops = qml.transforms.make_tape(qfunc)().operations
-        transformed_ops = qml.transforms.make_tape(transformed_qfunc)().operations
-
-        assert len(transformed_ops) == 1
-
         # Compare matrices
-        matrix_expected = compute_matrix_from_ops_one_qubit(original_ops)
-        matrix_obtained = compute_matrix_from_ops_one_qubit(transformed_ops)
+        compute_matrix = get_unitary_matrix(qfunc, [0])
+        matrix_expected = compute_matrix()
+
+        compute_transformed_matrix = get_unitary_matrix(transformed_qfunc, [0])
+        matrix_obtained = compute_transformed_matrix()
         assert check_matrix_equivalence(matrix_expected, matrix_obtained)
 
     def test_single_qubit_fusion_no_gates_after(self):
@@ -127,8 +126,11 @@ class TestSingleQubitFusion:
         compare_operation_lists(transformed_ops, names_expected, wires_expected)
 
         # Compare matrices
-        matrix_expected = compute_matrix_from_ops_two_qubit(original_ops, wire_order=[0, 1])
-        matrix_obtained = compute_matrix_from_ops_two_qubit(transformed_ops, wire_order=[0, 1])
+        compute_matrix = get_unitary_matrix(qfunc, [0, 1])
+        matrix_expected = compute_matrix()
+
+        compute_transformed_matrix = get_unitary_matrix(transformed_qfunc, [0, 1])
+        matrix_obtained = compute_transformed_matrix()
         assert check_matrix_equivalence(matrix_expected, matrix_obtained)
 
     def test_single_qubit_fusion_multiple_qubits(self):
@@ -154,8 +156,12 @@ class TestSingleQubitFusion:
         compare_operation_lists(transformed_ops, names_expected, wires_expected)
 
         # Check matrix representation
-        matrix_expected = compute_matrix_from_ops_two_qubit(original_ops, ["a", "b"])
-        matrix_obtained = compute_matrix_from_ops_two_qubit(transformed_ops, ["a", "b"])
+        compute_matrix = get_unitary_matrix(qfunc, ["a", "b"])
+        matrix_expected = compute_matrix()
+
+        compute_transformed_matrix = get_unitary_matrix(transformed_qfunc, ["a", "b"])
+        matrix_obtained = compute_transformed_matrix()
+
         assert check_matrix_equivalence(matrix_expected, matrix_obtained)
 
 
diff --git a/tests/transforms/test_optimization/utils.py b/tests/transforms/test_optimization/utils.py
index 19e1c9f03..101aa7f04 100644
--- a/tests/transforms/test_optimization/utils.py
+++ b/tests/transforms/test_optimization/utils.py
@@ -18,51 +18,6 @@ from pennylane import numpy as np
 from gate_data import I, SWAP
 
 
-def compute_matrix_from_ops_one_qubit(ops):
-    """Given a list of single-qubit operations, construct its matrix representation."""
-
-    mat = I
-
-    for op in ops:
-        mat = qml.math.dot(op.matrix, mat)
-    return mat
-
-
-def compute_matrix_from_ops_two_qubit(ops, wire_order):
-    """Given a list of two-qubit operations, construct its matrix representation."""
-
-    mat = np.eye(4)
-
-    wire_order = qml.wires.Wires(wire_order)
-
-    for op in ops:
-        op_wires = qml.wires.Wires(op.wires)
-
-        if len(op_wires) == 1:
-            # These first two cases are to cover tensorflow quirks
-            if op.name == "RZ":
-                op_mat = qml.RZ(*qml.math.unwrap(op.parameters), wires=0).matrix
-            elif op.name == "Rot":
-                op_mat = qml.Rot(*qml.math.unwrap(op.parameters), wires=0).matrix
-            else:
-                op_mat = qml.math.unwrap(op.matrix)
-
-            if op_wires[0] == wire_order[0]:
-                tensor_prod = np.kron(op_mat, I)
-            else:
-                tensor_prod = np.kron(I, op_mat)
-
-            mat = np.dot(tensor_prod, mat)
-
-        else:
-            if op_wires == wire_order:
-                mat = np.dot(op.matrix, mat)
-            else:
-                mat = np.linalg.multi_dot([SWAP, op.matrix, SWAP, mat])
-
-    return mat
-
-
 def check_matrix_equivalence(matrix_expected, matrix_obtained, atol=1e-8):
     """Takes two matrices and checks if multiplying one by the conjugate
     transpose of the other gives the identity."""