Repository: PennyLaneAI/pennylane
Pull request #2287 (2022-03-07 14:23:19)

Work on consistency of `Operator`s
**Context:**
In the context of preparing a tool for users that allows to analyze operations for consistency, correctness and completeness, some minor inconsistencies were found in PennyLane's operations and observables.

**Description of the Change:**
- Very minor speedup of string comparisons in `PauliError`
- Formatting/Clean up of some error messages
- Addition of `num_params=0` to `WireCut`
- ~~Change `SnapShot.wires` from `[]` to `qml.operation.AllWires`~~
- Remove `MultiControlledX` from parametrized tests in `test_parametric_ops.py`
- Remove `qml.SWAP.basis="X"`

**Benefits:**
Code quality and consistency.

**Possible Drawbacks:**
N/A


github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
dwierichs: I suppose this would be ready to merge once CodeFactor converges on an opinion :)

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 7f3f2a216..e471892bb 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -298,9 +298,9 @@
 
 <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/)
+* Some error messages have been improved and operation properties made
+  more consistent.
+  [(#2287)](https://github.com/PennyLaneAI/pennylane/pull/2287/)
 
 * No two-term parameter-shift rule is assumed anymore by default.
   [(#2227)](https://github.com/PennyLaneAI/pennylane/pull/2227)
@@ -366,6 +366,10 @@
 
 <h3>Bug fixes</h3>
 
+* The `basis` property of `qml.SWAP` was set to `"X"`, which is incorrect; it is
+  now set to `None`.
+  [(#2287)](https://github.com/PennyLaneAI/pennylane/pull/2287/)
+
 * The `qml.RandomLayers` template now decomposes when the weights are a list of lists.
   [(#2266)](https://github.com/PennyLaneAI/pennylane/pull/2266/)
 
@@ -543,6 +547,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, Olivia Di Matteo, Anthony Hayes, David Ittah, Josh Izaac, Christina Lee, Angus Lowe,
+Sam Banning, Thomas Bromley, 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/ops/channel.py b/pennylane/ops/channel.py
index 4ebc87a91..37c551b6c 100644
--- a/pennylane/ops/channel.py
+++ b/pennylane/ops/channel.py
@@ -79,7 +79,7 @@ class AmplitudeDamping(Channel):
          array([[0., 0.70710678], [0., 0.]])]
         """
         if not 0.0 <= gamma <= 1.0:
-            raise ValueError("gamma must be between [0,1].")
+            raise ValueError("gamma must be in the interval [0,1].")
 
         K0 = np.diag([1, np.sqrt(1 - gamma)])
         K1 = np.sqrt(gamma) * np.array([[0, 1], [0, 0]])
@@ -160,10 +160,10 @@ class GeneralizedAmplitudeDamping(Channel):
          array([[0.        , 0.        ], [0.34641016, 0.        ]])]
         """
         if not 0.0 <= gamma <= 1.0:
-            raise ValueError("gamma must be between [0,1].")
+            raise ValueError("gamma must be in the interval [0,1].")
 
         if not 0.0 <= p <= 1.0:
-            raise ValueError("p must be between [0,1].")
+            raise ValueError("p must be in the interval [0,1].")
 
         K0 = np.sqrt(p) * np.diag([1, np.sqrt(1 - gamma)])
         K1 = np.sqrt(p) * np.sqrt(gamma) * np.array([[0, 1], [0, 0]])
@@ -227,7 +227,7 @@ class PhaseDamping(Channel):
          array([[0.        , 0.        ], [0.        , 0.70710678]])]
         """
         if not 0.0 <= gamma <= 1.0:
-            raise ValueError("gamma must be between [0,1].")
+            raise ValueError("gamma must be in the interval [0,1].")
 
         K0 = np.diag([1, np.sqrt(1 - gamma)])
         K1 = np.diag([0, np.sqrt(gamma)])
@@ -306,7 +306,7 @@ class DepolarizingChannel(Channel):
          array([[ 0.40824829,  0.        ], [ 0.        , -0.40824829]])]
         """
         if not 0.0 <= p <= 1.0:
-            raise ValueError("p must be between [0,1]")
+            raise ValueError("p must be in the interval [0,1]")
 
         K0 = np.sqrt(1 - p) * np.eye(2)
         K1 = np.sqrt(p / 3) * np.array([[0, 1], [1, 0]])
@@ -372,7 +372,7 @@ class BitFlip(Channel):
          array([[0.        , 0.70710678], [0.70710678, 0.        ]])]
         """
         if not 0.0 <= p <= 1.0:
-            raise ValueError("p must be between [0,1]")
+            raise ValueError("p must be in the interval [0,1]")
 
         K0 = np.sqrt(1 - p) * np.eye(2)
         K1 = np.sqrt(p) * np.array([[0, 1], [1, 0]])
@@ -459,13 +459,13 @@ class ResetError(Channel):
          array([[0.        , 0.        ], [0.        , 0.54772256]])]
         """
         if not 0.0 <= p_0 <= 1.0:
-            raise ValueError("p_0 must be between [0,1]")
+            raise ValueError("p_0 must be in the interval [0,1]")
 
         if not 0.0 <= p_1 <= 1.0:
-            raise ValueError("p_1 must be between [0,1]")
+            raise ValueError("p_1 must be in the interval [0,1]")
 
         if not 0.0 <= p_0 + p_1 <= 1.0:
-            raise ValueError("p_0 + p_1 must be between [0,1]")
+            raise ValueError("p_0 + p_1 must be in the interval [0,1]")
 
         K0 = np.sqrt(1 - p_0 - p_1) * np.eye(2)
         K1 = np.sqrt(p_0) * np.array([[1, 0], [0, 0]])
@@ -533,12 +533,12 @@ class PauliError(Channel):
         super().__init__(operators, p, wires=wires, do_queue=do_queue, id=id)
 
         # check if the specified operators are legal
-        if not all(c in "XYZ" for c in operators):
+        if not set(operators).issubset({"X", "Y", "Z"}):
             raise ValueError("The specified operators need to be either of 'X', 'Y' or 'Z'")
 
         # check if probabilities are legal
         if not 0.0 <= p <= 1.0:
-            raise ValueError("p must be between [0,1]")
+            raise ValueError("p must be in the interval [0,1]")
 
         # check if the number of operators matches the number of wires
         if len(self.wires) != len(operators):
@@ -644,7 +644,7 @@ class PhaseFlip(Channel):
          array([[ 0.70710678,  0.        ], [ 0.        , -0.70710678]])]
         """
         if not 0.0 <= p <= 1.0:
-            raise ValueError("p must be between [0,1]")
+            raise ValueError("p must be in the interval [0,1]")
 
         K0 = np.sqrt(1 - p) * np.eye(2)
         K1 = np.sqrt(p) * np.array([[1, 0], [0, -1]])
diff --git a/pennylane/ops/qubit/non_parametric_ops.py b/pennylane/ops/qubit/non_parametric_ops.py
index 357634eb3..47da41c3b 100644
--- a/pennylane/ops/qubit/non_parametric_ops.py
+++ b/pennylane/ops/qubit/non_parametric_ops.py
@@ -1140,8 +1140,6 @@ class SWAP(Operation):
     num_params = 0
     """int: Number of trainable parameters that the operator depends on."""
 
-    basis = "X"
-
     @staticmethod
     def compute_matrix():  # pylint: disable=arguments-differ
         r"""Representation of the operator as a canonical matrix in the computational basis (static method).
@@ -1758,14 +1756,17 @@ class MultiControlledX(Operation):
                 wires = Wires(wires[-1])
             else:
                 raise ValueError(
-                    "MultiControlledX: wrong number of wires. 1 wire given. Need at least 2."
+                    "MultiControlledX: wrong number of wires. "
+                    f"{len(wires)} wire(s) given. Need at least 2."
                 )
         else:
             wires = Wires(wires)
             control_wires = Wires(control_wires)
 
             warnings.warn(
-                "The control_wires keyword will be removed soon. Use wires = (control_wires,target_wire). See the documentation for more information.",
+                "The control_wires keyword will be removed soon. "
+                "Use wires = (control_wires, target_wire) instead. "
+                "See the documentation for more information.",
                 category=UserWarning,
             )
 
@@ -1773,10 +1774,9 @@ class MultiControlledX(Operation):
                 raise ValueError("MultiControlledX accepts a single target wire.")
 
         work_wires = Wires([]) if work_wires is None else Wires(work_wires)
+        total_wires = control_wires + wires
 
-        if Wires.shared_wires([wires, work_wires]) or Wires.shared_wires(
-            [control_wires, work_wires]
-        ):
+        if Wires.shared_wires([total_wires, work_wires]):
             raise ValueError("The work wires must be different from the control and target wires")
 
         if not control_values:
@@ -1786,8 +1786,6 @@ class MultiControlledX(Operation):
         self.hyperparameters["work_wires"] = work_wires
         self.hyperparameters["control_values"] = control_values
 
-        total_wires = control_wires + wires
-
         super().__init__(wires=total_wires, do_queue=do_queue)
 
     def label(self, decimals=None, base_label=None, cache=None):
@@ -1834,12 +1832,12 @@ class MultiControlledX(Operation):
                 raise ValueError("Length of control bit string must equal number of control wires.")
 
             # Make sure all values are either 0 or 1
-            if any(x not in ["0", "1"] for x in control_values):
+            if not set(control_values).issubset({"1", "0"}):
                 raise ValueError("String of control values can contain only '0' or '1'.")
 
             control_int = int(control_values, 2)
         else:
-            raise ValueError("Alternative control values must be passed as a binary string.")
+            raise ValueError("Control values must be passed as a string.")
 
         padding_left = control_int * 2
         padding_right = 2 ** (len(control_wires) + 1) - 2 - padding_left
@@ -1852,8 +1850,7 @@ class MultiControlledX(Operation):
 
     def adjoint(self):
         return MultiControlledX(
-            control_wires=self.wires[:-1],
-            wires=self.wires[-1],
+            wires=self.wires,
             control_values=self.hyperparameters["control_values"],
         )
 
@@ -2069,6 +2066,7 @@ class WireCut(Operation):
     Args:
         wires (Sequence[int] or int): the wires the operation acts on
     """
+    num_params = 0
     num_wires = AnyWires
     grad_method = None
 
diff --git a/pennylane/ops/qubit/parametric_ops.py b/pennylane/ops/qubit/parametric_ops.py
index 231d00dd7..cbbbad0c1 100644
--- a/pennylane/ops/qubit/parametric_ops.py
+++ b/pennylane/ops/qubit/parametric_ops.py
@@ -770,11 +770,11 @@ class MultiRZ(Operation):
 
     def __init__(self, theta, wires=None, do_queue=True, id=None):
         wires = Wires(wires)
-        self.hyperparameters["n_wires"] = len(wires)
+        self.hyperparameters["num_wires"] = len(wires)
         super().__init__(theta, wires=wires, do_queue=do_queue, id=id)
 
     @staticmethod
-    def compute_matrix(theta, n_wires):  # pylint: disable=arguments-differ
+    def compute_matrix(theta, num_wires):  # pylint: disable=arguments-differ
         r"""Representation of the operator as a canonical matrix in the computational basis (static method).
 
         The canonical matrix is the textbook matrix representation that does not consider wires.
@@ -784,7 +784,7 @@ class MultiRZ(Operation):
 
         Args:
             theta (tensor_like or float): rotation angle
-            n_wires (int): number of wires the rotation acts on
+            num_wires (int): number of wires the rotation acts on
 
         Returns:
             tensor_like: canonical matrix
@@ -797,7 +797,7 @@ class MultiRZ(Operation):
                 [0.0000+0.0000j, 0.0000+0.0000j, 0.9988+0.0500j, 0.0000+0.0000j],
                 [0.0000+0.0000j, 0.0000+0.0000j, 0.0000+0.0000j, 0.9988-0.0500j]])
         """
-        eigs = qml.math.convert_like(pauli_eigs(n_wires), theta)
+        eigs = qml.math.convert_like(pauli_eigs(num_wires), theta)
 
         if qml.math.get_interface(theta) == "tensorflow":
             theta = qml.math.cast_like(theta, 1j)
@@ -810,7 +810,7 @@ class MultiRZ(Operation):
         return -0.5 * functools.reduce(matmul, [qml.PauliZ(w) for w in self.wires])
 
     @staticmethod
-    def compute_eigvals(theta, n_wires):  # pylint: disable=arguments-differ
+    def compute_eigvals(theta, num_wires):  # pylint: disable=arguments-differ
         r"""Eigenvalues of the operator in the computational basis (static method).
 
         If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U`,
@@ -827,7 +827,7 @@ class MultiRZ(Operation):
 
         Args:
             theta (tensor_like or float): rotation angle
-            n_wires (int): number of wires the rotation acts on
+            num_wires (int): number of wires the rotation acts on
 
         Returns:
             tensor_like: eigenvalues
@@ -838,7 +838,7 @@ class MultiRZ(Operation):
         tensor([0.9689-0.2474j, 0.9689+0.2474j, 0.9689+0.2474j, 0.9689-0.2474j,
                 0.9689+0.2474j, 0.9689-0.2474j, 0.9689-0.2474j, 0.9689+0.2474j])
         """
-        eigs = qml.math.convert_like(pauli_eigs(n_wires), theta)
+        eigs = qml.math.convert_like(pauli_eigs(num_wires), theta)
 
         if qml.math.get_interface(theta) == "tensorflow":
             theta = qml.math.cast_like(theta, 1j)
diff --git a/pennylane/transforms/decompositions/single_qubit_unitary.py b/pennylane/transforms/decompositions/single_qubit_unitary.py
index 07d875a6e..0ba5d8362 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"""Convert a 2x2 unitary matrix to :math:`SU(2)`.
+    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.
@@ -29,26 +29,29 @@ 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]
 
-    exp_angle = -1j * math.cast_like(math.angle(det), 1j) / 2
-    return U * math.exp(exp_angle)
+    # 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 can be converted to a single :class:`.RZ` gate, while non-diagonal
+    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)`.
 
-    .. 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.
@@ -78,24 +81,27 @@ def zyz_decomposition(U, wire):
     """
     U = _convert_to_su2(U)
 
-    # 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)]
+    # 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)]
diff --git a/pennylane/transforms/decompositions/two_qubit_unitary.py b/pennylane/transforms/decompositions/two_qubit_unitary.py
index 0c01fc28b..41dd10258 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"""Convert a 4x4 matrix to :math:`SU(4)`.
+    r"""Check unitarity of a 4x4 matrix and convert it to :math:`SU(4)` if the determinant is not 1.
 
     Args:
         U (array[complex]): A matrix, presumed to be :math:`4 \times 4` and unitary.
@@ -94,11 +94,19 @@ 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)
 
-    exp_angle = -1j * math.cast_like(math.angle(det), 1j) / 4
-    return math.cast_like(U, det) * math.exp(exp_angle)
+    # 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
 
 
 def _compute_num_cnots(U):
diff --git a/pennylane/transforms/optimization/merge_rotations.py b/pennylane/transforms/optimization/merge_rotations.py
index 6667cb265..b0a8c2da0 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, is_abstract
+from pennylane.math import allclose, stack, cast_like, zeros
 
 from pennylane.ops.qubit.attributes import composable_rotations
 from .optimization_utils import find_next_gate, fuse_rot_angles
@@ -146,14 +146,8 @@ 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 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):
+        if not allclose(cumulative_angles, zeros(len(cumulative_angles)), atol=atol, rtol=0):
             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 32486cd07..7b1e8b38c 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,import-outside-toplevel
-from pennylane.math import allclose, sin, cos, arccos, arctan2, stack, _multi_dispatch, is_abstract
+# pylint: disable=too-many-return-statements
+from pennylane.math import allclose, sin, cos, arccos, arctan2, zeros, cast_like, stack
 from pennylane.wires import Wires
 
 
@@ -40,32 +40,51 @@ def find_next_gate(wires, op_list):
     return next_gate_idx
 
 
-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]))
+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.
 
-    return stack([qw, qx, qy, qz])
+    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
 
+    .. math::
 
-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]
+       RY(y_1) RZ(z) RY(y_2)
 
-    return stack([qw, qx, qy, qz])
+    into the form
 
+    .. math::
 
-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))
+       RZ(z_1) RY(y) RZ(z_2)
 
-    # Convert the product back into the angles fed to Rot
+    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/
     z1_arg1 = 2 * (qy * qz - qw * qx)
     z1_arg2 = 2 * (qx * qz + qw * qy)
     z1 = arctan2(z1_arg1, z1_arg2)
@@ -73,20 +92,12 @@ def _fuse(angles_1, angles_2):
     y = arccos(qw**2 - qx**2 - qy**2 + qz**2)
 
     z2_arg1 = 2 * (qy * qz + qw * qx)
-    z2_arg2 = 2 * (qw * qy - qx * qz)
+    z2_arg2 = -2 * (qx * qz - qw * qy)
     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.
@@ -100,30 +111,48 @@ 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]: Rotation angles for a single ``qml.Rot`` operation that
-        implements the same operation as the two sets of input angles.
+        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.
     """
-
-    # 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)
+    # 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])
diff --git a/pennylane/transforms/optimization/single_qubit_fusion.py b/pennylane/transforms/optimization/single_qubit_fusion.py
index 77d10b9b0..de6482009 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, stack, is_abstract
+from pennylane.math import allclose, cast_like, stack, zeros
 
 from .optimization_utils import find_next_gate, fuse_rot_angles
 
@@ -127,29 +127,20 @@ 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 = stack(next_gate.single_qubit_rot_angles())
+                next_gate_angles = next_gate.single_qubit_rot_angles()
             except (NotImplementedError, AttributeError):
                 break
 
-            cumulative_angles = fuse_rot_angles(cumulative_angles, stack(next_gate_angles))
+            cumulative_angles = fuse_rot_angles(
+                cumulative_angles, cast_like(stack(next_gate_angles), cumulative_angles)
+            )
 
             list_copy.pop(next_gate_idx + 1)
             next_gate_idx = find_next_gate(current_gate.wires, list_copy[1:])
 
-        # If we are tracing/jitting, don't perform any conditional checks and
-        # apply the rotation regardless of the angles.
-        if is_abstract(cumulative_angles):
+        # Only apply if the cumulative angle is not close to 0
+        if not allclose(cumulative_angles, zeros(3), atol=atol, rtol=0):
             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 ea4dac5f9..2f96b0f12 100644
--- a/tests/ops/qubit/test_matrix_ops.py
+++ b/tests/ops/qubit/test_matrix_ops.py
@@ -160,31 +160,23 @@ class TestQubitUnitary:
 
     @pytest.mark.parametrize(
         "U,expected_gate,expected_params",
-        [
-            (I, qml.RZ, [0.0]),
+        [  # 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]),
-            (
-                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]),
+            (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
             (
-                np.exp(1j * 0.02) * qml.Rot(-1.0, 2.0, -3.0, wires=0).matrix,
+                np.array([[0, -0.98310193 + 0.18305901j], [0.98310193 + 0.18305901j, 0]]),
                 qml.Rot,
-                [-1.0, 2.0, -3.0],
+                [0, -np.pi, -5.914991017809059],
             ),
+            (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):
@@ -194,9 +186,9 @@ class TestQubitUnitary:
 
         assert len(decomp) == 1 == len(decomp2)
         assert isinstance(decomp[0], expected_gate)
-        assert np.allclose(decomp[0].parameters, expected_params, atol=1e-7)
+        assert np.allclose(decomp[0].parameters, expected_params)
         assert isinstance(decomp2[0], expected_gate)
-        assert np.allclose(decomp2[0].parameters, expected_params, atol=1e-7)
+        assert np.allclose(decomp2[0].parameters, expected_params)
 
     def test_qubit_unitary_decomposition_multiqubit_invalid(self):
         """Test that QubitUnitary is not decomposed for more than two qubits."""
diff --git a/tests/ops/qubit/test_non_parametric_ops.py b/tests/ops/qubit/test_non_parametric_ops.py
index e868e9ab6..8924018b8 100644
--- a/tests/ops/qubit/test_non_parametric_ops.py
+++ b/tests/ops/qubit/test_non_parametric_ops.py
@@ -613,18 +613,18 @@ class TestMultiControlledX:
                 None,
                 [0, 1, 2],
                 [0, 1],
-                "Alternative control values must be passed as a binary string.",
+                "Control values must be passed as a string.",
             ),
             (
                 None,
                 [1],
                 "1",
-                "MultiControlledX: wrong number of wires. 1 wire given. Need at least 2.",
+                r"MultiControlledX: wrong number of wires. 1 wire\(s\) given. Need at least 2.",
             ),
             ([0], None, "", "Must specify the wires where the operation acts on"),
             ([0, 1], 2, "ab", "String of control values can contain only '0' or '1'."),
             ([0, 1], 2, "011", "Length of control bit string must equal number of control wires."),
-            ([0, 1], 2, [0, 1], "Alternative control values must be passed as a binary string."),
+            ([0, 1], 2, [0, 1], "Control values must be passed as a string."),
             ([0, 1], [2, 3], "10", "MultiControlledX accepts a single target wire."),
         ],
     )
diff --git a/tests/ops/qubit/test_parametric_ops.py b/tests/ops/qubit/test_parametric_ops.py
index aaa34a774..3455a9294 100644
--- a/tests/ops/qubit/test_parametric_ops.py
+++ b/tests/ops/qubit/test_parametric_ops.py
@@ -47,8 +47,6 @@ PARAMETRIZED_OPERATIONS = [
     qml.QubitUnitary(np.eye(2) * 1j, wires=0),
     qml.DiagonalQubitUnitary(np.array([1.0, 1.0j]), wires=1),
     qml.ControlledQubitUnitary(np.eye(2) * 1j, wires=[0], control_wires=[2]),
-    qml.MultiControlledX(control_wires=[0, 1], wires=2, control_values="01"),
-    qml.MultiControlledX(wires=[0, 1, 2], control_values="01"),
     qml.SingleExcitation(0.123, wires=[0, 3]),
     qml.SingleExcitationPlus(0.123, wires=[0, 3]),
     qml.SingleExcitationMinus(0.123, wires=[0, 3]),
@@ -75,6 +73,7 @@ NON_PARAMETRIZED_OPERATIONS = [
     qml.PauliZ(wires=0),
     qml.PauliY(wires=0),
     qml.MultiControlledX(control_wires=[0, 1], wires=2, control_values="01"),
+    qml.MultiControlledX(wires=[0, 1, 2], control_values="01"),
     qml.QubitSum(wires=[0, 1, 2]),
 ]
 
diff --git a/tests/ops/qubit/test_sparse.py b/tests/ops/qubit/test_sparse.py
index 269023bfc..9c61bb9bf 100644
--- a/tests/ops/qubit/test_sparse.py
+++ b/tests/ops/qubit/test_sparse.py
@@ -59,7 +59,7 @@ class TestSparse:
 
     def test_label(self):
         """Test label method returns 𝓗"""
-        H = qml.SparseHamiltonian(coo_matrix((np.array([[1, 0], [-1.5, 0]]))), 1)
+        H = qml.SparseHamiltonian(coo_matrix(np.array([[1, 0], [-1.5, 0]])), 1)
         assert H.label() == "𝓗"
 
     @pytest.mark.parametrize("sparse_hamiltonian", SPARSE_HAMILTONIAN_TEST_DATA)
diff --git a/tests/ops/test_channel_ops.py b/tests/ops/test_channel_ops.py
index 22398b8f5..3947bbfb0 100644
--- a/tests/ops/test_channel_ops.py
+++ b/tests/ops/test_channel_ops.py
@@ -82,7 +82,7 @@ class TestAmplitudeDamping:
         assert np.allclose(op(0.1, wires=0).kraus_matrices(), expected, atol=tol, rtol=0)
 
     def test_gamma_invalid_parameter(self):
-        with pytest.raises(ValueError, match="gamma must be between"):
+        with pytest.raises(ValueError, match="gamma must be in the interval"):
             channel.AmplitudeDamping(1.5, wires=0).kraus_matrices()
 
 
@@ -110,11 +110,11 @@ class TestGeneralizedAmplitudeDamping:
         assert np.allclose(op(0.1, 0.5, wires=0).kraus_matrices()[3], expected_K3, atol=tol, rtol=0)
 
     def test_gamma_invalid_parameter(self):
-        with pytest.raises(ValueError, match="gamma must be between"):
+        with pytest.raises(ValueError, match="gamma must be in the interval"):
             channel.GeneralizedAmplitudeDamping(1.5, 0.0, wires=0).kraus_matrices()
 
     def test_p_invalid_parameter(self):
-        with pytest.raises(ValueError, match="p must be between"):
+        with pytest.raises(ValueError, match="p must be in the interval"):
             channel.GeneralizedAmplitudeDamping(0.0, 1.5, wires=0).kraus_matrices()
 
 
@@ -137,7 +137,7 @@ class TestPhaseDamping:
         assert np.allclose(op(0.1, wires=0).kraus_matrices(), expected, atol=tol, rtol=0)
 
     def test_gamma_invalid_parameter(self):
-        with pytest.raises(ValueError, match="gamma must be between"):
+        with pytest.raises(ValueError, match="gamma must be in the interval"):
             channel.PhaseDamping(1.5, wires=0).kraus_matrices()
 
 
@@ -174,7 +174,7 @@ class TestBitFlip:
         assert np.allclose(gradient, (-2 * np.cos(angle)))
 
     def test_p_invalid_parameter(self):
-        with pytest.raises(ValueError, match="p must be between"):
+        with pytest.raises(ValueError, match="p must be in the interval"):
             channel.BitFlip(1.5, wires=0).kraus_matrices()
 
 
@@ -211,7 +211,7 @@ class TestPhaseFlip:
         assert np.allclose(gradient, 0.0)
 
     def test_p_invalid_parameter(self):
-        with pytest.raises(ValueError, match="p must be between"):
+        with pytest.raises(ValueError, match="p must be in the interval"):
             channel.PhaseFlip(1.5, wires=0).kraus_matrices()
 
 
@@ -250,7 +250,7 @@ class TestDepolarizingChannel:
         assert np.allclose(gradient, -(4 / 3) * np.cos(angle))
 
     def test_p_invalid_parameter(self):
-        with pytest.raises(ValueError, match="p must be between"):
+        with pytest.raises(ValueError, match="p must be in the interval"):
             channel.DepolarizingChannel(1.5, wires=0).kraus_matrices()
 
 
@@ -278,15 +278,15 @@ class TestResetError:
         assert np.allclose(op(p_0, p_1, wires=0).kraus_matrices()[4], expected_K4, atol=tol, rtol=0)
 
     def test_p0_invalid_parameter(self):
-        with pytest.raises(ValueError, match="p_0 must be between"):
+        with pytest.raises(ValueError, match="p_0 must be in the interval"):
             channel.ResetError(1.5, 0.0, wires=0).kraus_matrices()
 
     def test_p1_invalid_parameter(self):
-        with pytest.raises(ValueError, match="p_1 must be between"):
+        with pytest.raises(ValueError, match="p_1 must be in the interval"):
             channel.ResetError(0.0, 1.5, wires=0).kraus_matrices()
 
     def test_p0_p1_sum_not_normalized(self):
-        with pytest.raises(ValueError, match="must be between"):
+        with pytest.raises(ValueError, match="must be in the interval"):
             channel.ResetError(1.0, 1.0, wires=0).kraus_matrices()
 
     @pytest.mark.parametrize("angle", np.linspace(0, 2 * np.pi, 7))
@@ -333,7 +333,7 @@ class TestPauliError:
     EXPECTED_ERRORS = [ValueError, ValueError, ValueError, WireError]
     EXPECTED_MESSAGES = [
         "The number of operators must match the number of wires",
-        "p must be between \\[0,1\\]",
+        "p must be in the interval \\[0,1\\]",
         "The specified operators need to be either of 'X', 'Y' or 'Z'",
         "Wires must be unique",
     ]
diff --git a/tests/transforms/test_compile.py b/tests/transforms/test_compile.py
index 831a6f1dd..b3496b73a 100644
--- a/tests/transforms/test_compile.py
+++ b/tests/transforms/test_compile.py
@@ -20,13 +20,7 @@ import pennylane as qml
 from pennylane.wires import Wires
 
 from pennylane.transforms.compile import compile
-from pennylane.transforms import unitary_to_rot
-from pennylane.transforms.optimization import (
-    cancel_inverses,
-    commute_controlled,
-    merge_rotations,
-    single_qubit_fusion,
-)
+from pennylane.transforms.optimization import cancel_inverses, commute_controlled, merge_rotations
 
 from test_optimization.utils import compare_operation_lists
 
@@ -507,45 +501,3 @@ 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 b9f623eee..0e8b01861 100644
--- a/tests/transforms/test_decompositions.py
+++ b/tests/transforms/test_decompositions.py
@@ -42,19 +42,15 @@ 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]),
-    # Next set of gates are non-diagonal and decomposed as Rots
+    # This is an off-diagonal gate that gets converted to Rot, but one RZ angle is 0
     (
-        np.array(
-            [
-                [0, -9.831019270939975e-01 + 0.1830590094588862j],
-                [9.831019270939975e-01 + 0.1830590094588862j, 0],
-            ]
-        ),
+        np.array([[0, -0.98310193 + 0.18305901j], [0.98310193 + 0.18305901j, 0]]),
         qml.Rot,
-        [-0.18409714468526372, np.pi, 0.18409714468526372],
+        [0, -np.pi, -5.914991017809059],
     ),
+    # Next set of gates are non-diagonal and decomposed as Rots
     (H, qml.Rot, [np.pi, np.pi / 2, 0.0]),
-    (X, qml.Rot, [np.pi / 2, np.pi, -np.pi / 2]),
+    (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.0, 2.0, -3.0, wires=0).get_matrix(),
@@ -67,6 +63,11 @@ 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."""
@@ -75,15 +76,14 @@ 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, atol=1e-7)
+        assert qml.math.allclose(obtained_gates[0].parameters, expected_params)
 
         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, atol=1e-7)
+        assert check_matrix_equivalence(obtained_mat, U)
 
     @pytest.mark.parametrize("U,expected_gate,expected_params", single_qubit_decomps)
     def test_zyz_decomposition_torch(self, U, expected_gate, expected_params):
@@ -97,17 +97,14 @@ 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, atol=1e-7
-        )
+        assert qml.math.allclose(qml.math.unwrap(obtained_gates[0].parameters), expected_params)
 
         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), atol=1e-7)
+        assert check_matrix_equivalence(obtained_mat, qml.math.unwrap(U))
 
     @pytest.mark.parametrize("U,expected_gate,expected_params", single_qubit_decomps)
     def test_zyz_decomposition_tf(self, U, expected_gate, expected_params):
@@ -121,7 +118,6 @@ 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:
@@ -132,7 +128,7 @@ class TestQubitUnitaryZYZDecomposition:
                 *qml.math.unwrap(obtained_gates[0].parameters), wires=0
             ).get_matrix()
 
-        assert check_matrix_equivalence(obtained_mat, U, atol=1e-7)
+        assert check_matrix_equivalence(obtained_mat, U)
 
     @pytest.mark.parametrize("U,expected_gate,expected_params", single_qubit_decomps)
     def test_zyz_decomposition_jax(self, U, expected_gate, expected_params):
@@ -153,16 +149,14 @@ 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, atol=1e-7
-        )
+        assert qml.math.allclose(qml.math.unwrap(obtained_gates[0].parameters), expected_params)
 
         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, atol=1e-7)
+        assert check_matrix_equivalence(obtained_mat, U)
 
 
 # Randomly generated set (scipy.unitary_group) of five U(4) operations.
@@ -179,132 +173,132 @@ samples_3_cnots = [
     ],
     [
         [
-            0.5636620105552962 + 0.0965426329099377j,
-            -0.5027332564280158 - 0.1635419007447941j,
-            -0.3847256335260943 + 0.3721938458489832j,
-            -0.0112168300627589 - 0.3268181271182378j,
+            -0.07016275 - 0.11813399j,
+            -0.46476569 - 0.36887134j,
+            -0.18641714 + 0.66322739j,
+            -0.36131479 - 0.15452521j,
         ],
         [
-            0.212792517319857 - 0.3244649017713778j,
-            0.1634380572825265 + 0.0960335485305775j,
-            -0.3139573850066476 + 0.2101838768898215j,
-            0.5975048356996334 + 0.5601329312675812j,
+            -0.21347395 - 0.47461873j,
+            0.45781338 - 0.21930172j,
+            0.16991759 - 0.12283128j,
+            -0.59626019 + 0.26831674j,
         ],
         [
-            -0.0126840675187197 + 0.0099226785860024j,
-            0.400697758142422 - 0.5848686938298225j,
-            0.0812309251127118 + 0.561309695149665j,
-            -0.3604590752381415 + 0.2133372693485148j,
+            -0.50100034 + 0.47448914j,
+            0.14346598 + 0.41837463j,
+            -0.39898589 + 0.01284601j,
+            -0.40027516 - 0.09308024j,
         ],
         [
-            -0.238184683798628 + 0.6822134865739269j,
-            -0.0963245167970645 + 0.4155362603090703j,
-            0.0955731566481913 + 0.4898062484085562j,
-            0.1317657014485092 + 0.1717481576104292j,
+            0.47310527 + 0.10157557j,
+            -0.4411336 - 0.00466688j,
+            -0.2328747 - 0.51752603j,
+            -0.46274728 + 0.18717456j,
         ],
     ],
     [
         [
-            0.5135439721958217 + 0.4814385791961366j,
-            -0.0469575920616757 - 0.5084648743399148j,
-            0.1487574414769012 - 0.2081434671214162j,
-            0.4170176952238572 + 0.0662699159748308j,
+            -0.41189319 + 0.06007113j,
+            0.15316396 + 0.1458654j,
+            -0.17064243 + 0.33405919j,
+            0.60457809 + 0.52513855j,
         ],
         [
-            -0.1694469031236633 - 0.5709199999770778j,
-            0.1385353297280915 - 0.0224196352047464j,
-            0.3755479552341257 - 0.3411293791379575j,
-            0.4963367242427539 + 0.3491242747200203j,
+            0.08694733 + 0.64367692j,
+            -0.17413963 + 0.61038522j,
+            -0.27698195 - 0.31363701j,
+            -0.01169698 - 0.0012099j,
         ],
         [
-            -0.1340989096401365 - 0.3533610802773053j,
-            0.120577125957884 - 0.8395701125658934j,
-            -0.0914053397143667 + 0.1816589401138052j,
-            -0.3090484419414308 - 0.0295087451853903j,
+            -0.39488729 - 0.41241129j,
+            0.2202738 + 0.19315024j,
+            -0.45149927 - 0.30746997j,
+            0.15876369 - 0.51435213j,
         ],
         [
-            -0.001618511029519 - 0.0835367795138499j,
-            0.0053302150713874 + 0.0108880702515039j,
-            -0.4304338310290907 + 0.6712032562215353j,
-            0.5750401423577074 + 0.1625231252605274j,
+            -0.27620095 - 0.05049386j,
+            0.47854591 + 0.48737626j,
+            0.6201202 - 0.03549654j,
+            -0.25966725 + 0.03722163j,
         ],
     ],
     [
         [
-            -0.0031072463860293 - 0.6094283173470628j,
-            -0.1706041375326247 - 0.5095579204727588j,
-            0.4717848248508053 - 0.1360136245848086j,
-            -0.3088104665532103 + 0.0582056156280258j,
+            -0.34812515 + 0.37427723j,
+            -0.11092236 + 0.47565307j,
+            0.13724183 - 0.29504039j,
+            -0.56249794 - 0.27908375j,
         ],
         [
-            0.0649795149056778 - 0.0807006955559929j,
-            0.3985010217077551 - 0.4878571486498078j,
-            -0.2593500452019464 - 0.369166183684014j,
-            0.5611176872453115 + 0.2721364042662399j,
+            -0.14408107 + 0.1693212j,
+            -0.20483797 - 0.10707915j,
+            -0.85376825 + 0.3175112j,
+            -0.1054503 - 0.23726165j,
         ],
         [
-            0.6706814505431677 - 0.3020432793307177j,
-            -0.445130459060895 + 0.2625311675303597j,
-            -0.1884345816738757 + 0.0153310963671089j,
-            0.1742601869852907 + 0.3546593618463598j,
+            0.03106625 - 0.04236712j,
+            0.78292822 + 0.03053768j,
+            0.01814738 + 0.16830002j,
+            0.03513342 - 0.59451003j,
         ],
         [
-            -0.0330058354970538 + 0.2751715747427885j,
-            -0.2166842655029602 - 0.0205300200732348j,
-            0.6878324043735908 + 0.2156784765639485j,
-            0.5870150498641511 + 0.1077050421736986j,
+            -0.00087219 - 0.82857442j,
+            0.04840206 + 0.30294214j,
+            -0.1884474 + 0.01468393j,
+            -0.41353861 + 0.11227088j,
         ],
     ],
     [
         [
-            -0.1517007161518054 - 0.3477896338392819j,
-            -0.1996932364673036 - 0.3437390801195139j,
-            -0.5626622855975293 - 0.1026521198436937j,
-            0.5737100114799175 + 0.2042690304628801j,
+            -0.05780187 - 0.06284269j,
+            0.13559069 + 0.19399748j,
+            0.12381697 + 0.01612151j,
+            0.71416466 - 0.64114599j,
         ],
         [
-            -0.1782045187028614 - 0.1798746573123555j,
-            0.6710156796470433 + 0.1482075706151732j,
-            -0.4090747335747313 + 0.4717438839366156j,
-            -0.2113406078464153 + 0.1706207743654816j,
+            -0.31103029 - 0.06658675j,
+            -0.50183231 + 0.49812898j,
+            -0.58061141 - 0.20451914j,
+            -0.07379796 - 0.12030957j,
         ],
         [
-            0.0355053593911394 - 0.6453216466524461j,
-            -0.1081082592116457 + 0.4107406869820365j,
-            0.4773230440109602 + 0.2293807227814043j,
-            0.2490211402335333 + 0.2438025872570191j,
+            0.47241806 - 0.79298028j,
+            0.14041019 + 0.06342211j,
+            -0.27789855 + 0.19625469j,
+            -0.07716877 - 0.05067088j,
         ],
         [
-            -0.5071192657577112 - 0.3421211957731553j,
-            -0.3794415508044536 - 0.2129960836701763j,
-            -0.0100377824272845 - 0.0493528476711767j,
-            -0.6523577435362742 + 0.0912843325184198j,
+            0.11114093 - 0.16488557j,
+            -0.12688073 + 0.63574829j,
+            0.68327072 - 0.15122624j,
+            -0.21697355 + 0.05813823j,
         ],
     ],
     [
         [
-            -0.0196408936733692 + 0.0436166575003405j,
-            -0.4114171776549808 + 0.1604631422478851j,
-            0.8090489042142334 + 0.1854903095749424j,
-            -0.0285563913316359 - 0.336031332694153j,
+            0.32457875 - 0.36309659j,
+            -0.21084067 + 0.48248995j,
+            -0.21588245 - 0.42368088j,
+            -0.0474262 + 0.50714809j,
         ],
         [
-            0.6111824978145968 - 0.0373559219203255j,
-            -0.0607088813252194 - 0.4399072558564469j,
-            0.0652009339801155 - 0.3851081946614576j,
-            0.4498048829332472 - 0.2701344319737953j,
+            -0.24688996 - 0.11890225j,
+            0.16113004 + 0.1518989j,
+            -0.40132738 - 0.28678782j,
+            -0.63810805 - 0.4747406j,
         ],
         [
-            -0.010769501670078 + 0.4210446448669616j,
-            -0.1955139271138671 + 0.5982440344653875j,
-            -0.3421820362428211 - 0.0107432223507508j,
-            0.4746580414982144 - 0.2897923140475613j,
+            -0.14474527 - 0.46561401j,
+            0.47151308 - 0.32560877j,
+            0.51600239 - 0.28298318j,
+            -0.18491473 + 0.23103107j,
         ],
         [
-            0.598418329142766 + 0.2954204464071747j,
-            -0.4534575233858022 + 0.0777057134279078j,
-            -0.1459545057858948 + 0.1413048185423547j,
-            -0.4168424055846296 + 0.3576772276508062j,
+            0.42624962 - 0.51795827j,
+            -0.17138618 - 0.56213399j,
+            -0.36133453 + 0.23168462j,
+            -0.0167845 - 0.14191731j,
         ],
     ],
 ]
@@ -317,82 +311,82 @@ samples_2_cnots = [
     # (A \otimes B) CNOT10 CNOT01 (C \otimes D)
     [
         [
-            -0.004631217256064 - 0.268534699868393j,
-            -0.2794044288596784 - 0.2842241448919913j,
-            0.1218071742216182 - 0.7879690440506752j,
-            0.2778952482254939 + 0.2367688668871671j,
+            0.06593717 - 0.26170549j,
+            -0.70347203 - 0.07117793j,
+            0.23473954 - 0.2369032j,
+            0.5244829 - 0.20227516j,
         ],
         [
-            0.4392070528530516 + 0.6904136673657634j,
-            -0.1826413830164512 + 0.2570714399459114j,
-            -0.1294472995785671 - 0.1887672822622143j,
-            -0.0785276120371355 + 0.4152424069619657j,
+            -0.20854782 + 0.32788173j,
+            -0.43987873 + 0.2876307j,
+            -0.14557311 + 0.70166825j,
+            0.10554244 + 0.21933444j,
         ],
         [
-            0.1620021717921561 + 0.1884711066839195j,
-            0.4872686757195456 - 0.6397614642921917j,
-            -0.4411324161480765 - 0.1891338003177969j,
-            -0.2059859403259958 - 0.1367802403050884j,
+            -0.10107 - 0.78553725j,
+            0.30719144 + 0.10962565j,
+            -0.08963713 + 0.42499178j,
+            0.27869654 - 0.00267013j,
         ],
         [
-            0.3022798492102187 - 0.3242755471407363j,
-            -0.1551409732131101 + 0.2662930443472138j,
-            -0.0292247033180649 - 0.2839975729154084j,
-            -0.7363616289920063 - 0.2911260927728274j,
+            0.06589238 - 0.38018178j,
+            -0.33913438 + 0.04939232j,
+            0.4184943 + 0.10702358j,
+            -0.7203013 + 0.16805593j,
         ],
     ],
     # CNOT01 (A \otimes B) CNOT01 (C \otimes D)
     [
         [
-            0.4952546422732242 - 0.3874950570842591j,
-            -0.565120684005787 - 0.0812920957189095j,
-            0.0599739067738637 - 0.3382684721070105j,
-            0.2002937293624679 - 0.3470743738556477j,
+            0.31547811 - 0.58699408j,
+            -0.34881989 - 0.41836625j,
+            0.30342551 - 0.04638021j,
+            -0.28264065 - 0.29172243j,
         ],
         [
-            0.3076748055295139 + 0.1668987026621435j,
-            0.1113148363873911 - 0.355244639751341j,
-            -0.5552367793191325 + 0.2360956317049145j,
-            -0.4402702690966316 - 0.4254695664213973j,
+            -0.16129481 + 0.35079944j,
+            -0.03709334 - 0.15894938j,
+            0.0736614 + 0.55424822j,
+            -0.71185497 + 0.0702021j,
         ],
         [
-            -0.0570544619126448 - 0.5997099301388422j,
-            0.3690979468442732 - 0.489643527399396j,
-            -0.2915922328649661 - 0.1955865908229054j,
-            0.1605581200806802 + 0.3347381888568993j,
+            -0.30829764 + 0.23330248j,
+            -0.34470925 - 0.74564132j,
+            -0.21494836 - 0.09732758j,
+            0.28996211 + 0.18964071j,
         ],
         [
-            0.3161801218120408 - 0.1384715702727579j,
-            0.1624040016253136 + 0.3648031287938079j,
-            -0.2674750917123022 + 0.5684412664492574j,
-            0.5714177142050261 - 0.0145866404362501j,
+            -0.44032041 + 0.25194387j,
+            0.04287028 + 0.00320109j,
+            0.72290688 - 0.12204478j,
+            0.15915489 - 0.4218703j,
         ],
     ],
     # (A \otimes B) CNOT10 (C \otimes D) CNOT01 (E \otimes F)
     [
         [
-            -0.5032838832977308 + 0.008462334360567j,
-            -0.1994567215589622 + 0.4166138291684532j,
-            -0.2927933889967656 + 0.3307978265454508j,
-            0.0980847185817039 - 0.5731560630175854j,
+            0.26680917 + 0.24320208j,
+            0.05925905 + 0.84940479j,
+            0.04305845 + 0.16580725j,
+            -0.03847889 + 0.33740004j,
         ],
         [
-            0.429674866329586 + 0.1517854863780976j,
-            -0.260061924207509 - 0.0913938324295813j,
-            0.4678474940708733 + 0.4691638242216676j,
-            0.4693562225846957 - 0.238881566880469j,
+            -0.05713611 - 0.69307062j,
+            -0.14217609 + 0.26062907j,
+            -0.12264675 - 0.59272591j,
+            -0.14185396 + 0.20434836j,
         ],
         [
-            -0.3407683276656533 - 0.1781416508017804j,
-            0.648130427328909 + 0.1856223456906182j,
-            0.3247150132557211 - 0.0284712498269712j,
-            0.5276154053566222 + 0.1139531710411058j,
+            0.18147763 + 0.27516979j,
+            -0.2868943 - 0.08518162j,
+            -0.09780205 - 0.39614393j,
+            0.76677616 + 0.21758277j,
         ],
         [
-            -0.6092485765643428 - 0.1411845180727529j,
-            -0.3732001418286821 - 0.3418447668630915j,
-            0.4811869468078607 + 0.1675038147841858j,
-            -0.1994479499053229 + 0.2310173648263156j,
+            -0.52667927 - 0.00325491j,
+            -0.05056066 + 0.30779487j,
+            0.65170215 - 0.11435361j,
+            0.2628016 - 0.34416154j,
         ],
     ],
 ]
@@ -404,163 +398,163 @@ samples_1_cnot = [
     # CNOT10 (A \otimes B)
     [
         [
-            0.1386390404432302 - 0.334570112999545j,
-            0.0148174707054368 - 0.0412250336805805j,
-            -0.4992326027350552 + 0.7779384196647746j,
-            -0.0558235265127082 + 0.0968780293767465j,
+            -0.40013501 + 0.6346169j,
+            0.36588428 + 0.27955713j,
+            -0.34061004 - 0.21797615j,
+            -0.15200951 + 0.19619931j,
         ],
         [
-            0.1110962138948646 + 0.0126194273429225j,
-            -0.922378240060485 - 0.0603304126213831j,
-            0.0419677169072708 + 0.0125607156508953j,
-            -0.3515157103454069 - 0.0871472858321581j,
+            0.19619931 - 0.15200951j,
+            0.21797615 + 0.34061004j,
+            -0.27955713 - 0.36588428j,
+            0.6346169 - 0.40013501j,
         ],
         [
-            -0.7766806433421076 - 0.5011871453677099j,
-            -0.0967373552333893 - 0.0560669491829709j,
-            -0.2647978260425042 - 0.2470628571248854j,
-            -0.0334214548399949 - 0.0283207908867259j,
+            -0.04618249 + 0.40174118j,
+            0.24303208 + 0.05036332j,
+            0.74472919 + 0.09069221j,
+            0.09646726 - 0.45024168j,
         ],
         [
-            0.0029378421621844 + 0.0437084651240645j,
-            -0.0415219650227516 - 0.3597691905488813j,
-            0.0123400367724977 - 0.1111275941460545j,
-            -0.0580109080696657 + 0.922527024519649j,
+            0.45024168 - 0.09646726j,
+            0.09069221 + 0.74472919j,
+            0.05036332 + 0.24303208j,
+            -0.40174118 + 0.04618249j,
         ],
     ],
     # CNOT01 (A \otimes B)
     [
         [
-            0.1275272250362362 - 0.3343853948894475j,
-            0.6337142336419559 - 0.1759440058732874j,
-            0.2947329435072354 - 0.116226264283067j,
-            0.4932993532409792 + 0.309278009290073j,
+            -0.65867779 - 0.39965996j,
+            0.10480488 - 0.11993597j,
+            0.33775094 + 0.50135151j,
+            -0.11410732 + 0.0509635j,
         ],
         [
-            0.3999023544680882 - 0.5221380373177482j,
-            -0.3578174135830141 - 0.0065942382104973j,
-            0.5800639860980087 - 0.0502285871838717j,
-            -0.2092751655259081 - 0.2378654190750314j,
+            -0.15620406 + 0.03112864j,
+            0.0184181 + 0.77022418j,
+            0.12151068 + 0.02920466j,
+            0.24022431 - 0.55472634j,
         ],
         [
-            -0.1565164704901467 + 0.5608027580207731j,
-            0.296176258834552 - 0.1124974497506721j,
-            0.5870241263928151 - 0.2965682015252186j,
-            -0.3194248225208941 - 0.1613832959085598j,
+            0.0509635 - 0.11410732j,
+            -0.50135151 - 0.33775094j,
+            0.11993597 - 0.10480488j,
+            -0.39965996 - 0.65867779j,
         ],
         [
-            0.0054101937368828 + 0.3167756024402226j,
-            -0.4627149049700612 + 0.3534007012683281j,
-            0.2601293856027128 - 0.2457834172330788j,
-            0.647226333094554 + 0.1168250695361578j,
+            0.55472634 - 0.24022431j,
+            0.02920466 + 0.12151068j,
+            0.77022418 + 0.0184181j,
+            -0.03112864 + 0.15620406j,
         ],
     ],
     # (A \otimes B) CNOT01
     [
         [
-            0.190458746413919 + 0.2870042726555157j,
-            0.502809168004646 + 0.6156507128238845j,
-            -0.3763228062935716 + 0.2616231647579805j,
-            -0.1734766279722236 + 0.096703042530068j,
+            0.61975417 + 0.05871326j,
+            0.18490393 - 0.03060346j,
+            0.11206333 - 0.18820514j,
+            0.51965346 - 0.50922726j,
         ],
         [
-            -0.7237411404174169 - 0.3286968531905625j,
-            0.3261596499092991 + 0.1107513854501445j,
-            -0.0781992451732389 + 0.1825664178909173j,
-            0.2212035729951513 - 0.4014156376735081j,
+            0.17568798 - 0.06526691j,
+            -0.50876306 + 0.35875148j,
+            -0.06315069 + 0.72482005j,
+            0.07415843 - 0.20610651j,
         ],
         [
-            -0.1821207882156272 + 0.0792315428784779j,
-            -0.4001569202892715 + 0.2234725352683322j,
-            0.3327888523333447 + 0.7218687134048248j,
-            0.1125960761054146 + 0.3255274339495012j,
+            0.72482005 - 0.06315069j,
+            0.20610651 - 0.07415843j,
+            -0.06526691 + 0.17568798j,
+            -0.35875148 + 0.50876306j,
         ],
         [
-            0.2594885215726799 - 0.3777978851802577j,
-            -0.0957194014420285 + 0.1740213063816186j,
-            0.2880779045831403 + 0.1888308965690472j,
-            -0.6184873664151125 - 0.4993157690429775j,
+            0.18820514 - 0.11206333j,
+            -0.50922726 + 0.51965346j,
+            -0.05871326 - 0.61975417j,
+            -0.03060346 + 0.18490393j,
         ],
     ],
     # (A \otimes B) CNOT10
     [
         [
-            -0.0885673151087696 - 0.0243844138270481j,
-            -0.4871418040378778 + 0.4830303555130133j,
-            -0.7088634125078255 - 0.1034876301900189j,
-            -0.0696323339281462 + 0.0537594674970065j,
+            0.19339868 + 0.08633274j,
+            0.01466384 + 0.70171345j,
+            -0.63869976 + 0.02205134j,
+            0.09769795 - 0.21108831j,
         ],
         [
-            0.0404303791439687 - 0.0781289108852639j,
-            0.3879597260952143 + 0.6022327441469053j,
-            0.2377071200057033 - 0.6435221726459008j,
-            0.0398405677965372 + 0.082773776650062j,
+            0.18782762 - 0.1372004j,
+            -0.56109965 - 0.30592619j,
+            -0.32755953 + 0.62074274j,
+            0.12258652 + 0.17271055j,
         ],
         [
-            -0.5889549838714252 + 0.4078345916760804j,
-            -0.0074407415198821 - 0.0876548780130169j,
-            0.0684934101349958 - 0.0612161883111329j,
-            -0.0267574702860782 + 0.6854994524959126j,
+            0.30592619 + 0.56109965j,
+            0.1372004 - 0.18782762j,
+            0.17271055 + 0.12258652j,
+            0.62074274 - 0.32755953j,
         ],
         [
-            -0.2651157710067235 - 0.6327235490732244j,
-            -0.0911552999919856 - 0.011378937851836j,
-            0.0437473475396883 + 0.0763211101667442j,
-            0.6944805007772098 + 0.1757664963376742j,
+            0.70171345 + 0.01466384j,
+            0.08633274 + 0.19339868j,
+            0.21108831 - 0.09769795j,
+            -0.02205134 + 0.63869976j,
         ],
     ],
     # (A \otimes B) CNOT10 (C \otimes D)
     [
         [
-            0.2211933286892437 + 0.0021586412960177j,
-            0.38444713219383 - 0.0581554847392023j,
-            -0.170894684604903 + 0.1031928176859452j,
-            -0.1036775219695835 - 0.8656121616196758j,
+            0.78385585 + 0.24581188j,
+            -0.08367428 + 0.34003966j,
+            -0.16227836 - 0.08719991j,
+            -0.222293 + 0.34520414j,
         ],
         [
-            0.3442924845446644 + 0.260151125836039j,
-            -0.1914328218155123 - 0.0746552589054714j,
-            -0.5744657965166804 - 0.6284808877028109j,
-            -0.2036000102332332 + 0.071498266549486j,
+            -0.40100084 - 0.10151918j,
+            0.07269317 + 0.76372534j,
+            -0.26873591 + 0.2188187j,
+            -0.34389609 - 0.04434368j,
         ],
         [
-            -0.1048402095394854 + 0.5447532602544797j,
-            0.7011096884305793 - 0.024182937133914j,
-            -0.1624036601402791 + 0.1839401392571305j,
-            -0.1162941782466869 + 0.3554990108217186j,
+            0.17914714 - 0.32406563j,
+            0.2474564 - 0.0645854j,
+            0.45026225 - 0.23004978j,
+            -0.66076185 - 0.32437111j,
         ],
         [
-            -0.6743897607764879 + 0.0479978633805106j,
-            -0.167244245931061 - 0.5352458890411946j,
-            -0.4128361178667754 + 0.06711548036308j,
-            0.1819390859257292 - 0.1428141352633625j,
+            0.04747281 + 0.12113138j,
+            0.4725868 + 0.00816913j,
+            -0.53596566 - 0.55049352j,
+            0.1451332 - 0.38510071j,
         ],
     ],
     # (A \otimes B) CNOT01 (C \otimes D)
     [
         [
-            -0.384175417575186 + 0.3923608184977355j,
-            0.1002767236561732 - 0.5802482358089508j,
-            -0.4052191818361661 + 0.4207032744597814j,
-            -0.1025417468296731 + 0.0031874535396264j,
+            0.38565489 - 0.50840093j,
+            -0.30818989 - 0.11768624j,
+            0.27524652 - 0.39173492j,
+            -0.5046883 - 0.00636884j,
         ],
         [
-            -0.4577221128473338 + 0.0586197038286386j,
-            -0.0937187062570979 + 0.6742419194459858j,
-            -0.0414783182526488 + 0.3776270025012934j,
-            -0.2437968467006254 + 0.3462792652304717j,
+            -0.05689053 - 0.29547371j,
+            0.51930406 - 0.09499063j,
+            -0.40250308 - 0.33321454j,
+            0.00587892 - 0.59806656j,
         ],
         [
-            0.3423554335462237 + 0.4267199288134111j,
-            0.1796472451791323 + 0.1419455069613175j,
-            -0.1490674637388668 - 0.2557411476563918j,
-            -0.7409292809972994 - 0.1080757005000481j,
+            0.56972698 - 0.14750915j,
+            -0.32434509 + 0.18005613j,
+            -0.62729996 + 0.30972331j,
+            0.15801266 - 0.04037939j,
         ],
         [
-            0.4279595694962425 - 0.0554066848293549j,
-            0.185704890878562 + 0.3208788487100644j,
-            -0.4049260830288736 + 0.5127316935400166j,
-            0.2406649865130532 - 0.4376672321481798j,
+            -0.37603427 + 0.12044912j,
+            -0.64676631 + 0.23778647j,
+            0.03588508 - 0.08379105j,
+            0.02300689 - 0.60033588j,
         ],
     ],
 ]
@@ -572,38 +566,48 @@ samples_su2_su2 = [
     (X, H),
     (
         [
-            [-0.8951484711495978 + 0.2618482674962029j, 0.1320728980584716 + 0.3357103647583892j],
-            [0.1145773085667637 + 0.3420770962361714j, -0.8540795475070558 + 0.3747044530324777j],
+            [0.21993927 - 0.1111822j, -0.27174921 - 0.93027824j],
+            [0.27174921 - 0.93027824j, 0.21993927 + 0.1111822j],
         ],
         [
-            [0.535259309608487 - 0.7568444138362017j, -0.0856403406729011 - 0.3651708323125105j],
-            [0.2346548909477215 + 0.2926108112778731j, -0.9258821807212569 - 0.0453671985080405j],
+            [0.86361715 + 0.11195238j, 0.35098794 - 0.34415046j],
+            [-0.35098794 - 0.34415046j, 0.86361715 - 0.11195238j],
         ],
     ),
     (
         [
-            [-0.3475218186197753 + 0.616967631919303j, -0.6918447778561189 - 0.1411748210564089j],
-            [0.4110483325261496 + 0.5741243724791375j, 0.4216047723512865 - 0.5689199321376678j],
+            [0.08652981 - 0.25976406j, 0.95775011 + 0.08803377j],
+            [-0.95775011 + 0.08803377j, 0.08652981 + 0.25976406j],
         ],
         [
-            [-0.1036172252382817 - 0.409539663695392j, 0.0173544935458401 + 0.9062226856837956j],
-            [-0.3763539496942646 + 0.8245595424490042j, -0.2579549535044457 + 0.3345422357017644j],
+            [-0.1668335 + 0.92244968j, -0.34571207 + 0.04166941j],
+            [0.34571207 + 0.04166941j, -0.1668335 - 0.92244968j],
         ],
     ),
     (
         [
-            [-0.6199757150883975 - 0.4615217464182165j, -0.5797949307229534 + 0.257809287248686j],
-            [0.5072854247342099 - 0.381168319953196j, 0.1016602350822234 + 0.7661836635685451j],
+            [-0.36170896 - 0.92179965j, -0.1390035 + 0.01140442j],
+            [0.1390035 + 0.01140442j, -0.36170896 + 0.92179965j],
         ],
         [
-            [-0.5741575760925136 - 0.4777766752909318j, -0.6560499370893111 + 0.1080324322072939j],
-            [0.3761599015364478 - 0.5482483514244275j, 0.1890587151173228 + 0.7226231907949495j],
+            [-0.26587616 + 0.09917768j, 0.731434 + 0.62006287j],
+            [-0.731434 + 0.62006287j, -0.26587616 - 0.09917768j],
         ],
     ),
     (
         [
-            [0.257810302107603 + 0.3791450460627684j, 0.7491086034534051 - 0.478141383280444j],
-            [-0.6878868994174759 + 0.5626673047041733j, -0.2798521896663711 - 0.3631802166496543j],
+            [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],
         ],
         [[np.exp(-1j * np.pi / 3), 0], [0, np.exp(1j * np.pi / 3)]],
     ),
@@ -613,14 +617,14 @@ samples_su2_su2 = [
             [-1j, 0],
         ],
         [
-            [-0.6429551068755086 - 0.2344852807169338j, 0.348170308379626 + 0.6406268961202303j],
-            [-0.0720224467655562 + 0.7255605769552991j, -0.5014071443907067 + 0.4657955472996208j],
+            [-0.05594177 - 0.02155518j, -0.83242096 - 0.55089131j],
+            [0.83242096 - 0.55089131j, -0.05594177 + 0.02155518j],
         ],
     ),
     (
         [
-            [0.7192114465877868 + 0.4233015158437373j, 0.5404227760496959 + 0.1072098172194949j],
-            [0.5300962072416202 + 0.1501623549680364j, -0.8330292881813453 + 0.0501147009427765j],
+            [0.90778489 + 0.27105877j, 0.17141891 - 0.27031333j],
+            [-0.17141891 - 0.27031333j, 0.90778489 - 0.27105877j],
         ],
         [
             [0, -1j],
@@ -642,6 +646,13 @@ 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 fdeff77fc..8020bb643 100644
--- a/tests/transforms/test_optimization/test_merge_rotations.py
+++ b/tests/transforms/test_optimization/test_merge_rotations.py
@@ -103,7 +103,9 @@ class TestMergeRotations:
             (0.3, -0.3, 0.7, -0.1, [qml.RY(0.6, wires=1)]),
         ],
     )
-    def test_two_qubits_merge(self, theta_11, theta_12, theta_21, theta_22, expected_ops):
+    def test_two_qubits_rotation_no_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():
@@ -396,29 +398,3 @@ 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 670f235db..059753366 100644
--- a/tests/transforms/test_optimization/test_optimization_utils.py
+++ b/tests/transforms/test_optimization/test_optimization_utils.py
@@ -17,13 +17,10 @@ import pytest
 import pennylane as qml
 from pennylane.transforms.optimization.optimization_utils import (
     find_next_gate,
-    _zyz_to_quat,
-    _quaternion_product,
+    _yzy_to_zyz,
     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
 
@@ -57,39 +54,32 @@ class TestRotGateFusion:
     """Test that utility functions for fusing two qml.Rot gates function as expected."""
 
     @pytest.mark.parametrize(
-        ("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]),
-        ],
+        ("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])],
     )
-    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)
+    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."""
 
-    @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)
+        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()
+
+        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"),
diff --git a/tests/transforms/test_optimization/test_single_qubit_fusion.py b/tests/transforms/test_optimization/test_single_qubit_fusion.py
index c1761bbb6..868871412 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(atol=1e-7)(qfunc)
+        transformed_qfunc = single_qubit_fusion()(qfunc)
         transformed_ops = qml.transforms.make_tape(transformed_qfunc)().operations
         assert len(transformed_ops) == 0
 
