Repository: PennyLaneAI/pennylane
Pull request #1894 (2021-11-12 17:03:43)

Make compilation transforms support `jax.jit`
**Context:** This is the first of a set of PRs which will adjust the compilation/optimization transforms to be jittable (or the equivalent) in the various interfaces.

**Description of the Change:** Removes conditional statements where possible in the optimization and compilation transforms  (e.g., `merge_rotations`, and `zyz_decomposition`), to enable `jax.jit` support for compiled QNodes.

**Benefits:** Nearly all compilation transforms of QNodes can now be jitted in JAX (missing is the two-qubit decomposition).  Furthermore, the removal of conditionals in the `zyz_decomposition` has led to some nice speedups of the single-qubit and two-qubit decomposition functions (roughly 30% in the two-qubit decomp, and up to ~75% in some of the single-qubit decomp cases).

**Possible Drawbacks:** The update to `zyz_decomposition` feels a little sketchy. It's also not as much faster as I would have hoped.

**Related GitHub Issues:** None

**Extended example**

Let's make a sizeable random circuit using JAX and try to compile it.

```python
import pennylane as qml
import jax
import jax.numpy as jnp
from jax.config import config
remember = config.read("jax_enable_x64")
config.update("jax_enable_x64", True)
from scipy.stats import unitary_group

n_wires = 5
dev = qml.device("default.qubit.jax", wires=n_wires)

key = jax.random.PRNGKey(42)

shape = qml.RandomLayers.shape(n_layers=7, n_rotations=7)
weights = jax.random.normal(key, shape)

random_U = [jnp.array(unitary_group.rvs(2), dtype=jnp.complex128) for _ in range(n_wires)]
random_U_2 = [jnp.array(unitary_group.rvs(2), dtype=jnp.complex128) for _ in range(n_wires)]
random_U_3 = [jnp.array(unitary_group.rvs(2), dtype=jnp.complex128) for _ in range(n_wires)]
```

Here is a circuit created using these random arrays:

```python
def my_circuit(weights):
    for idx, (U, U_2, U_3) in enumerate(zip(random_U, random_U_2, random_U_3)):
        qml.QubitUnitary(U, wires=idx)
        qml.QubitUnitary(U_2, wires=idx)
        qml.QubitUnitary(U_3, wires=idx)
                
    qml.RandomLayers(weights, wires=range(n_wires), seed=12)
    
    for idx, (U, U_2, U_3) in enumerate(zip(random_U, random_U_2, random_U_3)):
        qml.QubitUnitary(U, wires=idx)
        qml.QubitUnitary(U_2, wires=idx)
        qml.QubitUnitary(U_3, wires=idx)
        
    qml.RandomLayers(weights, wires=range(n_wires), seed=12)
    
    for idx, (U, U_2, U_3) in enumerate(zip(random_U, random_U_2, random_U_3)):
        qml.QubitUnitary(U, wires=idx)
        qml.QubitUnitary(U_2, wires=idx)
        qml.QubitUnitary(U_3, wires=idx)

    return qml.expval(qml.PauliZ(0))
```

We create 4 versions of this:  original, jitted original, compiled, and both jitted and compiled.
```python
# Original QNode
original_qnode = qml.QNode(my_circuit, dev, interface="jax")

# Jitted QNode
jitted_qnode = jax.jit(qml.QNode(my_circuit, dev, interface="jax"))

pipeline = [
    qml.transforms.unitary_to_rot,
    qml.transforms.commute_controlled(direction="right"),
    qml.transforms.single_qubit_fusion(),
]

# Compiled QNode
compiled_qfunc = qml.compile(pipeline=pipeline)(my_circuit)
compiled_qnode = qml.QNode(compiled_qfunc, dev, interface="jax")

# Jitted compiled QNode
jitted_compiled_qnode = jax.jit(compiled_qnode)
```

Timing information for the original version:

```
%%timeit
original_qnode(weights)
# 334 ms ± 6.92 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%%time
jitted_qnode(weights)
# CPU times: user 4.51 s, sys: 16.9 ms, total: 4.53 s
# Wall time: 4.51 s

%%timeit
jitted_qnode(weights)
# 39.6 µs ± 533 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
```

For the original QNode, the JIT takes a few seconds, but the new version is many orders of magnitude faster (8333x)

```
%%timeit
compiled_qnode(weights)
# 4.86 s ± 43.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%%time
jitted_compiled_qnode(weights)
# CPU times: user 45.1 s, sys: 470 ms, total: 45.6 s
# Wall time: 45.6 s

%%timeit
jitted_compiled_qnode(weights)
# 62.6 µs ± 1.17 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
```

The jitted, compiled version still takes about 2x as long on the simulator as the uncompiled jitted version; the jit time is also significantly longer. However, the jitted version (after jit), is 77635x faster than the unjitted version, so even with 45s overhead, we would absolutely want to do this if we were running the circuit multiple times.

