Repository: PennyLaneAI/pennylane
Pull request #1427 (2021-06-22 20:18:23)

Adds elementary gate decomposition for single-qubit `QubitUnitary`
**Context:** Currently, there is no decomposition/synthesis for the arbitrary `QubitUnitary` operation.

**Description of the Change:** Adds the math needed to express arbitrary single-qubit matrices as a `Rot` gate (or an `RZ` gate, if they are diagonal).

**Benefits:** This should allow `QubitUnitary` to run on more devices. Furthermore, I think that by using the decomposition we can make `QubitUnitary` fully differentiable, at least for the single-qubit case (see example below).

**Possible Drawbacks:** None

**Related GitHub Issues:** None

**Example:**

```python
import pennylane as qml
from pennylane import numpy as np

dev = qml.device('default.qubit', wires=1)

def qfunc(x):
    qml.Hadamard(wires=0)

    # Create a unitary matrix based on an input parameter
    Ux = np.array([[1, 0], [0, np.exp(1j * x)]])
    ops = qml.QubitUnitary.decomposition(Ux, wires=0)
    for op in ops:
        op.queue()

    return qml.expval(qml.PauliX(0))
```
```pycon
>>> original_qnode = qml.QNode(qfunc, dev, diff_method="adjoint")
>>> print(qml.grad(original_qnode)(0.2))
-0.19866933079506127
```

