Repository: PennyLaneAI/pennylane
Pull request #2241 (2022-02-28 08:42:09)

Add `qml.matrix(op)` function
**Context:**

Currently, to get the matrix of an operator, you must do `op.get_matrix()`. Alternatively, to get the matrix of a tape/qfunc, you must do `qml.transforms.get_unitary_matrix()`.

However, there are many downsides to this:

- `op.get_matrix()` is temporary, and we want to limit users relying on operator methods!

- `qml.transforms.get_unitary_matrix()` is hard to find, and also will fail if not all operations on the tape support `get_matrix()`.

It would be nice if there was a _single function_ that covered this singular functionality, while insulating the `op.get_matrix()` method from changes.

**Description of the Change:**

- Adds `qml.matrix()`

- This function works on operations, qfuncs, tapes, and QNodes.

- It accepts wire ordering.

**Benefits:** As above.

**Possible Drawbacks:** n/a

**Related GitHub Issues:** n/a


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

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

```diff
@@               Coverage Diff                @@
##           op-transform    #2241      +/-   ##
================================================
+ Coverage         96.15%   96.51%   +0.35%     
================================================
  Files               232      234       +2     
  Lines             18444    18465      +21     
================================================
+ Hits              17735    17821      +86     
+ Misses              709      644      -65     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2241?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/ops/functions/matrix.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2241/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9mdW5jdGlvbnMvbWF0cml4LnB5) | `94.44% <94.44%> (ø)` | |
| [pennylane/\_\_init\_\_.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2241/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL19faW5pdF9fLnB5) | `100.00% <100.00%> (ø)` | |
| [pennylane/ops/functions/\_\_init\_\_.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2241/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9mdW5jdGlvbnMvX19pbml0X18ucHk=) | `100.00% <100.00%> (ø)` | |
| [pennylane/operation.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2241/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wZXJhdGlvbi5weQ==) | `95.95% <0.00%> (ø)` | |
| [pennylane/transforms/op\_transform.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2241/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvb3BfdHJhbnNmb3JtLnB5) | `89.36% <0.00%> (+70.21%)` | :arrow_up: |

------

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

cvjjm: Extra awesome that `qml.matrix()` if auto differentiable!

Patch:
diff --git a/pennylane/__init__.py b/pennylane/__init__.py
index 4deb27a80..a8b179447 100644
--- a/pennylane/__init__.py
+++ b/pennylane/__init__.py
@@ -74,6 +74,7 @@ from pennylane.transforms import (
     is_commuting,
     simplify,
 )
+from pennylane.ops.functions import *
 from pennylane.optimize import *
 from pennylane.vqe import ExpvalCost, VQECost
 
diff --git a/pennylane/operation.py b/pennylane/operation.py
index 62250b0d5..dea58d9be 100644
--- a/pennylane/operation.py
+++ b/pennylane/operation.py
@@ -570,7 +570,7 @@ class Operator(abc.ABC):
         .. warning::
 
             The ``matrix`` property is deprecated and will be removed in
-            an upcoming release.
+            an upcoming release. Please use :class:`qml.matrix <.pennylane.matrix>` instead.
 
         **Example:**
 
@@ -583,7 +583,8 @@ class Operator(abc.ABC):
             array: matrix representation
         """
         warnings.warn(
-            "The 'matrix' property is deprecated and will be removed in an upcoming release.",
+            "The 'matrix' property is deprecated and will be removed in an upcoming release. "
+            "Please use 'qml.matrix' instead.",
             UserWarning,
         )
         return self.get_matrix()
diff --git a/pennylane/ops/functions/__init__.py b/pennylane/ops/functions/__init__.py
new file mode 100644
index 000000000..2f6c281c2
--- /dev/null
+++ b/pennylane/ops/functions/__init__.py
@@ -0,0 +1,17 @@
+# 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.
+"""
+This module contains functions that act on operators and tapes.
+"""
+from .matrix import matrix
diff --git a/pennylane/ops/functions/matrix.py b/pennylane/ops/functions/matrix.py
new file mode 100644
index 000000000..4a623032e
--- /dev/null
+++ b/pennylane/ops/functions/matrix.py
@@ -0,0 +1,145 @@
+# 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.
+"""
+This module contains the qml.matrix function.
+"""
+# pylint: disable=protected-access
+import pennylane as qml
+
+
+@qml.op_transform
+def matrix(op, *, wire_order=None):
+    r"""The matrix representation of an operation or quantum circuit.
+
+    Args:
+        op (.Operator, pennylane.QNode, .QuantumTape, or Callable): An operator, 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:
+        tensor_like or function: Function which accepts the same arguments as the QNode or quantum
+        function. When called, this function will return the unitary matrix in the appropriate
+        autodiff framework (Autograd, TensorFlow, PyTorch, JAX) given its parameters.
+
+    **Example**
+
+    Given an instantiated operator, ``qml.matrix`` returns the matrix representation:
+
+    >>> op = qml.RX(0.54, wires=0)
+    >>> qml.matrix(op)
+    [[0.9637709+0.j         0.       -0.26673144j]
+    [0.       -0.26673144j 0.9637709+0.j        ]]
+
+    It can also be used in a functional form:
+
+    >>> x = torch.tensor(0.6, requires_grad=True)
+    >>> matrix_fn = qml.matrix(qml.RX)
+    >>> matrix_fn(x)
+    tensor([[0.9553+0.0000j, 0.0000-0.2955j],
+            [0.0000-0.2955j, 0.9553+0.0000j]], grad_fn=<AddBackward0>)
+
+    In its functional form, it is fully differentiable with respect to gate arguments:
+
+    >>> loss = torch.real(torch.trace(matrix_fn(x, wires=0)))
+    >>> loss.backward()
+    >>> x.grad
+    tensor(-0.5910)
+
+    This operator transform can also be applied to QNodes, tapes, and quantum functions
+    that contain multiple operations; see Usage Details below for more details.
+
+    .. UsageDetails::
+
+        ``qml.matrix`` can also be used with QNodes, tapes, or quantum functions that
+        contain multiple operations.
+
+        Consider the following quantum function:
+
+        .. code-block:: python3
+
+            def circuit(theta):
+                qml.RX(theta, wires=1)
+                qml.PauliZ(wires=0)
+
+        We can use ``qml.matrix`` to generate a new function that returns the unitary matrix
+        corresponding to the function ``circuit``:
+
+        >>> matrix_fn = qml.matrix(circuit)
+        >>> theta = np.pi / 4
+        >>> matrix_fn(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 = qml.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)`.
+
+        This unitary matrix can also be used in differentiable calculations. For example, consider the
+        following cost function:
+
+        .. code-block:: python
+
+            def circuit(theta):
+                qml.RX(theta, wires=1) qml.PauliZ(wires=0)
+                qml.CNOT(wires=[0, 1])
+
+            def cost(theta):
+                matrix = qml.matrix(circuit)(theta)
+                return np.real(np.trace(matrix))
+
+        Since this cost function returns a real scalar as a function of ``theta``, we can differentiate
+        it:
+
+        >>> theta = np.array(0.3, requires_grad=True)
+        >>> cost(theta)
+        1.9775421558720845
+        >>> qml.grad(cost)(theta)
+        -0.14943813247359922
+    """
+    if isinstance(op, qml.operation.Tensor) and wire_order is not None:
+        op = 1.0 * op  # convert to a Hamiltonian
+
+    if isinstance(op, qml.Hamiltonian):
+        return qml.utils.sparse_hamiltonian(op, wires=wire_order).toarray()
+
+    return op.get_matrix(wire_order=wire_order)
+
+
+@matrix.tape_transform
+def _matrix(tape, wire_order=None):
+    """Defines how matrix works if applied to a tape containing multiple operations."""
+    params = tape.get_parameters(trainable_only=False)
+    interface = qml.math._multi_dispatch(params)
+
+    wire_order = wire_order or tape.wires
+    n_wires = len(wire_order)
+
+    # initialize the unitary matrix
+    unitary_matrix = qml.math.eye(2**n_wires, like=interface)
+
+    for op in tape.operations:
+        U = matrix(op, wire_order=wire_order)
+        unitary_matrix = qml.math.dot(U, unitary_matrix)
+
+    return unitary_matrix


Test patch:
diff --git a/tests/ops/functions/test_matrix.py b/tests/ops/functions/test_matrix.py
new file mode 100644
index 000000000..101aca5fc
--- /dev/null
+++ b/tests/ops/functions/test_matrix.py
@@ -0,0 +1,587 @@
+# 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
+
+import pennylane as qml
+from pennylane import numpy as np
+from pennylane.transforms.op_transforms import OperationTransformError
+
+from gate_data import I, X, Y, Z, H, S, CNOT, Roty as RY
+
+one_qubit_no_parameter = [
+    qml.PauliX,
+    qml.PauliY,
+    qml.PauliZ,
+    qml.Hadamard,
+    qml.S,
+    qml.T,
+    qml.SX,
+]
+
+
+one_qubit_one_parameter = [qml.RX, qml.RY, qml.RZ, qml.PhaseShift]
+
+
+class TestSingleOperation:
+    @pytest.mark.parametrize("op_class", one_qubit_no_parameter)
+    def test_non_parametric_instantiated(self, op_class):
+        """Verify that the matrices of non-parametric one qubit gates is correct
+        when provided as an instantiated operation"""
+        op = op_class(wires=0)
+        res = qml.matrix(op)
+        expected = op.get_matrix()
+        assert np.allclose(res, expected)
+
+    @pytest.mark.parametrize("op_class", one_qubit_no_parameter)
+    def test_non_parametric_gates_qfunc(self, op_class):
+        """Verify that the matrices of non-parametric one qubit gates is correct
+        when provided as a qfunc"""
+        res = qml.matrix(op_class)(wires=0)
+        expected = op_class(wires=0).get_matrix()
+        assert np.allclose(res, expected)
+
+    @pytest.mark.parametrize("op_class", one_qubit_no_parameter)
+    def test_non_parametric_gates_qnode(self, op_class):
+        """Verify that the matrices of non-parametric one qubit gates is correct
+        when provided as a QNode"""
+        dev = qml.device("default.qubit", wires=1)
+        qnode = qml.QNode(lambda: op_class(wires=0) and qml.probs(wires=0), dev)
+        res = qml.matrix(qnode)()
+        expected = op_class(wires=0).get_matrix()
+        assert np.allclose(res, expected)
+
+    @pytest.mark.parametrize("op_class", one_qubit_no_parameter)
+    def test_matrix_expansion(self, op_class):
+        """Verify that matrices are correctly expanded when a wire order is provided"""
+        res = qml.matrix(op_class, wire_order=[1, 0, 2])(wires=0)
+        expected = np.kron(np.eye(2), np.kron(op_class(wires=0).get_matrix(), np.eye(2)))
+        assert np.allclose(res, expected)
+
+    @pytest.mark.parametrize("op_class", one_qubit_one_parameter)
+    def test_parametric_gates_instantiated(self, op_class):
+        """Verify that the matrix of non-parametric one qubit gates is correct
+        when provided as an instantiated operation"""
+        op = op_class(0.54, wires=0)
+        res = qml.matrix(op)
+        expected = op.get_matrix()
+        assert np.allclose(res, expected)
+
+    @pytest.mark.parametrize("op_class", one_qubit_one_parameter)
+    def test_parametric_gates_qfunc(self, op_class):
+        """Verify that the matrices of non-parametric one qubit gates is correct
+        when provided as a qfunc"""
+        res = qml.matrix(op_class)(0.54, wires=0)
+        expected = op_class(0.54, wires=0).get_matrix()
+        assert np.allclose(res, expected)
+
+    @pytest.mark.parametrize("op_class", one_qubit_one_parameter)
+    def test_parametric_gates_qnode(self, op_class):
+        """Verify that the matrices of non-parametric one qubit gates is correct
+        when provided as a QNode"""
+        dev = qml.device("default.qubit", wires=1)
+        qnode = qml.QNode(lambda x: op_class(x, wires=0) and qml.probs(wires=0), dev)
+        res = qml.matrix(qnode)(0.54)
+        expected = op_class(0.54, wires=0).get_matrix()
+        assert np.allclose(res, expected)
+
+    @pytest.mark.parametrize("op_class", one_qubit_one_parameter)
+    def test_adjoint(self, op_class):
+        """Test that the adjoint is correctly taken into account"""
+        res = qml.matrix(qml.adjoint(op_class))(0.54, wires=0)
+        expected = op_class(-0.54, wires=0).get_matrix()
+        assert np.allclose(res, expected)
+
+    def test_ctrl(self):
+        """Test that the ctrl is correctly taken into account"""
+        res = qml.matrix(qml.ctrl(qml.PauliX, 0))(wires=1)
+        expected = CNOT
+        assert np.allclose(res, expected)
+
+    @pytest.mark.parametrize("target_wire", [0, 2, 3, 4])
+    def test_CNOT_permutations(self, target_wire):
+        """Test CNOT: 2-qubit gate with different target wires, some non-adjacent."""
+        res = qml.matrix(qml.CNOT, wire_order=[0, 1, 2, 3, 4])(wires=[1, target_wire])
+
+        # compute the expected matrix
+        perm = np.swapaxes(
+            np.swapaxes(np.arange(2**5).reshape([2] * 5), 0, 1), 0, target_wire
+        ).flatten()
+        expected = reduce(np.kron, [CNOT, I, I, I])[:, perm][perm]
+        assert np.allclose(res, expected)
+
+    def test_hamiltonian(self):
+        """Test that the matrix of a Hamiltonian is correctly returned"""
+        H = qml.PauliZ(0) @ qml.PauliY(1) - 0.5 * qml.PauliX(1)
+        mat = qml.matrix(H, wire_order=[1, 0, 2])
+        expected = reduce(np.kron, [Y, Z, I]) - 0.5 * reduce(np.kron, [X, I, I])
+
+    @pytest.mark.xfail(
+        reason="This test will fail because Hamiltonians are not queued to tapes yet!"
+    )
+    def test_hamiltonian_qfunc(self):
+        """Test that the matrix of a Hamiltonian is correctly returned"""
+
+        def ansatz(x):
+            return qml.PauliZ(0) @ qml.PauliY(1) - x * qml.PauliX(1)
+
+        x = 0.5
+        mat = qml.matrix(ansatz, wire_order=[1, 0, 2])(x)
+        expected = reduce(np.kron, [Y, Z, I]) - x * reduce(np.kron, [X, I, I])
+
+
+class TestMultipleOperations:
+    def test_multiple_operations_tape(self):
+        """Check the total matrix for a tape containing multiple gates"""
+        wire_order = ["a", "b", "c"]
+
+        with qml.tape.QuantumTape() as tape:
+            qml.PauliX(wires="a")
+            qml.S(wires="b")
+            qml.Hadamard(wires="c")
+            qml.CNOT(wires=["b", "c"])
+
+        matrix = qml.matrix(tape, wire_order)
+        expected_matrix = np.kron(I, CNOT) @ np.kron(X, np.kron(S, H))
+        assert np.allclose(matrix, expected_matrix)
+
+    def test_multiple_operations_qfunc(self):
+        """Check the total matrix for a qfunc containing multiple gates"""
+        wire_order = ["a", "b", "c"]
+
+        def testcircuit():
+            qml.PauliX(wires="a")
+            qml.S(wires="b")
+            qml.Hadamard(wires="c")
+            qml.CNOT(wires=["b", "c"])
+
+        matrix = qml.matrix(testcircuit, wire_order)()
+        expected_matrix = np.kron(I, CNOT) @ np.kron(X, np.kron(S, H))
+        assert np.allclose(matrix, expected_matrix)
+
+    def test_multiple_operations_qnode(self):
+        """Check the total matrix for a QNode containing multiple gates"""
+        dev = qml.device("default.qubit", wires=["a", "b", "c"])
+
+        @qml.qnode(dev)
+        def testcircuit():
+            qml.PauliX(wires="a")
+            qml.adjoint(qml.S)(wires="b")
+            qml.Hadamard(wires="c")
+            qml.CNOT(wires=["b", "c"])
+            return qml.expval(qml.PauliZ("a"))
+
+        matrix = qml.matrix(testcircuit)()
+        expected_matrix = np.kron(I, CNOT) @ np.kron(X, np.kron(np.linalg.inv(S), H))
+        assert np.allclose(matrix, expected_matrix)
+
+
+class TestCustomWireOrdering:
+    def test_tensor_wire_oder(self):
+        """Test wire order of a tensor product"""
+        H = qml.PauliZ(0) @ qml.PauliX(1)
+        res = qml.matrix(H, wire_order=[0, 2, 1])
+        expected = np.kron(Z, np.kron(I, X))
+        assert np.allclose(res, expected)
+
+    def test_tape_wireorder(self):
+        """Test changing the wire order when using a tape"""
+
+        with qml.tape.QuantumTape() as tape:
+            qml.PauliX(wires=0)
+            qml.PauliY(wires=1)
+            qml.PauliZ(wires=2)
+
+        matrix = qml.matrix(tape)
+        expected_matrix = np.kron(X, np.kron(Y, Z))
+        assert np.allclose(matrix, expected_matrix)
+
+        matrix = qml.matrix(tape, wire_order=[1, 0, 2])
+        expected_matrix = np.kron(Y, np.kron(X, Z))
+        assert np.allclose(matrix, expected_matrix)
+
+    def test_qfunc_wireorder(self):
+        """Test changing the wire order when using a qfunc"""
+
+        def testcircuit():
+            qml.PauliX(wires=0)
+            qml.PauliY(wires=1)
+            qml.PauliZ(wires=2)
+
+        matrix = qml.matrix(testcircuit)()
+        expected_matrix = np.kron(X, np.kron(Y, Z))
+        assert np.allclose(matrix, expected_matrix)
+
+        matrix = qml.matrix(testcircuit, wire_order=[1, 0, 2])()
+        expected_matrix = np.kron(Y, np.kron(X, Z))
+        assert np.allclose(matrix, expected_matrix)
+
+    def test_qnode_wireorder(self):
+        """Test changing the wire order when using a QNode"""
+        dev = qml.device("default.qubit", wires=[1, 0, 2, 3])
+
+        @qml.matrix()
+        @qml.qnode(dev)
+        def testcircuit(x):
+            qml.PauliX(wires=0)
+            qml.RY(x, wires=1)
+            qml.PauliZ(wires=2)
+            return qml.expval(qml.PauliZ(0))
+
+        x = 0.5
+
+        # default wire ordering will come from the device
+        expected_matrix = np.kron(RY(x), np.kron(X, np.kron(Z, I)))
+        assert np.allclose(testcircuit(x), expected_matrix)
+
+        @qml.matrix(wire_order=[1, 0, 2])
+        @qml.qnode(dev)
+        def testcircuit(x):
+            qml.PauliX(wires=0)
+            qml.RY(x, wires=1)
+            qml.PauliZ(wires=2)
+            return qml.expval(qml.PauliZ(0))
+
+        expected_matrix = np.kron(RY(x), np.kron(X, Z))
+        assert np.allclose(testcircuit(x), expected_matrix)
+
+
+class TestTemplates:
+    """These tests are useful as they test operators that might not have
+    matrix forms defined, requiring decomposition."""
+
+    def test_instantiated(self):
+        """Test an instantiated template"""
+        weights = np.array([[[0.1, 0.2, 0.3], [0.1, 0.2, 0.3]]])
+        op = qml.StronglyEntanglingLayers(weights, wires=[0, 1])
+        res = qml.matrix(op)
+
+        with qml.tape.QuantumTape() as tape:
+            op.decomposition()
+
+        expected = qml.matrix(tape)
+        np.allclose(res, expected)
+
+    def test_qfunc(self):
+        """Test a template used within a qfunc"""
+
+        def circuit(weights, x):
+            qml.StronglyEntanglingLayers(weights, wires=[0, 1])
+            qml.RX(x, wires=0)
+
+        weights = np.array([[[0.1, 0.2, 0.3], [0.1, 0.2, 0.3]]])
+        x = 0.54
+        res = qml.matrix(circuit)(weights, x)
+
+        op = qml.StronglyEntanglingLayers(weights, wires=[0, 1])
+
+        with qml.tape.QuantumTape() as tape:
+            op.decomposition()
+            qml.RX(x, wires=0)
+
+        expected = qml.matrix(tape)
+        np.allclose(res, expected)
+
+    def test_nested_instantiated(self):
+        """Test an operation that must be decomposed twice"""
+
+        class CustomOp(qml.operation.Operation):
+            num_params = 1
+            num_wires = 2
+
+            @staticmethod
+            def compute_decomposition(weights, wires):
+                return [qml.StronglyEntanglingLayers(weights, wires=wires)]
+
+        weights = np.array([[[0.1, 0.2, 0.3], [0.1, 0.2, 0.3]]])
+        op = CustomOp(weights, wires=[0, 1])
+        res = qml.matrix(op)
+
+        op = qml.StronglyEntanglingLayers(weights, wires=[0, 1])
+        with qml.tape.QuantumTape() as tape:
+            op.decomposition()
+
+        expected = qml.matrix(tape)
+        np.allclose(res, expected)
+
+    def test_nested_qfunc(self):
+        """Test an operation that must be decomposed twice"""
+
+        class CustomOp(qml.operation.Operation):
+            num_params = 1
+            num_wires = 2
+
+            @staticmethod
+            def compute_decomposition(weights, wires):
+                return [qml.StronglyEntanglingLayers(weights, wires=wires)]
+
+        def circuit(weights, x):
+            CustomOp(weights, wires=[0, 1])
+            qml.RX(x, wires=0)
+
+        weights = np.array([[[0.1, 0.2, 0.3], [0.1, 0.2, 0.3]]])
+        x = 0.54
+        res = qml.matrix(circuit)(weights, x)
+
+        op = qml.StronglyEntanglingLayers(weights, wires=[0, 1])
+
+        with qml.tape.QuantumTape() as tape:
+            op.decomposition()
+            qml.RX(x, wires=0)
+
+        expected = qml.matrix(tape)
+        np.allclose(res, expected)
+
+
+class TestValidation:
+    def test_invalid_argument(self):
+        """Assert error raised when input is neither a tape, QNode, nor quantum function"""
+        with pytest.raises(
+            OperationTransformError,
+            match="Input is not an Operator, tape, QNode, or quantum function",
+        ):
+            qml.matrix(None)(0.5)
+
+    def test_wrong_function(self):
+        """Assert error raised when input function is not a quantum function"""
+
+        def testfunction(x):
+            return x
+
+        with pytest.raises(OperationTransformError, match="function contains no quantum operation"):
+            qml.matrix(testfunction)(0)
+
+    def test_inconsistent_wires(self):
+        """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"]
+
+        with pytest.raises(
+            OperationTransformError,
+            match=r"Wires in circuit \[1, 0\] are inconsistent with those in wire_order \[0, 'b'\]",
+        ):
+            matrix = qml.matrix(circuit, wire_order=wires)()
+
+
+class TestInterfaces:
+    def test_tf(self):
+        """Test with tensorflow interface"""
+        tf = pytest.importorskip("tensorflow")
+
+        @qml.matrix
+        def circuit(beta, theta):
+            qml.RZ(beta, wires=0)
+            qml.RZ(theta[0], wires=1)
+            qml.CRY(theta[1], wires=[1, 2])
+
+        beta = 0.1
+        # input tensorflow parameters
+        theta = tf.Variable([0.2, 0.3])
+        matrix = circuit(beta, theta)
+
+        # expected matrix
+        theta_np = theta.numpy()
+        matrix1 = np.kron(
+            qml.RZ(beta, wires=0).get_matrix(),
+            np.kron(qml.RZ(theta_np[0], wires=1).get_matrix(), I),
+        )
+        matrix2 = np.kron(I, qml.CRY(theta_np[1], wires=[1, 2]).get_matrix())
+        expected_matrix = matrix2 @ matrix1
+
+        assert np.allclose(matrix, expected_matrix)
+
+    def test_torch(self):
+        """Test with torch interface"""
+
+        torch = pytest.importorskip("torch", minversion="1.8")
+
+        dev = qml.device("default.qubit", wires=3)
+
+        @qml.matrix
+        @qml.qnode(dev, interface="torch")
+        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))
+
+        # input torch parameters
+        theta = torch.tensor([0.1, 0.2, 0.3])
+        matrix = circuit(theta)
+
+        # expected matrix
+        matrix1 = np.kron(
+            qml.RZ(theta[0], wires=0).get_matrix(),
+            np.kron(qml.RZ(theta[1], wires=1).get_matrix(), I),
+        )
+        matrix2 = np.kron(I, qml.CRY(theta[2], wires=[1, 2]).get_matrix())
+        expected_matrix = matrix2 @ matrix1
+
+        assert np.allclose(matrix, expected_matrix)
+
+    def test_autograd(self):
+        """Test with autograd interface"""
+
+        @qml.matrix
+        def circuit(theta):
+            qml.RZ(theta[0], wires=0)
+            qml.RZ(theta[1], wires=1)
+            qml.CRY(theta[2], wires=[1, 2])
+
+        # set input parameters
+        theta = np.array([0.1, 0.2, 0.3], requires_grad=True)
+        matrix = circuit(theta)
+
+        # expected matrix
+        matrix1 = np.kron(
+            qml.RZ(theta[0], wires=0).get_matrix(),
+            np.kron(qml.RZ(theta[1], wires=1).get_matrix(), I),
+        )
+        matrix2 = np.kron(I, qml.CRY(theta[2], wires=[1, 2]).get_matrix())
+        expected_matrix = matrix2 @ matrix1
+
+        assert np.allclose(matrix, expected_matrix)
+
+    def test_get_unitary_matrix_interface_jax(self):
+        """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)
+
+        @qml.matrix
+        def circuit(theta):
+            qml.RZ(theta[0], wires=0)
+            qml.RZ(theta[1], wires=1)
+            qml.CRY(theta[2], wires=[1, 2])
+
+        # input jax parameters
+        theta = jnp.array([0.1, 0.2, 0.3], dtype=jnp.float64)
+
+        matrix = circuit(theta)
+
+        # expected matrix
+        matrix1 = np.kron(
+            qml.RZ(theta[0], wires=0).get_matrix(),
+            np.kron(qml.RZ(theta[1], wires=1).get_matrix(), I),
+        )
+        matrix2 = np.kron(I, qml.CRY(theta[2], wires=[1, 2]).get_matrix())
+        expected_matrix = matrix2 @ matrix1
+
+        assert np.allclose(matrix, expected_matrix)
+
+
+class TestDifferentiation:
+    @pytest.mark.parametrize("v", np.linspace(0.2, 1.6, 8))
+    def test_jax(self, v):
+
+        jax = pytest.importorskip("jax")
+
+        def circuit(theta):
+            qml.RX(theta, wires=0)
+            qml.PauliZ(wires=0)
+            qml.CNOT(wires=[0, 1])
+
+        def loss(theta):
+            U = qml.matrix(circuit)(theta)
+            return qml.math.real(qml.math.trace(U))
+
+        x = jax.numpy.array(v)
+
+        l = loss(x)
+        dl = jax.grad(loss)(x)
+        matrix = qml.matrix(circuit)(x)
+
+        assert isinstance(matrix, jax.numpy.ndarray)
+        assert np.allclose(l, 2 * np.cos(v / 2))
+        assert np.allclose(dl, -np.sin(v / 2))
+
+    @pytest.mark.parametrize("v", np.linspace(0.2, 1.6, 8))
+    def test_torch(self, v):
+
+        torch = pytest.importorskip("torch")
+
+        def circuit(theta):
+            qml.RX(theta, wires=0)
+            qml.PauliZ(wires=0)
+            qml.CNOT(wires=[0, 1])
+
+        def loss(theta):
+            U = qml.matrix(circuit)(theta)
+            return qml.math.real(qml.math.trace(U))
+
+        x = torch.tensor(v, requires_grad=True)
+        l = loss(x)
+        l.backward()
+        dl = x.grad
+        matrix = qml.matrix(circuit)(x)
+
+        assert isinstance(matrix, torch.Tensor)
+        assert np.allclose(l.detach(), 2 * np.cos(v / 2))
+        assert np.allclose(dl.detach(), -np.sin(v / 2))
+
+    @pytest.mark.parametrize("v", np.linspace(0.2, 1.6, 8))
+    def test_tensorflow(self, v):
+
+        tf = pytest.importorskip("tensorflow")
+
+        def circuit(theta):
+            qml.RX(theta, wires=0)
+            qml.PauliZ(wires=0)
+            qml.CNOT(wires=[0, 1])
+
+        def loss(theta):
+            U = qml.matrix(circuit)(theta)
+            return qml.math.real(qml.math.trace(U))
+
+        x = tf.Variable(v)
+        with tf.GradientTape() as tape:
+            l = loss(x)
+        dl = tape.gradient(l, x)
+        matrix = qml.matrix(circuit)(x)
+
+        assert isinstance(matrix, tf.Tensor)
+        assert np.allclose(l, 2 * np.cos(v / 2))
+        assert np.allclose(dl, -np.sin(v / 2))
+
+    @pytest.mark.parametrize("v", np.linspace(0.2, 1.6, 8))
+    def test_get_unitary_matrix_autograd_differentiable(self, v):
+        def circuit(theta):
+            qml.RX(theta, wires=0)
+            qml.PauliZ(wires=0)
+            qml.CNOT(wires=[0, 1])
+
+        def loss(theta):
+            U = qml.matrix(circuit)(theta)
+            return qml.math.real(qml.math.trace(U))
+
+        x = np.array(v, requires_grad=True)
+        l = loss(x)
+        dl = qml.grad(loss)(x)
+        matrix = qml.matrix(circuit)(x)
+
+        assert isinstance(matrix, qml.numpy.tensor)
+        assert np.allclose(l, 2 * np.cos(v / 2))
+        assert np.allclose(dl, -np.sin(v / 2))