Finally, the difference in gate count/depth makes it worth it as well:
```
>>> print(qml.specs(original_qnode)(weights))
{'gate_sizes': defaultdict(int, {1: 143, 2: 32}),
 'gate_types': defaultdict(int,
             {'QubitUnitary': 45, 'CNOT': 32, 'RX': 30, 'RZ': 36, 'RY': 32}),
 'num_operations': 175,
 'num_observables': 1,
 'num_diagonalizing_gates': 0,
 'num_used_wires': 5,
 'depth': 55,
 'num_trainable_params': 143,
 'num_device_wires': 5,
 'device_name': 'default.qubit.jax',
 'expansion_strategy': 'device',
 'gradient_options': {},
 'interface': 'jax',
 'diff_method': 'best',
 'gradient_fn': 'backprop'}
>>> print(qml.specs(compiled_qnode)(weights))
{'gate_sizes': defaultdict(int, {1: 41, 2: 32}),
 'gate_types': defaultdict(int, {'Rot': 41, 'CNOT': 32}),
 'num_operations': 73,
 'num_observables': 1,
 'num_diagonalizing_gates': 0,
 'num_used_wires': 5,
 'depth': 35,
 'num_trainable_params': 123,
 'num_device_wires': 5,
 'device_name': 'default.qubit.jax',
 'expansion_strategy': 'gradient',
 'gradient_options': {},
 'interface': 'jax',
 'diff_method': 'best',
 'gradient_fn': 'backprop'}
```

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.
josh146: @glassnotes I recently added [`qml.math.is_abstract`](https://pennylane.readthedocs.io/en/latest/code/api/pennylane.math.is_abstract.html), which allows you to check if a tensor is being traced (and therefore has no value!).

This might be helpful, e.g., to avoid branches where the tensor value needs to be compared to a hardcoded value.

However, it is only a 'stopgap' solution I think; a better solution would be to actually re-write those branches in such a way as to be compatible with JAX/TF etc.

For example, using [`tf.cond`](https://www.tensorflow.org/api_docs/python/tf/cond) or [`jax.lax.cond`](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.cond.html) instead of `if` statements. Perhaps we could even add a high-level `qml.math.cond()`?
glassnotes: > However, it is only a 'stopgap' solution I think; a better solution would be to actually re-write those branches in such a way as to be compatible with JAX/TF etc.

Yes agreed, let's see how far we can push this; the rotation angle checks are trivial to remove, and we might be able to swing it for the `zyz_decomposition`, however I'm thinking it's not possible for the two-qubit decomposition.
codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1894?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#1894](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1894?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (3b4d481) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/9cf401dd66bd385f91b633179a3b570d8da4dca6?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (9cf401d) will **decrease** coverage by `0.00%`.
> The diff coverage is `100.00%`.

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

```diff
@@            Coverage Diff             @@
##           master    #1894      +/-   ##
==========================================
- Coverage   98.83%   98.82%   -0.01%     
==========================================
  Files         222      222              
  Lines       16979    16961      -18     
==========================================
- Hits        16781    16762      -19     
- Misses        198      199       +1     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1894?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/1894/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvZGVjb21wb3NpdGlvbnMvc2luZ2xlX3F1Yml0X3VuaXRhcnkucHk=) | `100.00% <100.00%> (ø)` | |
| [...ane/transforms/decompositions/two\_qubit\_unitary.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1894/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvZGVjb21wb3NpdGlvbnMvdHdvX3F1Yml0X3VuaXRhcnkucHk=) | `100.00% <100.00%> (ø)` | |
| [...nnylane/transforms/optimization/merge\_rotations.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1894/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvb3B0aW1pemF0aW9uL21lcmdlX3JvdGF0aW9ucy5weQ==) | `100.00% <100.00%> (ø)` | |
| [...ane/transforms/optimization/single\_qubit\_fusion.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1894/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvb3B0aW1pemF0aW9uL3NpbmdsZV9xdWJpdF9mdXNpb24ucHk=) | `100.00% <100.00%> (ø)` | |
| [...lane/transforms/optimization/optimization\_utils.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1894/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvb3B0aW1pemF0aW9uL29wdGltaXphdGlvbl91dGlscy5weQ==) | `97.82% <0.00%> (-2.18%)` | :arrow_down: |

------

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

glassnotes: @josh146 this is ready for re-review. A few of the tests, however, are not passing - I have not been able to determine why as I can't reproduce the failures locally, all tests pass on my own machine(s).
josh146: > all tests pass on my own machine(s)

this is a sign there is a new package release that has broken something 😬😬😬😬 
josh146: Just checked - NumPy 1.22.2 recently came out. I upgraded locally, and those same tests fail 😬 
josh146: I've narrowed it down further --- with NumPy 1.22.2, the following line in the tests,

https://github.com/PennyLaneAI/pennylane/blob/bf1dac6549a8b396ce5438e60fe096bcd25f1a4d/tests/transforms/test_decompositions.py#L693-L699

is returning the following list of ops:

```
Rot(2.35619449, 3.14159257, -2.35619449, wires=[0])
Rot(-1.57079633, 1.57079633, -1.57079633, wires=[1])
CNOT(wires=[0, 1])
Rot(1.57079633, 1.57079633, 1.57079633, wires=[1])
Rot(0.78539816, nan, -0.78539816, wires=[0])
```

Note that `nan` in the last `Rot` 😬 
glassnotes: @josh146 incredible sleuthing :100: Weirdly, it still passed for me even after the upgrade. However, TIL that `math.clip` is differentiable in all frameworks, which entirely solves the problem and the tests are now passing :tada: 
josh146: 🎆 I'm just signing off for the night, but will give a review tomorrow!
glassnotes: @josh146 @maliasadi checks are now passing (minus the doc build, but that was a concurrency limit issue). Let me know if you want to review once more (I incorporated some of the additional changes you suggested Josh), otherwise I'll pull in from master and merge on Monday :+1: 
josh146: @glassnotes this should be good to go on my side!
glassnotes: Thanks @josh146 ! I think I will need a codecov override; it is complaining that line 153 of `merge_rotations.py` is not covered by tests; however this should be covered in `test_merge_rotations_jax_jit` at line 400 of `test_merge_rotations.py`.

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 25eff427e..177b16afc 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -297,6 +297,10 @@
 
 <h3>Improvements</h3>
 
+* Most compilation transforms, and relevant subroutines, have been updated to
+  support just-in-time compilation with `jax.jit`.
+  [(#1894)](https://github.com/PennyLaneAI/pennylane/pull/1894/)
+
 * No two-term parameter-shift rule is assumed anymore by default.
   [(#2227)](https://github.com/PennyLaneAI/pennylane/pull/2227)
 
@@ -538,6 +542,6 @@ The Operator class has undergone a major refactor with the following changes:
 
 This release contains contributions from (in alphabetical order):
 
-Sam Banning, Thomas Bromley, Anthony Hayes, David Ittah, Josh Izaac, Christina Lee, Angus Lowe,
+Sam Banning, Thomas Bromley, Olivia Di Matteo, Anthony Hayes, David Ittah, Josh Izaac, Christina Lee, Angus Lowe,
 Maria Fernanda Morris, Romain Moyard, Zeyue Niu, Maria Schuld, Jay Soni,
 Antal Száva, David Wierichs
diff --git a/pennylane/transforms/decompositions/single_qubit_unitary.py b/pennylane/transforms/decompositions/single_qubit_unitary.py
index 0ba5d8362..07d875a6e 100644
--- a/pennylane/transforms/decompositions/single_qubit_unitary.py
+++ b/pennylane/transforms/decompositions/single_qubit_unitary.py
@@ -14,13 +14,13 @@
 """Contains transforms and helpers functions for decomposing arbitrary unitary
 operations into elementary gates.
 """
-import numpy as np
+
 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.
+    r"""Convert a 2x2 unitary matrix to :math:`SU(2)`.
 
     Args:
         U (array[complex]): A matrix, presumed to be :math:`2 \times 2` and unitary.
@@ -29,29 +29,26 @@ def _convert_to_su2(U):
         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
+    exp_angle = -1j * math.cast_like(math.angle(det), 1j) / 2
+    return U * math.exp(exp_angle)
 
 
 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
+    Diagonal operations can 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)`.
 
+    .. warning::
+
+        When used with ``jax.jit``, all unitaries will be converted to :class:`.Rot` gates,
+        including those that are diagonal.
+
     Args:
         U (tensor): A 2 x 2 unitary matrix.
         wire (Union[Wires, Sequence[int] or int]): The wire on which to apply the operation.
@@ -81,27 +78,24 @@ def zyz_decomposition(U, wire):
     """
     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 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 = 0.0
-        theta = -np.pi
-        omega = 1j * math.log(U[0, 1] / U[1, 0]) - np.pi
-    else:
-        # 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))
-
-        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)]
+    # If the value of U is not abstract, we can include a conditional statement
+    # that will check if the off-diagonal elements are 0; if so, just use one RZ
+    if not math.is_abstract(U):
+        if math.allclose(U[0, 1], 0.0):
+            return [qml.RZ(2 * math.angle(U[1, 1]), wires=wire)]
+
+    # Derive theta from the off-diagonal element. Clip to ensure valid arcsin input
+    element = math.clip(math.abs(U[0, 1]), 0, 1)
+    theta = 2 * math.arcsin(element)
+
+    # Compute phi and omega from the angles of the top row; use atan2 to keep
+    # the angle within -np.pi and np.pi, and add very small values to avoid the
+    # undefined case of 0/0. We add a smaller value to the imaginary part than
+    # the real part because it is imag / real in the definition of atan2.
+    angle_U00 = math.arctan2(math.imag(U[0, 0]) + 1e-128, math.real(U[0, 0]) + 1e-64)
+    angle_U10 = math.arctan2(math.imag(U[1, 0]) + 1e-128, math.real(U[1, 0]) + 1e-64)
+
+    phi = -angle_U10 - angle_U00
+    omega = angle_U10 - angle_U00
+
+    return [qml.Rot(phi, theta, omega, wires=wire)]
diff --git a/pennylane/transforms/decompositions/two_qubit_unitary.py b/pennylane/transforms/decompositions/two_qubit_unitary.py
index 41dd10258..0c01fc28b 100644
--- a/pennylane/transforms/decompositions/two_qubit_unitary.py
+++ b/pennylane/transforms/decompositions/two_qubit_unitary.py
@@ -85,7 +85,7 @@ q_one_cnot = (1 / np.sqrt(2)) * np.array(
 
 
 def _convert_to_su4(U):
-    r"""Check unitarity of a 4x4 matrix and convert it to :math:`SU(4)` if the determinant is not 1.
+    r"""Convert a 4x4 matrix to :math:`SU(4)`.
 
     Args:
         U (array[complex]): A matrix, presumed to be :math:`4 \times 4` and unitary.
@@ -94,19 +94,11 @@ def _convert_to_su4(U):
         array[complex]: A :math:`4 \times 4` matrix in :math:`SU(4)` 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(4), atol=1e-7):
-        raise ValueError("Operator must be unitary.")
-
     # Compute the determinant
     det = math.linalg.det(U)
 
-    # Convert to SU(4) if it's not close to 1
-    if not math.allclose(det, 1.0):
-        exp_angle = -1j * math.cast_like(math.angle(det), 1j) / 4
-        U = math.cast_like(U, det) * math.exp(exp_angle)
-
-    return U
+    exp_angle = -1j * math.cast_like(math.angle(det), 1j) / 4
+    return math.cast_like(U, det) * math.exp(exp_angle)
 
 
 def _compute_num_cnots(U):
diff --git a/pennylane/transforms/optimization/merge_rotations.py b/pennylane/transforms/optimization/merge_rotations.py
index b0a8c2da0..6667cb265 100644
--- a/pennylane/transforms/optimization/merge_rotations.py
+++ b/pennylane/transforms/optimization/merge_rotations.py
@@ -12,10 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Transform for merging adjacent rotations of the same type in a quantum circuit."""
-
+# pylint: disable=too-many-branches
 from pennylane import apply
 from pennylane.transforms import qfunc_transform
-from pennylane.math import allclose, stack, cast_like, zeros
+from pennylane.math import allclose, stack, cast_like, zeros, is_abstract
 
 from pennylane.ops.qubit.attributes import composable_rotations
 from .optimization_utils import find_next_gate, fuse_rot_angles
@@ -146,8 +146,14 @@ def merge_rotations(tape, atol=1e-8, include_gates=None):
             # If we did merge, look now at the next gate
             next_gate_idx = find_next_gate(current_gate.wires, list_copy[1:])
 
-        if not allclose(cumulative_angles, zeros(len(cumulative_angles)), atol=atol, rtol=0):
+        # If we are tracing/jitting, don't perform any conditional checks and
+        # apply the operation regardless of the angles. Otherwise, only apply if
+        # the rotation angle is non-trivial.
+        if is_abstract(cumulative_angles):
             current_gate.__class__(*cumulative_angles, wires=current_gate.wires)
+        else:
+            if not allclose(cumulative_angles, zeros(len(cumulative_angles)), atol=atol, rtol=0):
+                current_gate.__class__(*cumulative_angles, wires=current_gate.wires)
 
         # Remove the first gate gate from the working list
         list_copy.pop(0)
diff --git a/pennylane/transforms/optimization/optimization_utils.py b/pennylane/transforms/optimization/optimization_utils.py
index 7b1e8b38c..32486cd07 100644
--- a/pennylane/transforms/optimization/optimization_utils.py
+++ b/pennylane/transforms/optimization/optimization_utils.py
@@ -12,8 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Utility functions for circuit optimization."""
-# pylint: disable=too-many-return-statements
-from pennylane.math import allclose, sin, cos, arccos, arctan2, zeros, cast_like, stack
+# pylint: disable=too-many-return-statements,import-outside-toplevel
+from pennylane.math import allclose, sin, cos, arccos, arctan2, stack, _multi_dispatch, is_abstract
 from pennylane.wires import Wires
 
 
@@ -40,51 +40,32 @@ def find_next_gate(wires, op_list):
     return next_gate_idx
 
 
-def _yzy_to_zyz(middle_yzy):
-    r"""Converts a set of angles representing a sequence of rotations RY, RZ, RY into
-    an equivalent sequence of the form RZ, RY, RZ.
+def _zyz_to_quat(angles):
+    """Converts a set of Euler angles in ZYZ format to a quaternion."""
+    qw = cos(angles[1] / 2) * cos(0.5 * (angles[0] + angles[2]))
+    qx = -sin(angles[1] / 2) * sin(0.5 * (angles[0] - angles[2]))
+    qy = sin(angles[1] / 2) * cos(0.5 * (angles[0] - angles[2]))
+    qz = cos(angles[1] / 2) * sin(0.5 * (angles[0] + angles[2]))
 
-    Any rotation in 3-dimensional space (or, equivalently, any single-qubit unitary)
-    can be expressed as a sequence of rotations about 3 axes in 12 different ways.
-    Typically, the arbitrary single-qubit rotation is expressed as RZ(a) RY(b) RZ(c),
-    but there are some situations, e.g., composing two such rotations, where we need
-    to convert between representations. This function converts the angles of a sequence
+    return stack([qw, qx, qy, qz])
 
-    .. math::
 
-       RY(y_1) RZ(z) RY(y_2)
+def _quaternion_product(q1, q2):
+    """Compute the product of two quaternions, q = q1 * q2."""
+    qw = q1[0] * q2[0] - q1[1] * q2[1] - q1[2] * q2[2] - q1[3] * q2[3]
+    qx = q1[0] * q2[1] + q1[1] * q2[0] + q1[2] * q2[3] - q1[3] * q2[2]
+    qy = q1[0] * q2[2] - q1[1] * q2[3] + q1[2] * q2[0] + q1[3] * q2[1]
+    qz = q1[0] * q2[3] + q1[1] * q2[2] - q1[2] * q2[1] + q1[3] * q2[0]
 
-    into the form
+    return stack([qw, qx, qy, qz])
 
-    .. math::
 
-       RZ(z_1) RY(y) RZ(z_2)
+def _fuse(angles_1, angles_2):
+    """Perform fusion of two angle sets. Separated out so we can do JIT with conditionals."""
+    # Compute the product of the quaternions
+    qw, qx, qy, qz = _quaternion_product(_zyz_to_quat(angles_1), _zyz_to_quat(angles_2))
 
-    This is accomplished by first converting the rotation to quaternion form, and then
-    extracting the desired set of angles.
-
-    Args:
-        y1 (float): The angle of the first ``RY`` rotation.
-        z (float): The angle of the inner ``RZ`` rotation.
-        y2 (float): The angle of the second ``RY`` rotation.
-
-    Returns:
-        tuple[float, float, float]: A list of rotation angles in the ZYZ representation.
-    """
-    if allclose(stack(middle_yzy), cast_like(zeros(3), stack(middle_yzy))):
-        return stack([0.0, 0.0, 0.0])
-
-    y1, z, y2 = middle_yzy[0], middle_yzy[1], middle_yzy[2]
-
-    # First, compute the quaternion representation
-    # https://ntrs.nasa.gov/api/citations/19770024290/downloads/19770024290.pdf
-    qw = cos(z / 2) * cos(0.5 * (y1 + y2))
-    qx = sin(z / 2) * sin(0.5 * (y1 - y2))
-    qy = cos(z / 2) * sin(0.5 * (y1 + y2))
-    qz = sin(z / 2) * cos(0.5 * (y1 - y2))
-
-    # Now convert from YZY Euler angles to ZYZ angles
-    # Source: http://bediyap.com/programming/convert-quaternion-to-euler-rotations/
+    # Convert the product back into the angles fed to Rot
     z1_arg1 = 2 * (qy * qz - qw * qx)
     z1_arg2 = 2 * (qx * qz + qw * qy)
     z1 = arctan2(z1_arg1, z1_arg2)
@@ -92,12 +73,20 @@ def _yzy_to_zyz(middle_yzy):
     y = arccos(qw**2 - qx**2 - qy**2 + qz**2)
 
     z2_arg1 = 2 * (qy * qz + qw * qx)
-    z2_arg2 = -2 * (qx * qz - qw * qy)
+    z2_arg2 = 2 * (qw * qy - qx * qz)
     z2 = arctan2(z2_arg1, z2_arg2)
 
     return stack([z1, y, z2])
 
 
+def _no_fuse(angles_1, angles_2):
+    """Special case: do not perform fusion when both Y angles are zero:
+        Rot(a, 0, b) Rot(c, 0, d) = Rot(a + b + c + d, 0, 0)
+    The quaternion math itself will fail in this case without a conditional.
+    """
+    return stack([angles_1[0] + angles_1[2] + angles_2[0] + angles_2[2], 0.0, 0.0])
+
+
 def fuse_rot_angles(angles_1, angles_2):
     """Computed the set of rotation angles that is obtained when composing
     two ``qml.Rot`` operations.
@@ -111,48 +100,30 @@ def fuse_rot_angles(angles_1, angles_2):
         angles_2 (float): A set of three angles for the second ``qml.Rot`` operation.
 
     Returns:
-        array[float]: A tuple of rotation angles for a single ``qml.Rot`` operation
-        that implements the same operation as the two sets of input angles.
+        array[float]: Rotation angles for a single ``qml.Rot`` operation that
+        implements the same operation as the two sets of input angles.
     """
-    # Check for all-zero instances; if there are some, we can just return the sum.
-    are_angles_1_zero = allclose(angles_1, zeros(3))
-    are_angles_2_zero = allclose(angles_2, zeros(3))
-
-    if are_angles_1_zero or are_angles_2_zero:
-        return stack([angles_1[i] + angles_2[i] for i in range(3)])
-
-    # RZ(a) RY(b) RZ(c) fused with RZ(d) RY(e) RZ(f)
-    # first produces RZ(a) RY(b) RZ(c+d) RY(e) RZ(f)
-    left_z = angles_1[0]
-    middle_yzy = angles_1[1], angles_1[2] + angles_2[0], angles_2[1]
-    right_z = angles_2[2]
-
-    # There are a few other cases to consider where things can be 0 and
-    # avoid having to use the quaternion conversion routine
-    # If b = 0, then we have RZ(a + c + d) RY(e) RZ(f)
-    if allclose(middle_yzy[0], 0.0):
-        # Then if e is close to zero, return a single rotation RZ(a + c + d + f)
-        if allclose(middle_yzy[2], 0.0):
-            return [left_z + middle_yzy[1] + right_z, 0.0, 0.0]
-        return stack([left_z + middle_yzy[1], middle_yzy[2], right_z])
-
-    # If c + d is close to 0, then we have the case RZ(a) RY(b + e) RZ(f)
-    if allclose(middle_yzy[1], 0.0):
-        # If b + e is 0, we have RZ(a + f)
-        if allclose(middle_yzy[0] + middle_yzy[2], 0.0):
-            return [left_z + right_z, 0.0, 0.0]
-        return stack([left_z, middle_yzy[0] + middle_yzy[2], right_z])
-
-    # If e is close to 0, then we have the case RZ(a) RY(b) RZ(c + d + f)
-    # The case where b is 0 actually already covered in the first loop,
-    # so only one case here.
-    if allclose(middle_yzy[2], 0.0):
-        return stack([left_z, middle_yzy[0], middle_yzy[1] + right_z])
-
-    # Otherwise, we need to turn the RY(b) RZ(c+d) RY(e) into something
-    # of the form RZ(u) RY(v) RZ(w)
-    u, v, w = _yzy_to_zyz(middle_yzy)
-
-    # Then we can combine to create
-    # RZ(a + u) RY(v) RZ(w + f)
-    return stack([left_z + u, v, w + right_z])
+
+    # Check if we are tracing; if so, use the special conditionals
+    if is_abstract(angles_1) or is_abstract(angles_2):
+        interface = _multi_dispatch([angles_1, angles_2])
+
+        # TODO: implement something similar for torch and tensorflow interfaces
+        # If the interface is JAX, use jax.lax.cond so that we can jit even with conditionals
+        if interface == "jax":
+            from jax.lax import cond
+
+            return cond(
+                allclose(angles_1[1], 0.0) * allclose(angles_2[1], 0.0),
+                _no_fuse,
+                _fuse,
+                angles_1,
+                angles_2,
+            )
+
+    # For other interfaces where we would not be jitting or tracing, we can simply check
+    # if we are dealing with the special case of Rot(a, 0, b) Rot(c, 0, d).
+    if allclose(angles_1[1], 0.0) and allclose(angles_2[1], 0.0):
+        return _no_fuse(angles_1, angles_2)
+
+    return _fuse(angles_1, angles_2)
diff --git a/pennylane/transforms/optimization/single_qubit_fusion.py b/pennylane/transforms/optimization/single_qubit_fusion.py
index de6482009..77d10b9b0 100644
--- a/pennylane/transforms/optimization/single_qubit_fusion.py
+++ b/pennylane/transforms/optimization/single_qubit_fusion.py
@@ -16,7 +16,7 @@
 from pennylane import apply
 from pennylane.transforms import qfunc_transform
 from pennylane.ops.qubit import Rot
-from pennylane.math import allclose, cast_like, stack, zeros
+from pennylane.math import allclose, stack, is_abstract
 
 from .optimization_utils import find_next_gate, fuse_rot_angles
 
@@ -127,20 +127,29 @@ def single_qubit_fusion(tape, atol=1e-8, exclude_gates=None):
             # the gate in question, only valid single-qubit gates on the same
             # wire as the current gate will be fused.
             try:
-                next_gate_angles = next_gate.single_qubit_rot_angles()
+                next_gate_angles = stack(next_gate.single_qubit_rot_angles())
             except (NotImplementedError, AttributeError):
                 break
 
-            cumulative_angles = fuse_rot_angles(
-                cumulative_angles, cast_like(stack(next_gate_angles), cumulative_angles)
-            )
+            cumulative_angles = fuse_rot_angles(cumulative_angles, stack(next_gate_angles))
 
             list_copy.pop(next_gate_idx + 1)
             next_gate_idx = find_next_gate(current_gate.wires, list_copy[1:])
 
-        # Only apply if the cumulative angle is not close to 0
-        if not allclose(cumulative_angles, zeros(3), atol=atol, rtol=0):
+        # If we are tracing/jitting, don't perform any conditional checks and
+        # apply the rotation regardless of the angles.
+        if is_abstract(cumulative_angles):
             Rot(*cumulative_angles, wires=current_gate.wires)
+        # If not tracing, check whether all angles are 0 (or equivalently, if the RY
+        # angle is close to 0, and so is the sum of the RZ angles
+        else:
+            if not allclose(
+                stack([cumulative_angles[0] + cumulative_angles[2], cumulative_angles[1]]),
+                [0.0, 0.0],
+                atol=atol,
+                rtol=0,
+            ):
+                Rot(*cumulative_angles, wires=current_gate.wires)
 
         # Remove the starting gate from the list
         list_copy.pop(0)


Test patch:
diff --git a/tests/ops/qubit/test_matrix_ops.py b/tests/ops/qubit/test_matrix_ops.py
index 2f96b0f12..ea4dac5f9 100644
--- a/tests/ops/qubit/test_matrix_ops.py
+++ b/tests/ops/qubit/test_matrix_ops.py
@@ -160,23 +160,31 @@ class TestQubitUnitary:
 
     @pytest.mark.parametrize(
         "U,expected_gate,expected_params",
-        [  # First set of gates are diagonal and converted to RZ
-            (I, qml.RZ, [0]),
+        [
+            (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).get_matrix(), qml.RZ, [0.3]),
-            (qml.RZ(-0.5, wires=0).get_matrix(), qml.RZ, [-0.5]),
-            # Next set of gates are non-diagonal and decomposed as Rots
+            (qml.RZ(0.3, wires=0).matrix, qml.RZ, [0.3]),
+            (qml.RZ(-0.5, wires=0).matrix, qml.RZ, [-0.5]),
+            (
+                np.array(
+                    [
+                        [0, -9.831019270939975e-01 + 0.1830590094588862j],
+                        [9.831019270939975e-01 + 0.1830590094588862j, 0],
+                    ]
+                ),
+                qml.Rot,
+                [-0.18409714468526372, np.pi, 0.18409714468526372],
+            ),
+            (H, qml.Rot, [np.pi, np.pi / 2, 0.0]),
+            (X, qml.Rot, [np.pi / 2, np.pi, -np.pi / 2]),
+            (qml.Rot(0.2, 0.5, -0.3, wires=0).matrix, qml.Rot, [0.2, 0.5, -0.3]),
             (
-                np.array([[0, -0.98310193 + 0.18305901j], [0.98310193 + 0.18305901j, 0]]),
+                np.exp(1j * 0.02) * qml.Rot(-1.0, 2.0, -3.0, wires=0).matrix,
                 qml.Rot,
-                [0, -np.pi, -5.914991017809059],
+                [-1.0, 2.0, -3.0],
             ),
-            (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).get_matrix(), qml.Rot, [0.2, 0.5, -0.3]),
-            (np.exp(1j * 0.02) * qml.Rot(-1, 2, -3, wires=0).get_matrix(), qml.Rot, [-1, 2, -3]),
         ],
     )
     def test_qubit_unitary_decomposition(self, U, expected_gate, expected_params):
@@ -186,9 +194,9 @@ class TestQubitUnitary:
 
         assert len(decomp) == 1 == len(decomp2)
         assert isinstance(decomp[0], expected_gate)
-        assert np.allclose(decomp[0].parameters, expected_params)
+        assert np.allclose(decomp[0].parameters, expected_params, atol=1e-7)
         assert isinstance(decomp2[0], expected_gate)
-        assert np.allclose(decomp2[0].parameters, expected_params)
+        assert np.allclose(decomp2[0].parameters, expected_params, atol=1e-7)
 
     def test_qubit_unitary_decomposition_multiqubit_invalid(self):
         """Test that QubitUnitary is not decomposed for more than two qubits."""
diff --git a/tests/transforms/test_compile.py b/tests/transforms/test_compile.py
index b3496b73a..831a6f1dd 100644
--- a/tests/transforms/test_compile.py
+++ b/tests/transforms/test_compile.py
@@ -20,7 +20,13 @@ import pennylane as qml
 from pennylane.wires import Wires
 
 from pennylane.transforms.compile import compile
-from pennylane.transforms.optimization import cancel_inverses, commute_controlled, merge_rotations
+from pennylane.transforms import unitary_to_rot
+from pennylane.transforms.optimization import (
+    cancel_inverses,
+    commute_controlled,
+    merge_rotations,
+    single_qubit_fusion,
+)
 
 from test_optimization.utils import compare_operation_lists
 
@@ -501,3 +507,45 @@ class TestCompileInterfaces:
         # Check operation list
         ops = transformed_qnode.qtape.operations
         compare_operation_lists(ops, expected_op_list, expected_wires_list)
+
+    @pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
+    def test_compile_jax_jit(self, diff_method):
+        """Test that compilation pipelines work with jax.jit, unitary_to_rot, and fusion."""
+        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=2)
+
+        def test_qfunc(x):
+            qml.Rot(x, x + 1, x + 2, wires=0)
+            qml.RX(2 * x, wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.QubitUnitary(qml.RX.compute_matrix(x), wires=1)
+            qml.QubitUnitary(qml.RZ.compute_matrix(x + 1), wires=1)
+            return qml.expval(qml.PauliX(0) @ qml.PauliZ(1))
+
+        original_qnode = qml.QNode(test_qfunc, dev, interface="jax", diff_method=diff_method)
+
+        pipeline = [cancel_inverses, unitary_to_rot, single_qubit_fusion]
+
+        compiled_qfunc = qml.compile(pipeline=pipeline)(test_qfunc)
+        compiled_qnode = qml.QNode(compiled_qfunc, dev, interface="jax", diff_method=diff_method)
+
+        jitted_compiled_qnode = jax.jit(compiled_qnode)
+
+        x = jnp.array(0.1, dtype=jnp.float64)
+
+        # Check that the numerical output is the same
+        assert qml.math.allclose(original_qnode(x), jitted_compiled_qnode(x))
+
+        # Check that the gradient is the same
+        assert qml.math.allclose(
+            jax.grad(original_qnode)(x),
+            jax.grad(jitted_compiled_qnode)(x),
+            atol=1e-7,
+        )
diff --git a/tests/transforms/test_decompositions.py b/tests/transforms/test_decompositions.py
index 0e8b01861..b9f623eee 100644
--- a/tests/transforms/test_decompositions.py
+++ b/tests/transforms/test_decompositions.py
@@ -42,15 +42,19 @@ single_qubit_decomps = [
     (T, qml.RZ, [np.pi / 4]),
     (qml.RZ(0.3, wires=0).get_matrix(), qml.RZ, [0.3]),
     (qml.RZ(-0.5, wires=0).get_matrix(), qml.RZ, [-0.5]),
-    # This is an off-diagonal gate that gets converted to Rot, but one RZ angle is 0
+    # Next set of gates are non-diagonal and decomposed as Rots
     (
-        np.array([[0, -0.98310193 + 0.18305901j], [0.98310193 + 0.18305901j, 0]]),
+        np.array(
+            [
+                [0, -9.831019270939975e-01 + 0.1830590094588862j],
+                [9.831019270939975e-01 + 0.1830590094588862j, 0],
+            ]
+        ),
         qml.Rot,
-        [0, -np.pi, -5.914991017809059],
+        [-0.18409714468526372, np.pi, 0.18409714468526372],
     ),
-    # 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]),
+    (X, qml.Rot, [np.pi / 2, np.pi, -np.pi / 2]),
     (qml.Rot(0.2, 0.5, -0.3, wires=0).get_matrix(), qml.Rot, [0.2, 0.5, -0.3]),
     (
         np.exp(1j * 0.02) * qml.Rot(-1.0, 2.0, -3.0, wires=0).get_matrix(),
@@ -63,11 +67,6 @@ single_qubit_decomps = [
 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."""
@@ -76,14 +75,15 @@ class TestQubitUnitaryZYZDecomposition:
         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)
+
+        assert qml.math.allclose(obtained_gates[0].parameters, expected_params, atol=1e-7)
 
         if obtained_gates[0].num_params == 1:
             obtained_mat = qml.RZ(obtained_gates[0].parameters[0], wires=0).get_matrix()
         else:
             obtained_mat = qml.Rot(*obtained_gates[0].parameters, wires=0).get_matrix()
 
-        assert check_matrix_equivalence(obtained_mat, U)
+        assert check_matrix_equivalence(obtained_mat, U, atol=1e-7)
 
     @pytest.mark.parametrize("U,expected_gate,expected_params", single_qubit_decomps)
     def test_zyz_decomposition_torch(self, U, expected_gate, expected_params):
@@ -97,14 +97,17 @@ class TestQubitUnitaryZYZDecomposition:
         assert len(obtained_gates) == 1
         assert isinstance(obtained_gates[0], expected_gate)
         assert obtained_gates[0].wires == Wires("a")
-        assert qml.math.allclose(qml.math.unwrap(obtained_gates[0].parameters), expected_params)
+
+        assert qml.math.allclose(
+            qml.math.unwrap(obtained_gates[0].parameters), expected_params, atol=1e-7
+        )
 
         if obtained_gates[0].num_params == 1:
             obtained_mat = qml.RZ(obtained_gates[0].parameters[0], wires=0).get_matrix()
         else:
             obtained_mat = qml.Rot(*obtained_gates[0].parameters, wires=0).get_matrix()
 
-        assert check_matrix_equivalence(obtained_mat, qml.math.unwrap(U))
+        assert check_matrix_equivalence(obtained_mat, qml.math.unwrap(U), atol=1e-7)
 
     @pytest.mark.parametrize("U,expected_gate,expected_params", single_qubit_decomps)
     def test_zyz_decomposition_tf(self, U, expected_gate, expected_params):
@@ -118,6 +121,7 @@ class TestQubitUnitaryZYZDecomposition:
         assert len(obtained_gates) == 1
         assert isinstance(obtained_gates[0], expected_gate)
         assert obtained_gates[0].wires == Wires("a")
+
         assert qml.math.allclose(qml.math.unwrap(obtained_gates[0].parameters), expected_params)
 
         if obtained_gates[0].num_params == 1:
@@ -128,7 +132,7 @@ class TestQubitUnitaryZYZDecomposition:
                 *qml.math.unwrap(obtained_gates[0].parameters), wires=0
             ).get_matrix()
 