@@ -296,37 +296,3 @@ 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 ad68673fd..e07b10d97 100644
--- a/tests/transforms/test_unitary_to_rot.py
+++ b/tests/transforms/test_unitary_to_rot.py
@@ -39,17 +39,12 @@ 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, -9.831019270939975e-01 + 0.1830590094588862j],
-                [9.831019270939975e-01 + 0.1830590094588862j, 0],
-            ]
-        ),
+        np.array([[0, -0.98310193 + 0.18305901j], [0.98310193 + 0.18305901j, 0]]),
         qml.Rot,
-        [-0.18409714468526372, np.pi, 0.18409714468526372],
+        [0, -np.pi, -5.914991017809059],
     ),
     (H, qml.Rot, [np.pi, np.pi / 2, 0.0]),
-    (X, qml.Rot, [np.pi / 2, np.pi, -np.pi / 2]),
+    (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.0, 2.0, -3.0, wires=0).get_matrix(),
@@ -82,7 +77,7 @@ class TestDecomposeSingleQubitUnitaryTransform:
 
         assert isinstance(ops[1], expected_gate)
         assert ops[1].wires == Wires("a")
-        assert qml.math.allclose(ops[1].parameters, expected_params, atol=1e-7)
+        assert qml.math.allclose(ops[1].parameters, expected_params)
 
         assert isinstance(ops[2], qml.CNOT)
         assert ops[2].wires == Wires(["b", "a"])