github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [.github/CHANGELOG.md](/PennyLaneAI/pennylane/blob/master/.github/CHANGELOG.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
glassnotes: > Part of me wonder if this feels like a compilation transform? Becuase there will be use cases where users want to force a unitary decomposition on a device that natively supports it. But I don't know? it would add overhead compared to simply calling QubitUnitary.decomposition() manually.

I am wondering this too; so, for additional context, I'd like to do this for the 2-qubit case as well. That's going to make the `decomposition` method extremely hairy. If we use transforms, we can make things more flexible (e.g., we can have different transforms for `ZYZ` versus `ZXZ` decompositions), and better separate the 1- and 2-qubit cases from each other. Then we can have something that passes through the circuit, finds `QubitUnitary`s, and depending on the number of qubits calls the appropriate transform.
josh146: > I am wondering this too; so, for additional context, I'd like to do this for the 2-qubit case as well. 

Oh! Maybe the answer is... both?

- We have transforms to make it flexible and keep the code organized.
- The `QubitUnitary.decomposition()` method calls the transforms as needed, in order to support devices where `QubitUnitary` is not defined.
glassnotes: >     * We have transforms to make it flexible and keep the code organized.
> 
>     * The `QubitUnitary.decomposition()` method calls the transforms as needed, in order to support devices where `QubitUnitary` is not defined.

I like it! I'll give it a go and we can see what it looks like. After deciding on that I'll make the rest of the tests.
glassnotes: @josh146 Thanks for the review! Have addressed all your comments, and split the decomposition methods into their own submodule, separate from the transform. 
codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1427?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#1427](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1427?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (b9480be) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/614db57b02536bb09fbfbae853ceff8878b928af?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (614db57) will **decrease** coverage by `52.40%`.
> The diff coverage is `35.84%`.

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

```diff
@@             Coverage Diff             @@
##           master    #1427       +/-   ##
===========================================
- Coverage   98.23%   45.83%   -52.41%     
===========================================
  Files         160      163        +3     
  Lines       11966    12043       +77     
===========================================
- Hits        11755     5520     -6235     
- Misses        211     6523     +6312     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1427?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [.../transforms/decompositions/single\_qubit\_unitary.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1427/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvZGVjb21wb3NpdGlvbnMvc2luZ2xlX3F1Yml0X3VuaXRhcnkucHk=) | `19.23% <19.23%> (ø)` | |
| [pennylane/ops/qubit.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1427/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9xdWJpdC5weQ==) | `72.30% <28.57%> (-26.42%)` | :arrow_down: |
| [...ane/transforms/decompose\_single\_qubit\_unitaries.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1427/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvZGVjb21wb3NlX3NpbmdsZV9xdWJpdF91bml0YXJpZXMucHk=) | `46.66% <46.66%> (ø)` | |
| [pennylane/transforms/\_\_init\_\_.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1427/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvX19pbml0X18ucHk=) | `100.00% <100.00%> (ø)` | |
| [pennylane/transforms/decompositions/\_\_init\_\_.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1427/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvZGVjb21wb3NpdGlvbnMvX19pbml0X18ucHk=) | `100.00% <100.00%> (ø)` | |
| [pennylane/beta/\_\_init\_\_.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1427/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2JldGEvX19pbml0X18ucHk=) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| [pennylane/fourier/utils.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1427/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ZvdXJpZXIvdXRpbHMucHk=) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| [pennylane/devices/tf\_ops.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1427/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvdGZfb3BzLnB5) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| [pennylane/devices/jax\_ops.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1427/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvamF4X29wcy5weQ==) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| [pennylane/interfaces/torch.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1427/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvdG9yY2gucHk=) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| ... and [129 more](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1427/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | |

------

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


Patch:
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
index b2545a8e1..7c80c3c28 100644
--- a/.github/CHANGELOG.md
+++ b/.github/CHANGELOG.md
@@ -2,6 +2,61 @@
 
 <h3>New features since last release</h3>
 
+* A decomposition has been added to ``QubitUnitary`` that makes the
+  single-qubit case fully differentiable in all interfaces. Furthermore,
+  a quantum function transform, ``unitary_to_rot()``, has been added to decompose all
+  single-qubit instances of ``QubitUnitary`` in a quantum circuit.
+  [(#1427)](https://github.com/PennyLaneAI/pennylane/pull/1427)
+
+  Instances of ``QubitUnitary`` may now be decomposed directly to ``Rot``
+  operations, or ``RZ`` operations if the input matrix is diagonal. For
+  example, let
+
+  ```python
+  >>> U = np.array([
+      [-0.28829348-0.78829734j,  0.30364367+0.45085995j],
+      [ 0.53396245-0.10177564j,  0.76279558-0.35024096j]
+  ])
+  ```
+
+  Then, we can compute the decomposition as:
+
+  ```pycon
+  >>> qml.QubitUnitary.decomposition(U, wires=0)
+  [Rot(-0.24209530281458358, 1.1493817777199102, 1.733058145303424, wires=[0])]
+  ```
+
+  We can also apply the transform directly to a quantum function, and compute the
+  gradients of parameters used to construct the unitary matrices.
+
+  ```python
+  def qfunc_with_qubit_unitary(angles):
+      z, x = angles[0], angles[1]
+
+      Z_mat = np.array([[np.exp(-1j * z / 2), 0.0], [0.0, np.exp(1j * z / 2)]])
+
+      c = np.cos(x / 2)
+      s = np.sin(x / 2) * 1j
+      X_mat = np.array([[c, -s], [-s, c]])
+
+      qml.Hadamard(wires="a")
+      qml.QubitUnitary(Z_mat, wires="a")
+      qml.QubitUnitary(X_mat, wires="b")
+      qml.CNOT(wires=["b", "a"])
+      return qml.expval(qml.PauliX(wires="a"))
+  ```
+
+  ```pycon
+  >>> dev = qml.device("default.qubit", wires=["a", "b"])
+  >>> transformed_qfunc = qml.transforms.unitary_to_rot(qfunc_with_qubit_unitary)
+  >>> transformed_qnode = qml.QNode(transformed_qfunc, dev)
+  >>> input = np.array([0.3, 0.4], requires_grad=True)
+  >>> transformed_qnode(input)
+  tensor(0.95533649, requires_grad=True)
+  >>> qml.grad(transformed_qnode)(input)
+  array([-0.29552021,  0.        ])
+  ```
+
 * The new ``qml.apply`` function can be used to add operations that might have
   already been instantiated elsewhere to the QNode and other queuing contexts:
   [(#1433)](https://github.com/PennyLaneAI/pennylane/pull/1433)
diff --git a/pennylane/ops/qubit.py b/pennylane/ops/qubit.py
index c3db03e5b..5002aa12f 100644
--- a/pennylane/ops/qubit.py
+++ b/pennylane/ops/qubit.py
@@ -2159,6 +2159,17 @@ class QubitUnitary(Operation):
 
         return U
 
+    @staticmethod
+    def decomposition(U, wires):
+        # Decomposes arbitrary single-qubit unitaries as Rot gates (RZ - RY - RZ format),
+        # or a single RZ for diagonal matrices.
+        if qml.math.shape(U) == (2, 2):
+            wire = Wires(wires)[0]
+            decomp_ops = qml.transforms.decompositions.zyz_decomposition(U, wire)
+            return decomp_ops
+
+        raise NotImplementedError("Decompositions only supported for single-qubit unitaries")
+
     def adjoint(self):
         return QubitUnitary(qml.math.T(qml.math.conj(self.matrix)), wires=self.wires)
 
diff --git a/pennylane/transforms/__init__.py b/pennylane/transforms/__init__.py
index 0460240dd..9a18760fb 100644
--- a/pennylane/transforms/__init__.py
+++ b/pennylane/transforms/__init__.py
@@ -45,10 +45,19 @@ containing quantum operations) that are used to construct QNodes.
 
     ~adjoint
     ~ctrl
+    ~transforms.unitary_to_rot
     ~transforms.invisible
     ~apply_controlled_Q
     ~quantum_monte_carlo
 
+There are also utility functions and decompositions available that assist with
+both transforms, and decompositions within the larger PennyLane codebase.
+
+.. autosummary::
+    :toctree: api
+
+    ~transforms.zyz_decomposition
+
 Transforms that act on tapes
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
@@ -75,14 +84,17 @@ to help build custom QNode, quantum function, and tape transforms:
     ~qfunc_transform
     ~transforms.make_tape
 """
+# Import the decorators first to prevent circular imports when used in other transforms
+from .qfunc_transforms import make_tape, single_tape_transform, qfunc_transform
 from .adjoint import adjoint
 from .classical_jacobian import classical_jacobian
 from .control import ControlledOperation, ctrl
+from .decompositions import zyz_decomposition
 from .draw import draw
 from .hamiltonian_expand import hamiltonian_expand
 from .invisible import invisible
 from .measurement_grouping import measurement_grouping
 from .metric_tensor import metric_tensor, metric_tensor_tape
 from .specs import specs
-from .qfunc_transforms import make_tape, single_tape_transform, qfunc_transform
 from .qmc import apply_controlled_Q, quantum_monte_carlo
+from .unitary_to_rot import unitary_to_rot
diff --git a/pennylane/transforms/decompositions/__init__.py b/pennylane/transforms/decompositions/__init__.py
new file mode 100644
index 000000000..8e771a31d
--- /dev/null
+++ b/pennylane/transforms/decompositions/__init__.py
@@ -0,0 +1,18 @@
+# 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.
+r"""This module contains decompositions for (numerically-specified) arbitrary
+unitary operations into sequences of elementary operations.
+"""
+
+from .single_qubit_unitary import zyz_decomposition
diff --git a/pennylane/transforms/decompositions/single_qubit_unitary.py b/pennylane/transforms/decompositions/single_qubit_unitary.py
new file mode 100644
index 000000000..1ce322119
--- /dev/null
+++ b/pennylane/transforms/decompositions/single_qubit_unitary.py
@@ -0,0 +1,105 @@
+# Copyright 2018-2021 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Contains transforms and helpers functions for decomposing arbitrary unitary
+operations into elementary gates.
+"""
+import pennylane as qml
+from pennylane import math
+
+
+def _convert_to_su2(U):
+    r"""Check unitarity of a matrix and convert it to :math:`SU(2)` if possible.
+
+    Args:
+        U (array[complex]): A matrix, presumed to be :math:`2 \times 2` and unitary.
+
+    Returns:
+        array[complex]: A :math:`2 \times 2` matrix in :math:`SU(2)` that is
+        equivalent to U up to a global phase.
+    """
+    # Check unitarity
+    if not math.allclose(math.dot(U, math.T(math.conj(U))), math.eye(2), atol=1e-7):
+        raise ValueError("Operator must be unitary.")
+
+    # Compute the determinant
+    det = U[0, 0] * U[1, 1] - U[0, 1] * U[1, 0]
+
+    # Convert to SU(2) if it's not close to 1
+    if not math.allclose(det, [1.0]):
+        exp_angle = -1j * math.cast_like(math.angle(det), 1j) / 2
+        U = math.cast_like(U, exp_angle) * math.exp(exp_angle)
+
+    return U
+
+
+def zyz_decomposition(U, wire):
+    r"""Recover the decomposition of a single-qubit matrix :math:`U` in terms of
+    elementary operations.
+
+    Diagonal operations will be converted to a single :class:`.RZ` gate, while non-diagonal
+    operations will be converted to a :class:`.Rot` gate that implements the original operation
+    up to a global phase in the form :math:`RZ(\omega) RY(\theta) RZ(\phi)`.
+
+    Args:
+        U (tensor): A 2 x 2 unitary matrix.
+        wire (Union[Wires, Sequence[int] or int]): The wire on which to apply the operation.
+
+    Returns:
+        list[qml.Operation]: A ``Rot`` gate on the specified wire that implements ``U``
+        up to a global phase, or an equivalent ``RZ`` gate if ``U`` is diagonal.
+
+    **Example**
+
+    Suppose we would like to apply the following unitary operation:
+
+    .. code-block:: python3
+
+        U = np.array([
+            [-0.28829348-0.78829734j,  0.30364367+0.45085995j],
+            [ 0.53396245-0.10177564j,  0.76279558-0.35024096j]
+        ])
+
+    For PennyLane devices that cannot natively implement ``QubitUnitary``, we
+    can instead recover a ``Rot`` gate that implements the same operation, up
+    to a global phase:
+
+    >>> decomp = zyz_decomposition(U, 0)
+    >>> decomp
+    [Rot(-0.24209529417800013, 1.14938178234275, 1.7330581433950871, wires=[0])]
+    """
+    U = _convert_to_su2(U)
+
+    # Check if the matrix is diagonal; only need to check one corner.
+    # If it is diagonal, we don't need a full Rot, just return an RZ.
+    if math.allclose(U[0, 1], [0.0]):
+        omega = 2 * math.angle(U[1, 1])
+        return [qml.RZ(omega, wires=wire)]
+
+    # If not diagonal, compute the angle of the RY
+    cos2_theta_over_2 = math.abs(U[0, 0] * U[1, 1])
+    theta = 2 * math.arccos(math.sqrt(cos2_theta_over_2))
+
+    # If the top left element is 0, can only use the off-diagonal elements We
+    # have to be very careful with the math here to ensure things that get
+    # multiplied together are of the correct type in the different interfaces.
+    if math.allclose(U[0, 0], [0.0]):
+        phi = 1j * math.log(U[0, 1] / U[1, 0])
+        omega = -phi - math.cast_like(2 * math.angle(U[1, 0]), phi)
+    else:
+        el_division = U[0, 0] / U[1, 0]
+        tan_part = math.cast_like(math.tan(theta / 2), el_division)
+        omega = 1j * math.log(tan_part * el_division)
+        phi = -omega - math.cast_like(2 * math.angle(U[0, 0]), omega)
+
+    return [qml.Rot(math.real(phi), math.real(theta), math.real(omega), wires=wire)]
diff --git a/pennylane/transforms/unitary_to_rot.py b/pennylane/transforms/unitary_to_rot.py
new file mode 100644
index 000000000..cd05fccbd
--- /dev/null
+++ b/pennylane/transforms/unitary_to_rot.py
@@ -0,0 +1,82 @@
+# 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 for decomposing arbitrary single-qubit QubitUnitary gates into elementary gates.
+"""
+import pennylane as qml
+from pennylane.transforms import qfunc_transform
+from pennylane.transforms.decompositions import zyz_decomposition
+
+
+@qfunc_transform
+def unitary_to_rot(tape):
+    """Quantum function transform to decomposes all instances of single-qubit :class:`~.QubitUnitary`
+    operations to parametrized single-qubit operations.
+
+    Diagonal operations will be converted to a single :class:`.RZ` gate, while non-diagonal
+    operations will be converted to a :class:`.Rot` gate that implements the original operation
+    up to a global phase.
+
+    Args:
+        qfunc (function): a quantum function
+
+    **Example**
+
+    Suppose we would like to apply the following unitary operation:
+
+    .. code-block:: python3
+
+        U = np.array([
+            [-0.17111489+0.58564875j, -0.69352236-0.38309524j],
+            [ 0.25053735+0.75164238j,  0.60700543-0.06171855j]
+        ])
+
+    The ``unitary_to_rot`` transform enables us to decompose
+    such numerical operations (as well as unitaries that may be defined by parameters
+    within the QNode, and instantiated therein), while preserving differentiability.
+
+
+    .. code-block:: python3
+
+        def qfunc():
+            qml.QubitUnitary(U, wires=0)
+            return qml.expval(qml.PauliZ(0))
+
+    The original circuit is:
+
+    >>> dev = qml.device('default.qubit', wires=1)
+    >>> qnode = qml.QNode(qfunc, dev)
+    >>> print(qml.draw(qnode)())
+     0: ──U0──┤ ⟨Z⟩
+    U0 =
+    [[-0.17111489+0.58564875j -0.69352236-0.38309524j]
+     [ 0.25053735+0.75164238j  0.60700543-0.06171855j]]
+
+    We can use the transform to decompose the gate:
+
+    >>> transformed_qfunc = unitary_to_rot(qfunc)
+    >>> transformed_qnode = qml.QNode(transformed_qfunc, dev)
+    >>> print(qml.draw(transformed_qnode)())
+     0: ──Rot(-1.35, 1.83, -0.606)──┤ ⟨Z⟩
+
+    """
+    for op in tape.operations + tape.measurements:
+        if isinstance(op, qml.QubitUnitary):
+            # Only act on single-qubit unitary operations
+            if qml.math.shape(op.parameters[0]) != (2, 2):
+                continue
+
+            zyz_decomposition(op.parameters[0], op.wires[0])
+        else:
+            qml.apply(op)


Test patch:
diff --git a/tests/circuit_graph/test_qasm.py b/tests/circuit_graph/test_qasm.py
index 0c5962381..03f9a34d5 100644
--- a/tests/circuit_graph/test_qasm.py
+++ b/tests/circuit_graph/test_qasm.py
@@ -227,12 +227,12 @@ class TestToQasmUnitTests:
     def test_unsupported_gate(self):
         """Test an exception is raised if an unsupported operation is
         applied."""
-        U = np.array([[1, 1], [1, -1]]) / np.sqrt(2)
-
         with qml.tape.QuantumTape() as circuit:
-            qml.S(wires=0), qml.QubitUnitary(U, wires=[0, 1])
+            qml.S(wires=0), qml.DoubleExcitationPlus(0.5, wires=[0, 1, 2, 3])
 
-        with pytest.raises(ValueError, match="QubitUnitary not supported by the QASM serializer"):
+        with pytest.raises(
+            ValueError, match="DoubleExcitationPlus not supported by the QASM serializer"
+        ):
             res = circuit.to_openqasm()
 
     def test_rotations(self):
@@ -578,18 +578,19 @@ class TestQNodeQasmIntegrationTests:
     def test_unsupported_gate(self):
         """Test an exception is raised if an unsupported operation is
         applied."""
-        U = np.array([[1, 1], [1, -1]]) / np.sqrt(2)
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit", wires=4)
 
         @qml.qnode(dev)
         def qnode():
             qml.S(wires=0)
-            qml.QubitUnitary(U, wires=0)
+            qml.DoubleExcitationPlus(0.5, wires=[0, 1, 2, 3])
             return qml.expval(qml.PauliZ(0))
 
         qnode()
 
-        with pytest.raises(ValueError, match="QubitUnitary not supported by the QASM serializer"):
+        with pytest.raises(
+            ValueError, match="DoubleExcitationPlus not supported by the QASM serializer"
+        ):
             qnode.qtape.to_openqasm()
 
     def test_rotations(self):
diff --git a/tests/ops/test_qubit_ops.py b/tests/ops/test_qubit_ops.py
index c37cce5b5..d63994607 100644
--- a/tests/ops/test_qubit_ops.py
+++ b/tests/ops/test_qubit_ops.py
@@ -1566,6 +1566,37 @@ class TestOperations:
         with pytest.raises(ValueError, match="must be a square matrix"):
             qml.QubitUnitary(U, wires=0).matrix
 
+    @pytest.mark.parametrize(
+        "U,expected_gate,expected_params",
+        [  # First set of gates are diagonal and converted to RZ
+            (I, qml.RZ, [0]),
+            (Z, qml.RZ, [np.pi]),
+            (S, qml.RZ, [np.pi / 2]),
+            (T, qml.RZ, [np.pi / 4]),
+            (qml.RZ(0.3, wires=0).matrix, qml.RZ, [0.3]),
+            (qml.RZ(-0.5, wires=0).matrix, qml.RZ, [-0.5]),
+            # Next set of gates are non-diagonal and decomposed as Rots
+            (H, qml.Rot, [np.pi, np.pi / 2, 0]),
+            (X, qml.Rot, [0.0, np.pi, np.pi]),
+            (qml.Rot(0.2, 0.5, -0.3, wires=0).matrix, qml.Rot, [0.2, 0.5, -0.3]),
+            (np.exp(1j * 0.02) * qml.Rot(-1, 2, -3, wires=0).matrix, qml.Rot, [-1, 2, -3]),
+        ],
+    )
+    def test_qubit_unitary_decomposition(self, U, expected_gate, expected_params):
+        """Tests that single-qubit QubitUnitary decompositions are performed."""
+        decomp = qml.QubitUnitary.decomposition(U, wires=0)
+
+        assert len(decomp) == 1
+        assert isinstance(decomp[0], expected_gate)
+        assert np.allclose(decomp[0].parameters, expected_params)
+
+    def test_qubit_unitary_decomposition_multiqubit_invalid(self):
+        """Test that QubitUnitary is not decomposed for more than a single qubit."""
+        U = qml.CRZ(0.3, wires=[0, 1]).matrix
+
+        with pytest.raises(NotImplementedError, match="only supported for single-qubit"):
+            qml.QubitUnitary.decomposition(U, wires=[0, 1])
+
     def test_iswap_eigenval(self):
         """Tests that the ISWAP eigenvalue matches the numpy eigenvalues of the ISWAP matrix"""
         op = qml.ISWAP(wires=[0, 1])
diff --git a/tests/transforms/test_decompositions.py b/tests/transforms/test_decompositions.py
new file mode 100644
index 000000000..17d0ce188
--- /dev/null
+++ b/tests/transforms/test_decompositions.py
@@ -0,0 +1,107 @@
+# Copyright 2018-2021 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Tests for the QubitUnitary decomposition transforms.
+"""
+
+import pytest
+
+import pennylane as qml
+from pennylane import numpy as np
+
+from pennylane.wires import Wires
+from pennylane.transforms.decompositions import zyz_decomposition
+
+from gate_data import I, Z, S, T, H, X
+
+single_qubit_decomps = [
+    # First set of gates are diagonal and converted to RZ
+    (I, qml.RZ, [0.0]),
+    (Z, qml.RZ, [np.pi]),
+    (S, qml.RZ, [np.pi / 2]),
+    (T, qml.RZ, [np.pi / 4]),
+    (qml.RZ(0.3, wires=0).matrix, qml.RZ, [0.3]),
+    (qml.RZ(-0.5, wires=0).matrix, qml.RZ, [-0.5]),
+    # Next set of gates are non-diagonal and decomposed as Rots
+    (H, qml.Rot, [np.pi, np.pi / 2, 0.0]),
+    (X, qml.Rot, [0.0, np.pi, np.pi]),
+    (qml.Rot(0.2, 0.5, -0.3, wires=0).matrix, qml.Rot, [0.2, 0.5, -0.3]),
+    (np.exp(1j * 0.02) * qml.Rot(-1.0, 2.0, -3.0, wires=0).matrix, qml.Rot, [-1.0, 2.0, -3.0]),
+]
+
+
+class TestQubitUnitaryZYZDecomposition:
+    """Test that the decompositions are correct."""
+
+    def test_zyz_decomposition_invalid_input(self):
+        """Test that non-unitary operations throw errors when we try to decompose."""
+        with pytest.raises(ValueError, match="Operator must be unitary"):
+            zyz_decomposition(I + H, Wires("a"))
+
+    @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)
+
+    @pytest.mark.parametrize("U,expected_gate,expected_params", single_qubit_decomps)
+    def test_zyz_decomposition_torch(self, U, expected_gate, expected_params):
+        """Test that a one-qubit operation in Torch is correctly decomposed."""
+        torch = pytest.importorskip("torch")
+
+        U = torch.tensor(U, dtype=torch.complex64)
+
+        obtained_gates = zyz_decomposition(U, wire="a")
+
+        assert len(obtained_gates) == 1
+        assert isinstance(obtained_gates[0], expected_gate)
+        assert obtained_gates[0].wires == Wires("a")
+        assert qml.math.allclose(
+            [x.detach() for x in obtained_gates[0].parameters], expected_params
+        )
+
+    @pytest.mark.parametrize("U,expected_gate,expected_params", single_qubit_decomps)
+    def test_zyz_decomposition_tf(self, U, expected_gate, expected_params):
+        """Test that a one-qubit operation in Tensorflow is correctly decomposed."""
+        tf = pytest.importorskip("tensorflow")
+
+        U = tf.Variable(U, dtype=tf.complex64)
+
+        obtained_gates = zyz_decomposition(U, wire="a")
+
+        assert len(obtained_gates) == 1
+        assert isinstance(obtained_gates[0], expected_gate)
+        assert obtained_gates[0].wires == Wires("a")
+        assert qml.math.allclose([x.numpy() for x in obtained_gates[0].parameters], expected_params)
+
+    @pytest.mark.parametrize("U,expected_gate,expected_params", single_qubit_decomps)
+    def test_zyz_decomposition_jax(self, U, expected_gate, expected_params):
+        """Test that a one-qubit operation in JAX is correctly decomposed."""
+        jax = pytest.importorskip("jax")
+
+        U = jax.numpy.array(U, dtype=jax.numpy.complex64)
+
+        obtained_gates = zyz_decomposition(U, wire="a")
+
+        assert len(obtained_gates) == 1
+        assert isinstance(obtained_gates[0], expected_gate)
+        assert obtained_gates[0].wires == Wires("a")
+        assert qml.math.allclose(
+            [jax.numpy.asarray(x) for x in obtained_gates[0].parameters], expected_params
+        )
diff --git a/tests/transforms/test_unitary_to_rot.py b/tests/transforms/test_unitary_to_rot.py
new file mode 100644
index 000000000..72de3a5af
--- /dev/null
+++ b/tests/transforms/test_unitary_to_rot.py
@@ -0,0 +1,336 @@
+# Copyright 2018-2021 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Tests for the QubitUnitary decomposition transforms.
+"""
+
+import pytest
+
+from itertools import product
+
+import pennylane as qml
+from pennylane import numpy as np
+
+from pennylane.wires import Wires
+from pennylane.transforms import unitary_to_rot
+
+from gate_data import I, Z, S, T, H, X
+
+single_qubit_decomps = [
+    # First set of gates are diagonal and converted to RZ
+    (I, qml.RZ, [0.0]),
+    (Z, qml.RZ, [np.pi]),
+    (S, qml.RZ, [np.pi / 2]),
+    (T, qml.RZ, [np.pi / 4]),
+    (qml.RZ(0.3, wires=0).matrix, qml.RZ, [0.3]),
+    (qml.RZ(-0.5, wires=0).matrix, qml.RZ, [-0.5]),
+    # Next set of gates are non-diagonal and decomposed as Rots
+    (H, qml.Rot, [np.pi, np.pi / 2, 0.0]),
+    (X, qml.Rot, [0.0, np.pi, np.pi]),
+    (qml.Rot(0.2, 0.5, -0.3, wires=0).matrix, qml.Rot, [0.2, 0.5, -0.3]),
+    (np.exp(1j * 0.02) * qml.Rot(-1.0, 2.0, -3.0, wires=0).matrix, qml.Rot, [-1.0, 2.0, -3.0]),
+]
+
+# A simple quantum function for testing
+def qfunc(U):
+    qml.Hadamard(wires="a")
+    qml.QubitUnitary(U, wires="a")
+    qml.CNOT(wires=["b", "a"])
+
+
+class TestDecomposeSingleQubitUnitaryTransform:
+    """Tests to ensure the transform itself works in all interfaces."""
+
+    @pytest.mark.parametrize("U,expected_gate,expected_params", single_qubit_decomps)
+    def test_unitary_to_rot(self, U, expected_gate, expected_params):
+        """Test that the transform works in the autograd interface."""
+        transformed_qfunc = unitary_to_rot(qfunc)
+
+        ops = qml.transforms.make_tape(transformed_qfunc)(U).operations
+
+        assert len(ops) == 3
+
+        assert isinstance(ops[0], qml.Hadamard)
+        assert ops[0].wires == Wires("a")
+
+        assert isinstance(ops[1], expected_gate)
+        assert ops[1].wires == Wires("a")
+        assert qml.math.allclose(ops[1].parameters, expected_params)
+
+        assert isinstance(ops[2], qml.CNOT)
+        assert ops[2].wires == Wires(["b", "a"])
+
+    @pytest.mark.parametrize("U,expected_gate,expected_params", single_qubit_decomps)
+    def test_unitary_to_rot_torch(self, U, expected_gate, expected_params):
+        """Test that the transform works in the torch interface."""
+        torch = pytest.importorskip("torch")
+
+        U = torch.tensor(U, dtype=torch.complex64)
+
+        transformed_qfunc = unitary_to_rot(qfunc)
+
+        ops = qml.transforms.make_tape(transformed_qfunc)(U).operations
+
+        assert len(ops) == 3
+
+        assert isinstance(ops[0], qml.Hadamard)
+        assert ops[0].wires == Wires("a")
+
+        assert isinstance(ops[1], expected_gate)
+        assert ops[1].wires == Wires("a")
+        assert qml.math.allclose([x.detach() for x in ops[1].parameters], expected_params)
+
+        assert isinstance(ops[2], qml.CNOT)
+        assert ops[2].wires == Wires(["b", "a"])
+
+    @pytest.mark.parametrize("U,expected_gate,expected_params", single_qubit_decomps)
+    def test_unitary_to_rot_tf(self, U, expected_gate, expected_params):
+        """Test that the transform works in the Tensorflow interface."""
+        tf = pytest.importorskip("tensorflow")
+
+        U = tf.Variable(U, dtype=tf.complex64)
+
+        transformed_qfunc = unitary_to_rot(qfunc)
+
+        ops = qml.transforms.make_tape(transformed_qfunc)(U).operations
+
+        assert len(ops) == 3
+
+        assert isinstance(ops[0], qml.Hadamard)
+        assert ops[0].wires == Wires("a")
+
+        assert isinstance(ops[1], expected_gate)
+        assert ops[1].wires == Wires("a")
+        assert qml.math.allclose([x.numpy() for x in ops[1].parameters], expected_params)
+
+        assert isinstance(ops[2], qml.CNOT)
+        assert ops[2].wires == Wires(["b", "a"])
+
+    @pytest.mark.parametrize("U,expected_gate,expected_params", single_qubit_decomps)
+    def test_unitary_to_rot_jax(self, U, expected_gate, expected_params):
+        """Test that the transform works in the JAX interface."""
+        jax = pytest.importorskip("jax")
+
+        U = jax.numpy.array(U, dtype=jax.numpy.complex64)
+
+        transformed_qfunc = unitary_to_rot(qfunc)
+
+        ops = qml.transforms.make_tape(transformed_qfunc)(U).operations
+
+        assert len(ops) == 3
+
+        assert isinstance(ops[0], qml.Hadamard)
+        assert ops[0].wires == Wires("a")
+
+        assert isinstance(ops[1], expected_gate)
+        assert ops[1].wires == Wires("a")
+        assert qml.math.allclose([jax.numpy.asarray(x) for x in ops[1].parameters], expected_params)
+
+        assert isinstance(ops[2], qml.CNOT)
+        assert ops[2].wires == Wires(["b", "a"])
+
+
+# A simple circuit; we will test QubitUnitary on matrices constructed using trainable
+# parameters, and RZ/RX are easy to write the matrices for.
+def original_qfunc_for_grad(angles):
+    qml.Hadamard(wires="a")
+    qml.RZ(angles[0], wires="a")
+    qml.RX(angles[1], wires="b")
+    qml.CNOT(wires=["b", "a"])
+    return qml.expval(qml.PauliX(wires="a"))
+
+
+dev = qml.device("default.qubit", wires=["a", "b"])
+
+angle_pairs = [[0.3, 0.3], [np.pi, -0.65], [0.0, np.pi / 2], [np.pi / 3, 0.0]]
+diff_methods = ["parameter-shift", "backprop"]
+angle_diff_pairs = list(product(angle_pairs, diff_methods))
+
+
+class TestQubitUnitaryDifferentiability:
+    """Tests to ensure the transform is fully differentiable in all interfaces."""
+
+    @pytest.mark.parametrize("rot_angles,diff_method", angle_diff_pairs)
+    def test_gradient_unitary_to_rot(self, rot_angles, diff_method):
+        """Tests differentiability in autograd interface."""
+
+        def qfunc_with_qubit_unitary(angles):
+            z = angles[0]
+            x = angles[1]
+
+            Z_mat = np.array([[np.exp(-1j * z / 2), 0.0], [0.0, np.exp(1j * z / 2)]])
+
+            c = np.cos(x / 2)
+            s = np.sin(x / 2) * 1j
+            X_mat = np.array([[c, -s], [-s, c]])
+
+            qml.Hadamard(wires="a")
+            qml.QubitUnitary(Z_mat, wires="a")
+            qml.QubitUnitary(X_mat, wires="b")
+            qml.CNOT(wires=["b", "a"])
+            return qml.expval(qml.PauliX(wires="a"))
+
+        original_qnode = qml.QNode(original_qfunc_for_grad, dev, diff_method=diff_method)
+
+        transformed_qfunc = unitary_to_rot(qfunc_with_qubit_unitary)
+        transformed_qnode = qml.QNode(transformed_qfunc, dev, diff_method=diff_method)
+
+        input = np.array(rot_angles, requires_grad=True)
+        assert qml.math.allclose(original_qnode(input), transformed_qnode(input))
+
+        original_grad = qml.grad(original_qnode)(input)
+        transformed_grad = qml.grad(transformed_qnode)(input)
+
+        assert qml.math.allclose(original_grad, transformed_grad)
+
+    @pytest.mark.parametrize("rot_angles", angle_pairs)
+    def test_gradient_unitary_to_rot_torch(self, rot_angles):
+        """Tests differentiability in torch interface. Torch interface doesn't use
+        backprop so we test only with parameter-shift."""
+        torch = pytest.importorskip("torch", minversion="1.8")
+
+        def qfunc_with_qubit_unitary(angles):
+            z = angles[0]
+            x = angles[1]
+
+            # Had to do this in order to make a torch tensor of torch tensors
+            Z_mat = torch.stack(
+                [
+                    torch.exp(-1j * z / 2),
+                    torch.tensor(0.0),
+                    torch.tensor(0.0),
+                    torch.exp(1j * z / 2),
+                ]
+            ).reshape(2, 2)
+
+            # Variables need to be complex
+            c = torch.cos(x / 2).type(torch.complex64)
+            s = torch.sin(x / 2) * 1j
+
+            X_mat = torch.stack([c, -s, -s, c]).reshape(2, 2)
+
+            qml.Hadamard(wires="a")
+            qml.QubitUnitary(Z_mat, wires="a")
+            qml.QubitUnitary(X_mat, wires="b")
+            qml.CNOT(wires=["b", "a"])
+            return qml.expval(qml.PauliX(wires="a"))
+
+        original_qnode = qml.QNode(
+            original_qfunc_for_grad, dev, interface="torch", diff_method="parameter-shift"
+        )
+        original_input = torch.tensor(rot_angles, dtype=torch.float64, requires_grad=True)
+        original_result = original_qnode(original_input)
+
+        transformed_qfunc = unitary_to_rot(qfunc_with_qubit_unitary)
+        transformed_qnode = qml.QNode(
+            transformed_qfunc, dev, interface="torch", diff_method="parameter-shift"
+        )
+        transformed_input = torch.tensor(rot_angles, dtype=torch.float64, requires_grad=True)
+        transformed_result = transformed_qnode(transformed_input)
+
+        assert qml.math.allclose(original_result, transformed_result)
+
+        original_result.backward()
+        transformed_result.backward()
+
+        assert qml.math.allclose(original_input.grad, transformed_input.grad)
+
+    @pytest.mark.parametrize("rot_angles,diff_method", angle_diff_pairs)
+    def test_gradient_unitary_to_rot_tf(self, rot_angles, diff_method):
+        """Tests differentiability in tensorflow interface."""
+        tf = pytest.importorskip("tensorflow")
+
+        def qfunc_with_qubit_unitary(angles):
+            z = tf.cast(angles[0], tf.complex128)
+            x = tf.cast(angles[1], tf.complex128)
+
+            c = tf.cos(x / 2)
+            s = tf.sin(x / 2) * 1j
+
+            Z_mat = tf.convert_to_tensor([[tf.exp(-1j * z / 2), 0.0], [0.0, tf.exp(1j * z / 2)]])
+            X_mat = tf.convert_to_tensor([[c, -s], [-s, c]])
+
+            qml.Hadamard(wires="a")
+            qml.QubitUnitary(Z_mat, wires="a")
+            qml.QubitUnitary(X_mat, wires="b")
+            qml.CNOT(wires=["b", "a"])
+            return qml.expval(qml.PauliX(wires="a"))
+
+        original_qnode = qml.QNode(
+            original_qfunc_for_grad, dev, interface="tf", diff_method=diff_method
+        )
+        original_input = tf.Variable(rot_angles, dtype=tf.float64)
+        original_result = original_qnode(original_input)
+
+        transformed_qfunc = unitary_to_rot(qfunc_with_qubit_unitary)
+        transformed_qnode = qml.QNode(
+            transformed_qfunc, dev, interface="tf", diff_method=diff_method
+        )
+        transformed_input = tf.Variable(rot_angles, dtype=tf.float64)
+        transformed_result = transformed_qnode(transformed_input)
+
+        assert qml.math.allclose(original_result, transformed_result)
+
+        with tf.GradientTape() as tape:
+            loss = original_qnode(original_input)
+        original_grad = tape.gradient(loss, original_input)
+
+        with tf.GradientTape() as tape:
+            loss = transformed_qnode(transformed_input)
+
+        transformed_grad = tape.gradient(loss, transformed_input)
+
+        # For 64bit values, need to slightly increase the tolerance threshold
+        assert qml.math.allclose(original_grad, transformed_grad, atol=1e-7)
+
+    @pytest.mark.parametrize("rot_angles,diff_method", angle_diff_pairs)
+    def test_gradient_unitary_to_rot_jax(self, rot_angles, diff_method):
+        """Tests differentiability in jax interface."""
+        jax = pytest.importorskip("jax")
+        from jax import numpy as jnp
+
+        def qfunc_with_qubit_unitary(angles):
+            z = angles[0]
+            x = angles[1]
+
+            Z_mat = jnp.array([[jnp.exp(-1j * z / 2), 0.0], [0.0, jnp.exp(1j * z / 2)]])
+
+            c = jnp.cos(x / 2)
+            s = jnp.sin(x / 2) * 1j
+            X_mat = jnp.array([[c, -s], [-s, c]])
+
+            qml.Hadamard(wires="a")
+            qml.QubitUnitary(Z_mat, wires="a")
+            qml.QubitUnitary(X_mat, wires="b")
+            qml.CNOT(wires=["b", "a"])
+            return qml.expval(qml.PauliX(wires="a"))
+
+        # Setting the dtype to complex64 causes the gradients to be complex...
+        input = jnp.array(rot_angles, dtype=jnp.float64)
+
+        original_qnode = qml.QNode(
+            original_qfunc_for_grad, dev, interface="jax", diff_method=diff_method
+        )
+        original_result = original_qnode(input)
+
+        transformed_qfunc = unitary_to_rot(qfunc_with_qubit_unitary)
+        transformed_qnode = qml.QNode(
+            transformed_qfunc, dev, interface="jax", diff_method=diff_method
+        )
+        transformed_result = transformed_qnode(input)
+        assert qml.math.allclose(original_result, transformed_result)
+
+        original_grad = jax.grad(original_qnode)(input)
+        transformed_grad = jax.grad(transformed_qnode)(input)
+        assert qml.math.allclose(original_grad, transformed_grad, atol=1e-7)