-        assert check_matrix_equivalence(obtained_mat, U)
+        assert check_matrix_equivalence(obtained_mat, U, atol=1e-7)
 
     @pytest.mark.parametrize("U,expected_gate,expected_params", single_qubit_decomps)
     def test_zyz_decomposition_jax(self, U, expected_gate, expected_params):
@@ -149,14 +153,16 @@ class TestQubitUnitaryZYZDecomposition:
         assert isinstance(obtained_gates[0], expected_gate)
         assert obtained_gates[0].wires == Wires("a")
 
-        assert qml.math.allclose(qml.math.unwrap(obtained_gates[0].parameters), expected_params)
+        assert qml.math.allclose(
+            qml.math.unwrap(obtained_gates[0].parameters), expected_params, atol=1e-7
+        )
 
         if obtained_gates[0].num_params == 1:
             obtained_mat = qml.RZ(obtained_gates[0].parameters[0], wires=0).get_matrix()
         else:
             obtained_mat = qml.Rot(*obtained_gates[0].parameters, wires=0).get_matrix()
 
-        assert check_matrix_equivalence(obtained_mat, U)
+        assert check_matrix_equivalence(obtained_mat, U, atol=1e-7)
 
 
 # Randomly generated set (scipy.unitary_group) of five U(4) operations.