@@ -114,7 +109,7 @@ class TestDecomposeSingleQubitUnitaryTransform:
         """Test that the transform works in the torch interface."""
         torch = pytest.importorskip("torch")
 
-        U = torch.tensor(U, dtype=torch.complex128)
+        U = torch.tensor(U, dtype=torch.complex64)
 
         transformed_qfunc = unitary_to_rot(qfunc)
 
@@ -127,7 +122,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, atol=1e-7)
+        assert qml.math.allclose(qml.math.unwrap(ops[1].parameters), expected_params)
 
         assert isinstance(ops[2], qml.CNOT)
         assert ops[2].wires == Wires(["b", "a"])
@@ -137,7 +132,7 @@ class TestDecomposeSingleQubitUnitaryTransform:
         """Test that the transform works in the Tensorflow interface."""
         tf = pytest.importorskip("tensorflow")
 
-        U = tf.Variable(U, dtype=tf.complex128)
+        U = tf.Variable(U, dtype=tf.complex64)
 
         transformed_qfunc = unitary_to_rot(qfunc)
 
@@ -150,7 +145,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, atol=1e-7)
+        assert qml.math.allclose(qml.math.unwrap(ops[1].parameters), expected_params)
 
         assert isinstance(ops[2], qml.CNOT)
         assert ops[2].wires == Wires(["b", "a"])
@@ -166,7 +161,7 @@ class TestDecomposeSingleQubitUnitaryTransform:
         remember = config.read("jax_enable_x64")
         config.update("jax_enable_x64", True)
 
-        U = jax.numpy.array(U, dtype=jax.numpy.complex128)
+        U = jax.numpy.array(U, dtype=jax.numpy.complex64)
 
         transformed_qfunc = unitary_to_rot(qfunc)
 
@@ -179,40 +174,11 @@ 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, atol=1e-7)
+        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"])
 
-    @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.
@@ -261,12 +227,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), atol=1e-7)
+        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, atol=1e-7)
+        assert qml.math.allclose(original_grad, transformed_grad)
 
     @pytest.mark.parametrize("rot_angles,diff_method", angle_diff_pairs)
     def test_gradient_unitary_to_rot_torch(self, rot_angles, diff_method):
@@ -314,12 +280,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, atol=1e-7)
+        assert qml.math.allclose(original_result, transformed_result)
 
         original_result.backward()
         transformed_result.backward()
 
-        assert qml.math.allclose(original_input.grad, transformed_input.grad, atol=1e-7)
+        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):
@@ -420,12 +386,6 @@ 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