@@ -173,132 +179,132 @@ samples_3_cnots = [
     ],
     [
         [
-            -0.07016275 - 0.11813399j,
-            -0.46476569 - 0.36887134j,
-            -0.18641714 + 0.66322739j,
-            -0.36131479 - 0.15452521j,
+            0.5636620105552962 + 0.0965426329099377j,
+            -0.5027332564280158 - 0.1635419007447941j,
+            -0.3847256335260943 + 0.3721938458489832j,
+            -0.0112168300627589 - 0.3268181271182378j,
         ],
         [
-            -0.21347395 - 0.47461873j,
-            0.45781338 - 0.21930172j,
-            0.16991759 - 0.12283128j,
-            -0.59626019 + 0.26831674j,
+            0.212792517319857 - 0.3244649017713778j,
+            0.1634380572825265 + 0.0960335485305775j,
+            -0.3139573850066476 + 0.2101838768898215j,
+            0.5975048356996334 + 0.5601329312675812j,
         ],
         [
-            -0.50100034 + 0.47448914j,
-            0.14346598 + 0.41837463j,
-            -0.39898589 + 0.01284601j,
-            -0.40027516 - 0.09308024j,
+            -0.0126840675187197 + 0.0099226785860024j,
+            0.400697758142422 - 0.5848686938298225j,
+            0.0812309251127118 + 0.561309695149665j,
+            -0.3604590752381415 + 0.2133372693485148j,
         ],
         [
-            0.47310527 + 0.10157557j,
-            -0.4411336 - 0.00466688j,
-            -0.2328747 - 0.51752603j,
-            -0.46274728 + 0.18717456j,
+            -0.238184683798628 + 0.6822134865739269j,
+            -0.0963245167970645 + 0.4155362603090703j,
+            0.0955731566481913 + 0.4898062484085562j,
+            0.1317657014485092 + 0.1717481576104292j,
         ],
     ],
     [
         [
-            -0.41189319 + 0.06007113j,
-            0.15316396 + 0.1458654j,
-            -0.17064243 + 0.33405919j,
-            0.60457809 + 0.52513855j,
+            0.5135439721958217 + 0.4814385791961366j,
+            -0.0469575920616757 - 0.5084648743399148j,
+            0.1487574414769012 - 0.2081434671214162j,
+            0.4170176952238572 + 0.0662699159748308j,
         ],
         [
-            0.08694733 + 0.64367692j,
-            -0.17413963 + 0.61038522j,
-            -0.27698195 - 0.31363701j,
-            -0.01169698 - 0.0012099j,
+            -0.1694469031236633 - 0.5709199999770778j,
+            0.1385353297280915 - 0.0224196352047464j,
+            0.3755479552341257 - 0.3411293791379575j,
+            0.4963367242427539 + 0.3491242747200203j,
         ],
         [
-            -0.39488729 - 0.41241129j,
-            0.2202738 + 0.19315024j,
-            -0.45149927 - 0.30746997j,
-            0.15876369 - 0.51435213j,
+            -0.1340989096401365 - 0.3533610802773053j,
+            0.120577125957884 - 0.8395701125658934j,
+            -0.0914053397143667 + 0.1816589401138052j,
+            -0.3090484419414308 - 0.0295087451853903j,
         ],
         [
-            -0.27620095 - 0.05049386j,
-            0.47854591 + 0.48737626j,
-            0.6201202 - 0.03549654j,
-            -0.25966725 + 0.03722163j,
+            -0.001618511029519 - 0.0835367795138499j,
+            0.0053302150713874 + 0.0108880702515039j,
+            -0.4304338310290907 + 0.6712032562215353j,
+            0.5750401423577074 + 0.1625231252605274j,
         ],
     ],
     [
         [
-            -0.34812515 + 0.37427723j,
-            -0.11092236 + 0.47565307j,
-            0.13724183 - 0.29504039j,
-            -0.56249794 - 0.27908375j,
+            -0.0031072463860293 - 0.6094283173470628j,
+            -0.1706041375326247 - 0.5095579204727588j,
+            0.4717848248508053 - 0.1360136245848086j,
+            -0.3088104665532103 + 0.0582056156280258j,
         ],
         [
-            -0.14408107 + 0.1693212j,
-            -0.20483797 - 0.10707915j,
-            -0.85376825 + 0.3175112j,
-            -0.1054503 - 0.23726165j,
+            0.0649795149056778 - 0.0807006955559929j,
+            0.3985010217077551 - 0.4878571486498078j,
+            -0.2593500452019464 - 0.369166183684014j,
+            0.5611176872453115 + 0.2721364042662399j,
         ],
         [
-            0.03106625 - 0.04236712j,
-            0.78292822 + 0.03053768j,
-            0.01814738 + 0.16830002j,
-            0.03513342 - 0.59451003j,
+            0.6706814505431677 - 0.3020432793307177j,
+            -0.445130459060895 + 0.2625311675303597j,
+            -0.1884345816738757 + 0.0153310963671089j,
+            0.1742601869852907 + 0.3546593618463598j,
         ],
         [
-            -0.00087219 - 0.82857442j,
-            0.04840206 + 0.30294214j,
-            -0.1884474 + 0.01468393j,
-            -0.41353861 + 0.11227088j,
+            -0.0330058354970538 + 0.2751715747427885j,
+            -0.2166842655029602 - 0.0205300200732348j,
+            0.6878324043735908 + 0.2156784765639485j,
+            0.5870150498641511 + 0.1077050421736986j,
         ],
     ],
     [
         [
-            -0.05780187 - 0.06284269j,
-            0.13559069 + 0.19399748j,
-            0.12381697 + 0.01612151j,
-            0.71416466 - 0.64114599j,
+            -0.1517007161518054 - 0.3477896338392819j,
+            -0.1996932364673036 - 0.3437390801195139j,
+            -0.5626622855975293 - 0.1026521198436937j,
+            0.5737100114799175 + 0.2042690304628801j,
         ],
         [
-            -0.31103029 - 0.06658675j,
-            -0.50183231 + 0.49812898j,
-            -0.58061141 - 0.20451914j,
-            -0.07379796 - 0.12030957j,
+            -0.1782045187028614 - 0.1798746573123555j,
+            0.6710156796470433 + 0.1482075706151732j,
+            -0.4090747335747313 + 0.4717438839366156j,
+            -0.2113406078464153 + 0.1706207743654816j,
         ],
         [
-            0.47241806 - 0.79298028j,
-            0.14041019 + 0.06342211j,
-            -0.27789855 + 0.19625469j,
-            -0.07716877 - 0.05067088j,
+            0.0355053593911394 - 0.6453216466524461j,
+            -0.1081082592116457 + 0.4107406869820365j,
+            0.4773230440109602 + 0.2293807227814043j,
+            0.2490211402335333 + 0.2438025872570191j,
         ],
         [
-            0.11114093 - 0.16488557j,
-            -0.12688073 + 0.63574829j,
-            0.68327072 - 0.15122624j,
-            -0.21697355 + 0.05813823j,
+            -0.5071192657577112 - 0.3421211957731553j,
+            -0.3794415508044536 - 0.2129960836701763j,
+            -0.0100377824272845 - 0.0493528476711767j,
+            -0.6523577435362742 + 0.0912843325184198j,
         ],
     ],
     [
         [
-            0.32457875 - 0.36309659j,
-            -0.21084067 + 0.48248995j,
-            -0.21588245 - 0.42368088j,
-            -0.0474262 + 0.50714809j,
+            -0.0196408936733692 + 0.0436166575003405j,
+            -0.4114171776549808 + 0.1604631422478851j,
+            0.8090489042142334 + 0.1854903095749424j,
+            -0.0285563913316359 - 0.336031332694153j,
         ],
         [
-            -0.24688996 - 0.11890225j,
-            0.16113004 + 0.1518989j,
-            -0.40132738 - 0.28678782j,
-            -0.63810805 - 0.4747406j,
+            0.6111824978145968 - 0.0373559219203255j,
+            -0.0607088813252194 - 0.4399072558564469j,
+            0.0652009339801155 - 0.3851081946614576j,
+            0.4498048829332472 - 0.2701344319737953j,
         ],
         [
-            -0.14474527 - 0.46561401j,
-            0.47151308 - 0.32560877j,
-            0.51600239 - 0.28298318j,
-            -0.18491473 + 0.23103107j,
+            -0.010769501670078 + 0.4210446448669616j,
+            -0.1955139271138671 + 0.5982440344653875j,
+            -0.3421820362428211 - 0.0107432223507508j,
+            0.4746580414982144 - 0.2897923140475613j,
         ],
         [
-            0.42624962 - 0.51795827j,
-            -0.17138618 - 0.56213399j,
-            -0.36133453 + 0.23168462j,
-            -0.0167845 - 0.14191731j,
+            0.598418329142766 + 0.2954204464071747j,
+            -0.4534575233858022 + 0.0777057134279078j,
+            -0.1459545057858948 + 0.1413048185423547j,
+            -0.4168424055846296 + 0.3576772276508062j,
         ],
     ],
 ]
@@ -311,82 +317,82 @@ samples_2_cnots = [
     # (A \otimes B) CNOT10 CNOT01 (C \otimes D)
     [
         [
-            0.06593717 - 0.26170549j,
-            -0.70347203 - 0.07117793j,
-            0.23473954 - 0.2369032j,
-            0.5244829 - 0.20227516j,
+            -0.004631217256064 - 0.268534699868393j,
+            -0.2794044288596784 - 0.2842241448919913j,
+            0.1218071742216182 - 0.7879690440506752j,
+            0.2778952482254939 + 0.2367688668871671j,
         ],
         [
-            -0.20854782 + 0.32788173j,
-            -0.43987873 + 0.2876307j,
-            -0.14557311 + 0.70166825j,
-            0.10554244 + 0.21933444j,
+            0.4392070528530516 + 0.6904136673657634j,
+            -0.1826413830164512 + 0.2570714399459114j,
+            -0.1294472995785671 - 0.1887672822622143j,
+            -0.0785276120371355 + 0.4152424069619657j,
         ],
         [
-            -0.10107 - 0.78553725j,
-            0.30719144 + 0.10962565j,
-            -0.08963713 + 0.42499178j,
-            0.27869654 - 0.00267013j,
+            0.1620021717921561 + 0.1884711066839195j,
+            0.4872686757195456 - 0.6397614642921917j,
+            -0.4411324161480765 - 0.1891338003177969j,
+            -0.2059859403259958 - 0.1367802403050884j,
         ],
         [
-            0.06589238 - 0.38018178j,
-            -0.33913438 + 0.04939232j,
-            0.4184943 + 0.10702358j,
-            -0.7203013 + 0.16805593j,
+            0.3022798492102187 - 0.3242755471407363j,
+            -0.1551409732131101 + 0.2662930443472138j,
+            -0.0292247033180649 - 0.2839975729154084j,
+            -0.7363616289920063 - 0.2911260927728274j,
         ],
     ],
     # CNOT01 (A \otimes B) CNOT01 (C \otimes D)
     [
         [
-            0.31547811 - 0.58699408j,
-            -0.34881989 - 0.41836625j,
-            0.30342551 - 0.04638021j,
-            -0.28264065 - 0.29172243j,
+            0.4952546422732242 - 0.3874950570842591j,
+            -0.565120684005787 - 0.0812920957189095j,
+            0.0599739067738637 - 0.3382684721070105j,
+            0.2002937293624679 - 0.3470743738556477j,
         ],
         [
-            -0.16129481 + 0.35079944j,
-            -0.03709334 - 0.15894938j,
-            0.0736614 + 0.55424822j,
-            -0.71185497 + 0.0702021j,
+            0.3076748055295139 + 0.1668987026621435j,
+            0.1113148363873911 - 0.355244639751341j,
+            -0.5552367793191325 + 0.2360956317049145j,
+            -0.4402702690966316 - 0.4254695664213973j,
         ],
         [
-            -0.30829764 + 0.23330248j,
-            -0.34470925 - 0.74564132j,
-            -0.21494836 - 0.09732758j,
-            0.28996211 + 0.18964071j,
+            -0.0570544619126448 - 0.5997099301388422j,
+            0.3690979468442732 - 0.489643527399396j,
+            -0.2915922328649661 - 0.1955865908229054j,
+            0.1605581200806802 + 0.3347381888568993j,
         ],
         [
-            -0.44032041 + 0.25194387j,
-            0.04287028 + 0.00320109j,
-            0.72290688 - 0.12204478j,
-            0.15915489 - 0.4218703j,
+            0.3161801218120408 - 0.1384715702727579j,
+            0.1624040016253136 + 0.3648031287938079j,
+            -0.2674750917123022 + 0.5684412664492574j,
+            0.5714177142050261 - 0.0145866404362501j,
         ],
     ],
     # (A \otimes B) CNOT10 (C \otimes D) CNOT01 (E \otimes F)
     [
         [
-            0.26680917 + 0.24320208j,
-            0.05925905 + 0.84940479j,
-            0.04305845 + 0.16580725j,
-            -0.03847889 + 0.33740004j,
+            -0.5032838832977308 + 0.008462334360567j,
+            -0.1994567215589622 + 0.4166138291684532j,
+            -0.2927933889967656 + 0.3307978265454508j,
+            0.0980847185817039 - 0.5731560630175854j,
         ],
         [
-            -0.05713611 - 0.69307062j,
-            -0.14217609 + 0.26062907j,
-            -0.12264675 - 0.59272591j,
-            -0.14185396 + 0.20434836j,
+            0.429674866329586 + 0.1517854863780976j,
+            -0.260061924207509 - 0.0913938324295813j,
+            0.4678474940708733 + 0.4691638242216676j,
+            0.4693562225846957 - 0.238881566880469j,
         ],
         [
-            0.18147763 + 0.27516979j,
-            -0.2868943 - 0.08518162j,
-            -0.09780205 - 0.39614393j,
-            0.76677616 + 0.21758277j,
+            -0.3407683276656533 - 0.1781416508017804j,
+            0.648130427328909 + 0.1856223456906182j,
+            0.3247150132557211 - 0.0284712498269712j,
+            0.5276154053566222 + 0.1139531710411058j,
         ],
         [
-            -0.52667927 - 0.00325491j,
-            -0.05056066 + 0.30779487j,
-            0.65170215 - 0.11435361j,
-            0.2628016 - 0.34416154j,
+            -0.6092485765643428 - 0.1411845180727529j,
+            -0.3732001418286821 - 0.3418447668630915j,
+            0.4811869468078607 + 0.1675038147841858j,
+            -0.1994479499053229 + 0.2310173648263156j,
         ],
     ],
 ]
@@ -398,163 +404,163 @@ samples_1_cnot = [
     # CNOT10 (A \otimes B)
     [
         [
-            -0.40013501 + 0.6346169j,
-            0.36588428 + 0.27955713j,
-            -0.34061004 - 0.21797615j,
-            -0.15200951 + 0.19619931j,
+            0.1386390404432302 - 0.334570112999545j,
+            0.0148174707054368 - 0.0412250336805805j,
+            -0.4992326027350552 + 0.7779384196647746j,
+            -0.0558235265127082 + 0.0968780293767465j,
         ],
         [
-            0.19619931 - 0.15200951j,
-            0.21797615 + 0.34061004j,
-            -0.27955713 - 0.36588428j,
-            0.6346169 - 0.40013501j,
+            0.1110962138948646 + 0.0126194273429225j,
+            -0.922378240060485 - 0.0603304126213831j,
+            0.0419677169072708 + 0.0125607156508953j,
+            -0.3515157103454069 - 0.0871472858321581j,
         ],
         [
-            -0.04618249 + 0.40174118j,
-            0.24303208 + 0.05036332j,
-            0.74472919 + 0.09069221j,
-            0.09646726 - 0.45024168j,
+            -0.7766806433421076 - 0.5011871453677099j,
+            -0.0967373552333893 - 0.0560669491829709j,
+            -0.2647978260425042 - 0.2470628571248854j,
+            -0.0334214548399949 - 0.0283207908867259j,
         ],
         [
-            0.45024168 - 0.09646726j,
-            0.09069221 + 0.74472919j,
-            0.05036332 + 0.24303208j,
-            -0.40174118 + 0.04618249j,
+            0.0029378421621844 + 0.0437084651240645j,
+            -0.0415219650227516 - 0.3597691905488813j,
+            0.0123400367724977 - 0.1111275941460545j,
+            -0.0580109080696657 + 0.922527024519649j,
         ],
     ],
     # CNOT01 (A \otimes B)
     [
         [
-            -0.65867779 - 0.39965996j,
-            0.10480488 - 0.11993597j,
-            0.33775094 + 0.50135151j,
-            -0.11410732 + 0.0509635j,
+            0.1275272250362362 - 0.3343853948894475j,
+            0.6337142336419559 - 0.1759440058732874j,
+            0.2947329435072354 - 0.116226264283067j,
+            0.4932993532409792 + 0.309278009290073j,
         ],
         [
-            -0.15620406 + 0.03112864j,
-            0.0184181 + 0.77022418j,
-            0.12151068 + 0.02920466j,
-            0.24022431 - 0.55472634j,
+            0.3999023544680882 - 0.5221380373177482j,
+            -0.3578174135830141 - 0.0065942382104973j,
+            0.5800639860980087 - 0.0502285871838717j,
+            -0.2092751655259081 - 0.2378654190750314j,
         ],
         [
-            0.0509635 - 0.11410732j,
-            -0.50135151 - 0.33775094j,
-            0.11993597 - 0.10480488j,
-            -0.39965996 - 0.65867779j,
+            -0.1565164704901467 + 0.5608027580207731j,
+            0.296176258834552 - 0.1124974497506721j,
+            0.5870241263928151 - 0.2965682015252186j,
+            -0.3194248225208941 - 0.1613832959085598j,
         ],
         [
-            0.55472634 - 0.24022431j,
-            0.02920466 + 0.12151068j,
-            0.77022418 + 0.0184181j,
-            -0.03112864 + 0.15620406j,
+            0.0054101937368828 + 0.3167756024402226j,
+            -0.4627149049700612 + 0.3534007012683281j,
+            0.2601293856027128 - 0.2457834172330788j,
+            0.647226333094554 + 0.1168250695361578j,
         ],
     ],
     # (A \otimes B) CNOT01
     [
         [
-            0.61975417 + 0.05871326j,
-            0.18490393 - 0.03060346j,
-            0.11206333 - 0.18820514j,
-            0.51965346 - 0.50922726j,
+            0.190458746413919 + 0.2870042726555157j,
+            0.502809168004646 + 0.6156507128238845j,
+            -0.3763228062935716 + 0.2616231647579805j,
+            -0.1734766279722236 + 0.096703042530068j,
         ],
         [
-            0.17568798 - 0.06526691j,
-            -0.50876306 + 0.35875148j,
-            -0.06315069 + 0.72482005j,
-            0.07415843 - 0.20610651j,
+            -0.7237411404174169 - 0.3286968531905625j,
+            0.3261596499092991 + 0.1107513854501445j,
+            -0.0781992451732389 + 0.1825664178909173j,
+            0.2212035729951513 - 0.4014156376735081j,
         ],
         [
-            0.72482005 - 0.06315069j,
-            0.20610651 - 0.07415843j,
-            -0.06526691 + 0.17568798j,
-            -0.35875148 + 0.50876306j,
+            -0.1821207882156272 + 0.0792315428784779j,
+            -0.4001569202892715 + 0.2234725352683322j,
+            0.3327888523333447 + 0.7218687134048248j,
+            0.1125960761054146 + 0.3255274339495012j,
         ],
         [
-            0.18820514 - 0.11206333j,
-            -0.50922726 + 0.51965346j,
-            -0.05871326 - 0.61975417j,
-            -0.03060346 + 0.18490393j,
+            0.2594885215726799 - 0.3777978851802577j,
+            -0.0957194014420285 + 0.1740213063816186j,
+            0.2880779045831403 + 0.1888308965690472j,
+            -0.6184873664151125 - 0.4993157690429775j,
         ],
     ],
     # (A \otimes B) CNOT10
     [
         [
-            0.19339868 + 0.08633274j,
-            0.01466384 + 0.70171345j,
-            -0.63869976 + 0.02205134j,
-            0.09769795 - 0.21108831j,
+            -0.0885673151087696 - 0.0243844138270481j,
+            -0.4871418040378778 + 0.4830303555130133j,
+            -0.7088634125078255 - 0.1034876301900189j,
+            -0.0696323339281462 + 0.0537594674970065j,
         ],
         [
-            0.18782762 - 0.1372004j,
-            -0.56109965 - 0.30592619j,
-            -0.32755953 + 0.62074274j,
-            0.12258652 + 0.17271055j,
+            0.0404303791439687 - 0.0781289108852639j,
+            0.3879597260952143 + 0.6022327441469053j,
+            0.2377071200057033 - 0.6435221726459008j,
+            0.0398405677965372 + 0.082773776650062j,
         ],
         [
-            0.30592619 + 0.56109965j,
-            0.1372004 - 0.18782762j,
-            0.17271055 + 0.12258652j,
-            0.62074274 - 0.32755953j,
+            -0.5889549838714252 + 0.4078345916760804j,
+            -0.0074407415198821 - 0.0876548780130169j,
+            0.0684934101349958 - 0.0612161883111329j,
+            -0.0267574702860782 + 0.6854994524959126j,
         ],
         [
-            0.70171345 + 0.01466384j,
-            0.08633274 + 0.19339868j,
-            0.21108831 - 0.09769795j,
-            -0.02205134 + 0.63869976j,
+            -0.2651157710067235 - 0.6327235490732244j,
+            -0.0911552999919856 - 0.011378937851836j,
+            0.0437473475396883 + 0.0763211101667442j,
+            0.6944805007772098 + 0.1757664963376742j,
         ],
     ],
     # (A \otimes B) CNOT10 (C \otimes D)
     [
         [
-            0.78385585 + 0.24581188j,
-            -0.08367428 + 0.34003966j,
-            -0.16227836 - 0.08719991j,
-            -0.222293 + 0.34520414j,
+            0.2211933286892437 + 0.0021586412960177j,
+            0.38444713219383 - 0.0581554847392023j,
+            -0.170894684604903 + 0.1031928176859452j,
+            -0.1036775219695835 - 0.8656121616196758j,
         ],
         [
-            -0.40100084 - 0.10151918j,
-            0.07269317 + 0.76372534j,
-            -0.26873591 + 0.2188187j,
-            -0.34389609 - 0.04434368j,
+            0.3442924845446644 + 0.260151125836039j,
+            -0.1914328218155123 - 0.0746552589054714j,
+            -0.5744657965166804 - 0.6284808877028109j,
+            -0.2036000102332332 + 0.071498266549486j,
         ],
         [
-            0.17914714 - 0.32406563j,
-            0.2474564 - 0.0645854j,
-            0.45026225 - 0.23004978j,
-            -0.66076185 - 0.32437111j,
+            -0.1048402095394854 + 0.5447532602544797j,
+            0.7011096884305793 - 0.024182937133914j,
+            -0.1624036601402791 + 0.1839401392571305j,
+            -0.1162941782466869 + 0.3554990108217186j,
         ],
         [
-            0.04747281 + 0.12113138j,
-            0.4725868 + 0.00816913j,
-            -0.53596566 - 0.55049352j,
-            0.1451332 - 0.38510071j,
+            -0.6743897607764879 + 0.0479978633805106j,
+            -0.167244245931061 - 0.5352458890411946j,
+            -0.4128361178667754 + 0.06711548036308j,
+            0.1819390859257292 - 0.1428141352633625j,
         ],
     ],
     # (A \otimes B) CNOT01 (C \otimes D)
     [
         [
-            0.38565489 - 0.50840093j,
-            -0.30818989 - 0.11768624j,
-            0.27524652 - 0.39173492j,
-            -0.5046883 - 0.00636884j,
+            -0.384175417575186 + 0.3923608184977355j,
+            0.1002767236561732 - 0.5802482358089508j,
+            -0.4052191818361661 + 0.4207032744597814j,
+            -0.1025417468296731 + 0.0031874535396264j,
         ],
         [
-            -0.05689053 - 0.29547371j,
-            0.51930406 - 0.09499063j,
-            -0.40250308 - 0.33321454j,
-            0.00587892 - 0.59806656j,
+            -0.4577221128473338 + 0.0586197038286386j,
+            -0.0937187062570979 + 0.6742419194459858j,
+            -0.0414783182526488 + 0.3776270025012934j,
+            -0.2437968467006254 + 0.3462792652304717j,
         ],
         [
-            0.56972698 - 0.14750915j,
-            -0.32434509 + 0.18005613j,
-            -0.62729996 + 0.30972331j,
-            0.15801266 - 0.04037939j,
+            0.3423554335462237 + 0.4267199288134111j,
+            0.1796472451791323 + 0.1419455069613175j,
+            -0.1490674637388668 - 0.2557411476563918j,
+            -0.7409292809972994 - 0.1080757005000481j,
         ],
         [
-            -0.37603427 + 0.12044912j,
-            -0.64676631 + 0.23778647j,
-            0.03588508 - 0.08379105j,
-            0.02300689 - 0.60033588j,
+            0.4279595694962425 - 0.0554066848293549j,
+            0.185704890878562 + 0.3208788487100644j,
+            -0.4049260830288736 + 0.5127316935400166j,
+            0.2406649865130532 - 0.4376672321481798j,
         ],
     ],
 ]
@@ -566,48 +572,38 @@ samples_su2_su2 = [
     (X, H),
     (
         [
-            [0.21993927 - 0.1111822j, -0.27174921 - 0.93027824j],
-            [0.27174921 - 0.93027824j, 0.21993927 + 0.1111822j],
+            [-0.8951484711495978 + 0.2618482674962029j, 0.1320728980584716 + 0.3357103647583892j],
+            [0.1145773085667637 + 0.3420770962361714j, -0.8540795475070558 + 0.3747044530324777j],
         ],
         [
-            [0.86361715 + 0.11195238j, 0.35098794 - 0.34415046j],
-            [-0.35098794 - 0.34415046j, 0.86361715 - 0.11195238j],
+            [0.535259309608487 - 0.7568444138362017j, -0.0856403406729011 - 0.3651708323125105j],
+            [0.2346548909477215 + 0.2926108112778731j, -0.9258821807212569 - 0.0453671985080405j],
         ],
     ),
     (
         [
-            [0.08652981 - 0.25976406j, 0.95775011 + 0.08803377j],
-            [-0.95775011 + 0.08803377j, 0.08652981 + 0.25976406j],
+            [-0.3475218186197753 + 0.616967631919303j, -0.6918447778561189 - 0.1411748210564089j],
+            [0.4110483325261496 + 0.5741243724791375j, 0.4216047723512865 - 0.5689199321376678j],
         ],
         [
-            [-0.1668335 + 0.92244968j, -0.34571207 + 0.04166941j],
-            [0.34571207 + 0.04166941j, -0.1668335 - 0.92244968j],
+            [-0.1036172252382817 - 0.409539663695392j, 0.0173544935458401 + 0.9062226856837956j],
+            [-0.3763539496942646 + 0.8245595424490042j, -0.2579549535044457 + 0.3345422357017644j],
         ],
     ),
     (
         [
-            [-0.36170896 - 0.92179965j, -0.1390035 + 0.01140442j],
-            [0.1390035 + 0.01140442j, -0.36170896 + 0.92179965j],
+            [-0.6199757150883975 - 0.4615217464182165j, -0.5797949307229534 + 0.257809287248686j],
+            [0.5072854247342099 - 0.381168319953196j, 0.1016602350822234 + 0.7661836635685451j],
         ],
         [
-            [-0.26587616 + 0.09917768j, 0.731434 + 0.62006287j],
-            [-0.731434 + 0.62006287j, -0.26587616 - 0.09917768j],
+            [-0.5741575760925136 - 0.4777766752909318j, -0.6560499370893111 + 0.1080324322072939j],
+            [0.3761599015364478 - 0.5482483514244275j, 0.1890587151173228 + 0.7226231907949495j],
         ],
     ),
     (
         [
-            [0.90778489 + 0.27105877j, 0.17141891 - 0.27031333j],
-            [-0.17141891 - 0.27031333j, 0.90778489 - 0.27105877j],
-        ],
-        [
-            [0.72380105 - 0.58082557j, -0.18751423 + 0.32185728j],
-            [0.18751423 + 0.32185728j, 0.72380105 + 0.58082557j],
-        ],
-    ),
-    (
-        [
-            [-0.11704772 + 0.39598718j, -0.89480902 - 0.16973743j],
-            [0.89480902 - 0.16973743j, -0.11704772 - 0.39598718j],
+            [0.257810302107603 + 0.3791450460627684j, 0.7491086034534051 - 0.478141383280444j],
+            [-0.6878868994174759 + 0.5626673047041733j, -0.2798521896663711 - 0.3631802166496543j],
         ],
         [[np.exp(-1j * np.pi / 3), 0], [0, np.exp(1j * np.pi / 3)]],
     ),
@@ -617,14 +613,14 @@ samples_su2_su2 = [
             [-1j, 0],
         ],
         [
-            [-0.05594177 - 0.02155518j, -0.83242096 - 0.55089131j],
-            [0.83242096 - 0.55089131j, -0.05594177 + 0.02155518j],
+            [-0.6429551068755086 - 0.2344852807169338j, 0.348170308379626 + 0.6406268961202303j],
+            [-0.0720224467655562 + 0.7255605769552991j, -0.5014071443907067 + 0.4657955472996208j],
         ],
     ),
     (
         [
-            [0.90778489 + 0.27105877j, 0.17141891 - 0.27031333j],
-            [-0.17141891 - 0.27031333j, 0.90778489 - 0.27105877j],
+            [0.7192114465877868 + 0.4233015158437373j, 0.5404227760496959 + 0.1072098172194949j],
+            [0.5300962072416202 + 0.1501623549680364j, -0.8330292881813453 + 0.0501147009427765j],
         ],
         [
             [0, -1j],
@@ -646,13 +642,6 @@ class TestTwoQubitUnitaryDecomposition:
         assert qml.math.isclose(qml.math.linalg.det(U_su4), 1.0)
         assert check_matrix_equivalence(np.array(U), U_su4)
 
-    def test_convert_to_su4_invalid(self):
-        """Test a matrix in U(4) is correct converted to SU(4)."""
-        U = np.ones((4, 4))
-
-        with pytest.raises(ValueError, match="Operator must be unitary"):
-            _convert_to_su4(U)
-
     @pytest.mark.parametrize("U_pair", samples_su2_su2)
     def test_su2su2_to_tensor_products(self, U_pair):
         """Test SU(2) x SU(2) can be correctly factored into tensor products."""
diff --git a/tests/transforms/test_optimization/test_merge_rotations.py b/tests/transforms/test_optimization/test_merge_rotations.py
index 8020bb643..fdeff77fc 100644
--- a/tests/transforms/test_optimization/test_merge_rotations.py
+++ b/tests/transforms/test_optimization/test_merge_rotations.py
@@ -103,9 +103,7 @@ class TestMergeRotations:
             (0.3, -0.3, 0.7, -0.1, [qml.RY(0.6, wires=1)]),
         ],
     )
-    def test_two_qubits_rotation_no_merge(
-        self, theta_11, theta_12, theta_21, theta_22, expected_ops
-    ):
+    def test_two_qubits_merge(self, theta_11, theta_12, theta_21, theta_22, expected_ops):
         """Test that a two-qubit circuit with rotations on different qubits get merged."""
 
         def qfunc():
@@ -398,3 +396,29 @@ class TestMergeRotationsInterfaces:
         # Check operation list
         ops = transformed_qnode.qtape.operations
         compare_operation_lists(ops, expected_op_list, expected_wires_list)
+
+    def test_merge_rotations_jax_jit(self):
+        """Test that when using jax.jit, the conditional statement that checks for
+        0 rotation angles does not break things."""
+
+        jax = pytest.importorskip("jax")
+
+        # Enable float64 support
+        from jax.config import config
+
+        remember = config.read("jax_enable_x64")
+        config.update("jax_enable_x64", True)
+
+        dev = qml.device("default.qubit", wires=["w1", "w2"])
+
+        @jax.jit
+        @qml.qnode(dev, interface="jax")
+        @merge_rotations()
+        def qfunc():
+            qml.CRX(0.2, wires=["w1", "w2"])
+            qml.CRX(-0.2, wires=["w1", "w2"])
+            return qml.expval(qml.PauliZ("w1"))
+
+        res = qfunc()
+
+        assert qml.math.allclose(res, [1.0])
diff --git a/tests/transforms/test_optimization/test_optimization_utils.py b/tests/transforms/test_optimization/test_optimization_utils.py
index 059753366..670f235db 100644
--- a/tests/transforms/test_optimization/test_optimization_utils.py
+++ b/tests/transforms/test_optimization/test_optimization_utils.py
@@ -17,10 +17,13 @@ import pytest
 import pennylane as qml
 from pennylane.transforms.optimization.optimization_utils import (
     find_next_gate,
-    _yzy_to_zyz,
+    _zyz_to_quat,
+    _quaternion_product,
     fuse_rot_angles,
 )
 
+from pennylane import numpy as np
+
 from utils import check_matrix_equivalence
 from pennylane.transforms.get_unitary_matrix import get_unitary_matrix
 
@@ -54,32 +57,39 @@ class TestRotGateFusion:
     """Test that utility functions for fusing two qml.Rot gates function as expected."""
 
     @pytest.mark.parametrize(
-        ("angles"),
-        [([0.15, 0.25, -0.90]), ([0.0, 0.0, 0.0]), ([0.15, 0.25, -0.90]), ([0.05, -1.34, 4.12])],
+        ("angles", "expected_quat"),
+        # Examples generated at https://www.mathworks.com/help/robotics/ref/eul2quat.html
+        [
+            (
+                [0.15, 0.25, -0.90],
+                [0.923247491800509, -0.062488597726915, 0.107884031713695, -0.363414748929363],
+            ),
+            ([np.pi / 2, 0.0, 0.0], [1 / np.sqrt(2), 0.0, 0.0, 1 / np.sqrt(2)]),
+            ([0.0, 0.0, 0.0], [1.0, 0.0, 0.0, 0.0]),
+            ([0.15, 0, -0.90], [0.930507621912314, 0.0, 0.0, -0.366272529086048]),
+            ([0.0, 0.2, 0.0], [0.995004165278026, 0.0, 0.099833416646828, 0.0]),
+        ],
     )
-    def test_yzy_to_zyz(self, angles):
-        """Test that a set of rotations of the form YZY is correctly converted
-        to a sequence of the form ZYZ."""
-
-        def original_ops():
-            qml.RY(angles[0], wires=0),
-            qml.RZ(angles[1], wires=0),
-            qml.RY(angles[2], wires=0),
-
-        compute_matrix = get_unitary_matrix(original_ops, [0])
-        product_yzy = compute_matrix()
+    def test_zyz_to_quat(self, angles, expected_quat):
+        """Test that ZYZ Euler angles are correctly converted to quaternions."""
+        obtained_quat = _zyz_to_quat(angles)
+        assert qml.math.allclose(obtained_quat, expected_quat)
 
-        z1, y, z2 = _yzy_to_zyz(angles)
-
-        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)
+    @pytest.mark.parametrize(
+        ("angles_1", "angles_2", "expected_quat"),
+        # Examples generated at https://www.vcalc.com/wiki/vCalc/Quaternion+Multiplication
+        [
+            ([0.0, 0.0, 0.0, 0.0], [0.1, 0.2, 0.3, 0.4], [0.0, 0.0, 0.0, 0.0]),
+            ([1.0, 0.0, 0.0, 0.0], [1.0, 0.0, 0.0, 0.0], [1.0, 0.0, 0.0, 0.0]),
+            ([1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0], [-60.0, 12.0, 30.0, 24.0]),
+            ([0.1, 0.0, -0.2, 0.15], [1.0, 0.05, 1.65, -0.25], [0.4675, -0.1925, -0.0275, 0.135]),
+            ([0.1, 0.0, 0.0, 0.15], [1.0, 0.0, 0.0, -0.25], [0.1375, 0.0, 0.0, 0.125]),
+        ],
+    )
+    def test_quaternion_product(self, angles_1, angles_2, expected_quat):
+        """Test that products of quaternions produce expected results."""
+        obtained_quat = _quaternion_product(angles_1, angles_2)
+        assert qml.math.allclose(obtained_quat, expected_quat)
 
     @pytest.mark.parametrize(
         ("angles_1", "angles_2"),
diff --git a/tests/transforms/test_optimization/test_single_qubit_fusion.py b/tests/transforms/test_optimization/test_single_qubit_fusion.py
index 868871412..c1761bbb6 100644
--- a/tests/transforms/test_optimization/test_single_qubit_fusion.py
+++ b/tests/transforms/test_optimization/test_single_qubit_fusion.py
@@ -71,7 +71,7 @@ class TestSingleQubitFusion:
             qml.RX(-0.2, wires=0)
             qml.RZ(-0.1, wires=0)
 
-        transformed_qfunc = single_qubit_fusion()(qfunc)
+        transformed_qfunc = single_qubit_fusion(atol=1e-7)(qfunc)
         transformed_ops = qml.transforms.make_tape(transformed_qfunc)().operations
         assert len(transformed_ops) == 0
 
@@ -296,3 +296,37 @@ class TestSingleQubitFusionInterfaces:
         # Check operation list
         ops = transformed_qnode.qtape.operations
         compare_operation_lists(ops, expected_op_list, expected_wires_list)
+
+    def test_single_qubit_fusion_jax_jit(self):
+        """Test QNode and gradient in JAX interface with JIT."""
+        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)
+
+        original_qnode = qml.QNode(qfunc, dev, interface="jax")
+        jitted_qnode = jax.jit(original_qnode)
+
+        transformed_qnode = qml.QNode(transformed_qfunc, dev, interface="jax")
+        jitted_transformed_qnode = jax.jit(transformed_qnode)
+
+        input = jnp.array([0.1, 0.2, 0.3, 0.4], dtype=jnp.float64)
+
+        # Check that the numerical output is the same
+        original_output = original_qnode(input)
+        assert qml.math.allclose(jitted_qnode(input), original_output)
+        assert qml.math.allclose(transformed_qnode(input), original_output)
+        assert qml.math.allclose(jitted_transformed_qnode(input), original_output)
+
+        # Check that the gradients are the same even after jitting
+        original_gradient = jax.grad(original_qnode)(input)
+        assert qml.math.allclose(jax.grad(jitted_qnode)(input), original_gradient)
+        assert qml.math.allclose(jax.grad(transformed_qnode)(input), original_gradient)
+        assert qml.math.allclose(jax.grad(jitted_transformed_qnode)(input), original_gradient)
+
+        # Check operation list
+        ops = transformed_qnode.qtape.operations
+        compare_operation_lists(ops, expected_op_list, expected_wires_list)
diff --git a/tests/transforms/test_unitary_to_rot.py b/tests/transforms/test_unitary_to_rot.py
index e07b10d97..ad68673fd 100644
--- a/tests/transforms/test_unitary_to_rot.py
+++ b/tests/transforms/test_unitary_to_rot.py
@@ -39,12 +39,17 @@ single_qubit_decomps = [
     (qml.RZ(-0.5, wires=0).get_matrix(), qml.RZ, [-0.5]),
     # Next set of gates are non-diagonal and decomposed as Rots
     (
-        np.array([[0, -0.98310193 + 0.18305901j], [0.98310193 + 0.18305901j, 0]]),
+        np.array(
+            [
+                [0, -9.831019270939975e-01 + 0.1830590094588862j],
+                [9.831019270939975e-01 + 0.1830590094588862j, 0],
+            ]
+        ),
         qml.Rot,
-        [0, -np.pi, -5.914991017809059],
+        [-0.18409714468526372, np.pi, 0.18409714468526372],
     ),
     (H, qml.Rot, [np.pi, np.pi / 2, 0.0]),
-    (X, qml.Rot, [0.0, -np.pi, -np.pi]),
+    (X, qml.Rot, [np.pi / 2, np.pi, -np.pi / 2]),
     (qml.Rot(0.2, 0.5, -0.3, wires=0).get_matrix(), qml.Rot, [0.2, 0.5, -0.3]),
     (
         np.exp(1j * 0.02) * qml.Rot(-1.0, 2.0, -3.0, wires=0).get_matrix(),
@@ -77,7 +82,7 @@ class TestDecomposeSingleQubitUnitaryTransform:
 
         assert isinstance(ops[1], expected_gate)
         assert ops[1].wires == Wires("a")
-        assert qml.math.allclose(ops[1].parameters, expected_params)
+        assert qml.math.allclose(ops[1].parameters, expected_params, atol=1e-7)
 
         assert isinstance(ops[2], qml.CNOT)
         assert ops[2].wires == Wires(["b", "a"])
@@ -109,7 +114,7 @@ class TestDecomposeSingleQubitUnitaryTransform:
         """Test that the transform works in the torch interface."""
         torch = pytest.importorskip("torch")
 
-        U = torch.tensor(U, dtype=torch.complex64)
+        U = torch.tensor(U, dtype=torch.complex128)
 
         transformed_qfunc = unitary_to_rot(qfunc)
 
@@ -122,7 +127,7 @@ class TestDecomposeSingleQubitUnitaryTransform:
 
         assert isinstance(ops[1], expected_gate)
         assert ops[1].wires == Wires("a")
-        assert qml.math.allclose(qml.math.unwrap(ops[1].parameters), expected_params)
+        assert qml.math.allclose(qml.math.unwrap(ops[1].parameters), expected_params, atol=1e-7)
 
         assert isinstance(ops[2], qml.CNOT)
         assert ops[2].wires == Wires(["b", "a"])
@@ -132,7 +137,7 @@ class TestDecomposeSingleQubitUnitaryTransform:
         """Test that the transform works in the Tensorflow interface."""
         tf = pytest.importorskip("tensorflow")
 
-        U = tf.Variable(U, dtype=tf.complex64)
+        U = tf.Variable(U, dtype=tf.complex128)
 
         transformed_qfunc = unitary_to_rot(qfunc)
 
@@ -145,7 +150,7 @@ class TestDecomposeSingleQubitUnitaryTransform:
 
         assert isinstance(ops[1], expected_gate)
         assert ops[1].wires == Wires("a")
-        assert qml.math.allclose(qml.math.unwrap(ops[1].parameters), expected_params)
+        assert qml.math.allclose(qml.math.unwrap(ops[1].parameters), expected_params, atol=1e-7)
 
         assert isinstance(ops[2], qml.CNOT)
         assert ops[2].wires == Wires(["b", "a"])
@@ -161,7 +166,7 @@ class TestDecomposeSingleQubitUnitaryTransform:
         remember = config.read("jax_enable_x64")
         config.update("jax_enable_x64", True)
 
-        U = jax.numpy.array(U, dtype=jax.numpy.complex64)
+        U = jax.numpy.array(U, dtype=jax.numpy.complex128)
 
         transformed_qfunc = unitary_to_rot(qfunc)
 
@@ -174,11 +179,40 @@ class TestDecomposeSingleQubitUnitaryTransform:
 
         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 qml.math.allclose(qml.math.unwrap(ops[1].parameters), expected_params, atol=1e-7)
 
         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_jit(self, U, expected_gate, expected_params):
+        """Test that the transform works in the JAX interface with JIT."""
+        jax = pytest.importorskip("jax")
+
+        # Enable float64 support
+        from jax.config import config
+
+        remember = config.read("jax_enable_x64")
+        config.update("jax_enable_x64", True)
+
+        U = jax.numpy.array(U, dtype=jax.numpy.complex128)
+
+        dev = qml.device("default.qubit", wires=["a", "b"])
+
+        def test_qfunc(U):
+            qfunc(U)
+            return qml.expval(qml.PauliX("a") @ qml.PauliZ("b"))
+
+        original_qnode = qml.QNode(test_qfunc, dev, interface="jax")
+        transformed_qnode = qml.QNode(unitary_to_rot(test_qfunc), dev, interface="jax")
+        jitted_qnode = jax.jit(transformed_qnode)
+
+        original_result = original_qnode(U)
+        transformed_result = transformed_qnode(U)
+        jitted_result = jitted_qnode(U)
+        assert np.allclose(transformed_result, original_result)
+        assert np.allclose(jitted_result, original_result)
+
 
 # A simple circuit; we will test QubitUnitary on matrices constructed using trainable
 # parameters, and RZ/RX are easy to write the matrices for.
@@ -227,12 +261,12 @@ class TestQubitUnitaryDifferentiability:
         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))
+        assert qml.math.allclose(original_qnode(input), transformed_qnode(input), atol=1e-7)
 
         original_grad = qml.grad(original_qnode)(input)
         transformed_grad = qml.grad(transformed_qnode)(input)
 
-        assert qml.math.allclose(original_grad, transformed_grad)
+        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_torch(self, rot_angles, diff_method):
@@ -280,12 +314,12 @@ class TestQubitUnitaryDifferentiability:
         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)
+        assert qml.math.allclose(original_result, transformed_result, atol=1e-7)
 
         original_result.backward()
         transformed_result.backward()
 
-        assert qml.math.allclose(original_input.grad, transformed_input.grad)
+        assert qml.math.allclose(original_input.grad, transformed_input.grad, atol=1e-7)
 
     @pytest.mark.parametrize("rot_angles,diff_method", angle_diff_pairs)
     def test_gradient_unitary_to_rot_tf(self, rot_angles, diff_method):
@@ -386,6 +420,12 @@ class TestQubitUnitaryDifferentiability:
         transformed_grad = jax.grad(transformed_qnode)(input)
         assert qml.math.allclose(original_grad, transformed_grad, atol=1e-7)
 
+        # Check that we can also JIT
+        grad_of_jit = jax.grad(jax.jit(transformed_qnode))(input)
+        jit_of_grad = jax.jit(jax.grad(transformed_qnode))(input)
+        assert np.allclose(original_grad, jit_of_grad, atol=1e-7)
+        assert np.allclose(original_grad, grad_of_jit, atol=1e-7)
+
 
 test_two_qubit_unitaries = [
     # A tensor product of two operations