Repository: PennyLaneAI/pennylane
Pull request #4862 (2023-11-21 03:45:54)

XZX Single Qubit Decomposition
**Context:**

Implementing the XZX single qubit decomposition as requested in https://github.com/PennyLaneAI/pennylane/issues/4826 

**Description of the Change:**

A new single qubit decomposition is added which decomposes a 1-qubit unitary as 

$$U=e^{i\gamma}\cdot R_x(\phi)\cdot R_z(\theta)\cdot R_x(\lambda)$$

Additionally, cleaned up ```tests/transforms/test_decompositions.py```:

1. Generalized the _run_assertions method, removed duplicate code.
2. For tests with jax where ```remember = config.read("jax_enable_x64")``` was created to remember the current config before updating it to ```True```, but never actually used to restore the config, add ```config.update("jax_enable_x64", remember)``` at the end of the test case.
3. Renaming ```TestOneQubitUnitaryDecomposition``` to ```TestOneQubitRotDecomposition```, and removing test cases for other decomposition types from this test class. Since now each different type of decomposition has its own test class, the ```TestOneQubitRotDecomposition``` should now be dedicated to only the rot decomposition.
4. Added tests for tf, torch, jax for the rot transform.

**Benefits:**

Nice to have.

**Related GitHub Issues:**

https://github.com/PennyLaneAI/pennylane/issues/4826 


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4862?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Comparison is base [(`1f47116`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/1f471160bf6ec57cbed949a6b8c339613869cc3f?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.65% compared to head [(`0c4744b`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4862?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.64%.


<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #4862      +/-   ##
==========================================
- Coverage   99.65%   99.64%   -0.01%     
==========================================
  Files         382      382              
  Lines       34500    34261     -239     
==========================================
- Hits        34380    34140     -240     
- Misses        120      121       +1     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4862?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

rmoyard: Hi @astralcai thanks for the contribution! I will take a look at your PR this afternoon.

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 5f7b9c1cd..dade66c79 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -7,6 +7,9 @@
 * Approximate Quantum Fourier Transform (AQFT) is now available from `qml.AQFT`.
   [(#4656)](https://github.com/PennyLaneAI/pennylane/pull/4656)
 
+* XZX decomposition is added to the list of supported single-qubit unitary decompositions.
+  [(#4862)](https://github.com/PennyLaneAI/pennylane/pull/4862)
+
 * Iterative Quantum Phase Estimation is now available from `qml.iterative_qpe`.
   [(#4804)](https://github.com/PennyLaneAI/pennylane/pull/4804)
 
@@ -195,6 +198,7 @@
 This release contains contributions from (in alphabetical order):
 
 Guillermo Alonso,
+Astral Cai,
 Amintor Dusko,
 Lillian Frederiksen,
 Emiliano Godinez Ramirez,
diff --git a/pennylane/transforms/decompositions/single_qubit_unitary.py b/pennylane/transforms/decompositions/single_qubit_unitary.py
index 62969e36b..ef420acf1 100644
--- a/pennylane/transforms/decompositions/single_qubit_unitary.py
+++ b/pennylane/transforms/decompositions/single_qubit_unitary.py
@@ -82,7 +82,7 @@ def _rot_decomposition(U, wire):
     can instead recover a ``Rot`` gate that implements the same operation, up
     to a global phase:
 
-    >>> decomp = zyz_decomposition(U, 0)
+    >>> decomp = _rot_decomposition(U, 0)
     >>> decomp
     [Rot(-0.24209529417800013, 1.14938178234275, 1.7330581433950871, wires=[0])]
     """
@@ -261,6 +261,76 @@ def _xyx_decomposition(U, wire, return_global_phase=False):
     return operations
 
 
+def _xzx_decomposition(U, wire, return_global_phase=False):
+    r"""Computes the decomposition of a single-qubit matrix :math:`U` in terms
+    of elementary operations, as a product of Z and X rotations in the form
+    :math:`e^{i\gamma} RX(\phi) RZ(\theta) RX(\lambda)`. (batched operation)
+
+    Args:
+        U (tensor): A :math:`2 \times 2` unitary matrix.
+        wire (Union[Wires, Sequence[int] or int]): The wire on which to apply the operation.
+        return_global_phase (bool): Whether to return the global phase ``qml.GlobalPhase(-gamma)``
+            as the last element of the returned list of operations.
+
+    Returns:
+        list[Operation]: Returns a list of gates, an ``RX``, an ``RZ`` and
+        another ``RX`` gate, which when applied in the order of appearance in the list is
+        equivalent to the unitary :math:`U` up to a global phase. If `return_global_phase=True`,
+        the global phase is returned as the last element of the list.
+
+    **Example**
+
+    >>> U = np.array(
+    ...     [[-0.28829348-0.78829734j,  0.30364367+0.45085995j],
+    ...      [ 0.53396245-0.10177564j,  0.76279558-0.35024096j]]
+    ... )
+    >>> decompositions = _xzx_decomposition(U, 0, return_global_phase=True)
+    >>> decompositions
+    [RX(12.416147693665032, wires=[0]),
+     RZ(1.3974974090935608, wires=[0]),
+     RX(11.448040119199066, wires=[0]),
+     GlobalPhase((1.1759220332464762-0j), wires=[0])]
+
+    """
+
+    # Small number to add to denominators to avoid division by zero
+    EPS = 1e-64
+
+    # Choose gamma such that exp(-i*gamma)*U is special unitary (detU==1).
+    U = math.expand_dims(U, axis=0) if len(U.shape) == 2 else U
+    U_det1, gammas = _convert_to_su2(U, return_global_phase=return_global_phase)
+
+    # Compute \phi, \theta and \lambda after analytically solving for them from
+    # U_det1 = RX(\phi) RZ(\theta) RX(\lambda)
+    sum_diagonal_real = math.real(U_det1[:, 0, 0] + U_det1[:, 1, 1])
+    sum_off_diagonal_imag = math.imag(U_det1[:, 0, 1] + U_det1[:, 1, 0])
+    phi_plus_lambdas_d2 = math.arctan2(-sum_off_diagonal_imag, sum_diagonal_real + EPS)
+    diff_diagonal_imag = math.imag(U_det1[:, 0, 0] - U_det1[:, 1, 1])
+    diff_off_diagonal_real = math.real(U_det1[:, 0, 1] - U_det1[:, 1, 0])
+    phi_minus_lambdas_d2 = math.arctan2(diff_off_diagonal_real, -diff_diagonal_imag + EPS)
+    lams = phi_plus_lambdas_d2 - phi_minus_lambdas_d2
+    phis = phi_plus_lambdas_d2 + phi_minus_lambdas_d2
+
+    # Compute \theta
+    thetas = math.where(
+        math.isclose(math.sin(phi_plus_lambdas_d2), math.zeros_like(phi_plus_lambdas_d2)),
+        2 * math.arccos(sum_diagonal_real / (2 * math.cos(phi_plus_lambdas_d2) + EPS)),
+        2 * math.arccos(-sum_off_diagonal_imag / (2 * math.sin(phi_plus_lambdas_d2) + EPS)),
+    )
+
+    phis, thetas, lams, gammas = map(math.squeeze, [phis, thetas, lams, gammas])
+
+    phis = phis % (4 * np.pi)
+    thetas = thetas % (4 * np.pi)
+    lams = lams % (4 * np.pi)
+
+    operations = [qml.RX(lams, wire), qml.RZ(thetas, wire), qml.RX(phis, wire)]
+    if return_global_phase:
+        operations.append(qml.GlobalPhase(-gammas))
+
+    return operations
+
+
 def _zxz_decomposition(U, wire, return_global_phase=False):
     r"""Compute the decomposition of a single-qubit matrix :math:`U` in terms
     of elementary operations, as a product of X and Z rotations in the form
@@ -310,6 +380,7 @@ def _zxz_decomposition(U, wire, return_global_phase=False):
         math.real(U_det1[:, 0, 0]) / (math.cos(phis_plus_psis) + EPS),
         -math.imag(U_det1[:, 0, 0]) / (math.sin(phis_plus_psis) + EPS),
     )
+
     # Arcos is only defined between -1 and 1
     thetas = qml.math.clip(thetas, -1.0, 1.0)
     thetas = 2 * math.arccos(thetas)
@@ -334,7 +405,7 @@ def one_qubit_decomposition(U, wire, rotations="ZYZ", return_global_phase=False)
     Any one qubit unitary operation can be implemented up to a global phase by composing RX, RY,
     and RZ gates.
 
-    Currently supported values for ``rotations`` are "ZYZ", "XYX", and "ZXZ".
+    Currently supported values for ``rotations`` are "ZYZ", "XYX", "XZX", and "ZXZ".
 
     Args:
         U (tensor): A :math:`2 \times 2` unitary matrix.
@@ -363,6 +434,7 @@ def one_qubit_decomposition(U, wire, rotations="ZYZ", return_global_phase=False)
         "rot": _rot_decomposition,
         "ZYZ": _zyz_decomposition,
         "XYX": _xyx_decomposition,
+        "XZX": _xzx_decomposition,
         "ZXZ": _zxz_decomposition,
     }
 


Test patch:
diff --git a/tests/transforms/test_decompositions.py b/tests/transforms/test_decompositions.py
index b4bdda968..45965047e 100644
--- a/tests/transforms/test_decompositions.py
+++ b/tests/transforms/test_decompositions.py
@@ -35,6 +35,43 @@ from pennylane.transforms.decompositions.two_qubit_unitary import (
 )
 
 
+def _run_assertions(U, expected_gates, expected_params, obtained_gates):
+    """Runs assertions for unitary decomposition transforms."""
+
+    assert len(obtained_gates) == 4, "Incorrect number of gates"
+
+    # Check initial 3 gates
+    for i in range(3):
+        assert isinstance(obtained_gates[i], expected_gates[i]), "Incorrect type of gate"
+        assert obtained_gates[i].wires == Wires("a"), "Incorrect wire"
+
+    # Check the global phase
+    assert isinstance(obtained_gates[3], expected_gates[3]), "Incorrect type of gate"
+    assert qml.math.allclose(
+        qml.math.unwrap(obtained_gates[3].parameters[:1]), expected_params[3]
+    ), "Incorrect global phase"
+
+    # Now we check the rotation angles
+    assert qml.math.allclose(
+        [qml.math.unwrap(o.parameters)[0] for o in obtained_gates[:3]],
+        expected_params[:3],
+        atol=1e-7,
+    ), "Incorrect rotation angles"
+
+    obtained_mat = reduce(
+        np.matmul, [op.matrix(wire_order=["a"]) for op in reversed(obtained_gates)]
+    )
+
+    if len(obtained_mat.shape) == 2:
+        U = [U]
+        obtained_mat = [obtained_mat]
+
+    assert all(
+        check_matrix_equivalence(curr_obtained_mat, curr_U, atol=1e-7)
+        for curr_obtained_mat, curr_U in zip(obtained_mat, qml.math.unwrap(U))
+    )
+
+
 typeof_gates_zyz = (qml.RZ, qml.RY, qml.RZ, qml.GlobalPhase)
 single_qubit_decomps_zyz = [
     (I, typeof_gates_zyz, [0.0, 0.0, 0.0, 0]),
@@ -99,73 +136,40 @@ single_qubit_decomps_zyz = [
 class TestQubitUnitaryZYZDecomposition:
     """Test that the ZYZ decomposition is correct."""
 
-    def _run_assertions(self, U, expected_gates, expected_params, obtained_gates):
-        assert len(obtained_gates) == 4, "Incorrect number of gates"
-        # Check initial 3 gates
-        for i in range(3):
-            assert isinstance(obtained_gates[i], expected_gates[i]), "Incorrect type of gate"
-            assert obtained_gates[i].wires == Wires("a"), "Incorrect wire"
-
-        # Check the global phase
-        assert isinstance(obtained_gates[3], expected_gates[3]), "Incorrect type of gate"
-        assert qml.math.allclose(
-            qml.math.unwrap(obtained_gates[3].parameters[:1]), expected_params[3]
-        ), "Incorrect global phase"
-        # Now we check the ZYZ rotation angles
-        assert qml.math.allclose(
-            [qml.math.unwrap(o.parameters)[0] for o in obtained_gates[:3]],
-            expected_params[:3],
-            atol=1e-7,
-        ), "Incorrect ZYZ rotation angles"
-
-        obtained_mat = reduce(
-            np.matmul, [op.matrix(wire_order=["a"]) for op in reversed(obtained_gates)]
-        )
-
-        if len(obtained_mat.shape) == 2:
-            U = [U]
-            obtained_mat = [obtained_mat]
-
-        assert all(
-            check_matrix_equivalence(curr_obtained_mat, curr_U, atol=1e-7)
-            for curr_obtained_mat, curr_U in zip(obtained_mat, qml.math.unwrap(U))
-        )
-
     @pytest.mark.parametrize("U,expected_gates,expected_params", single_qubit_decomps_zyz)
     def test_zyz_decomposition(self, U, expected_gates, expected_params):
         """Test that a one-qubit matrix in isolation is correctly decomposed."""
-        obtained_gates = one_qubit_decomposition(U, Wires("a"), "ZYZ", return_global_phase=True)
 
-        self._run_assertions(U, expected_gates, expected_params, obtained_gates)
+        obtained_gates = one_qubit_decomposition(U, Wires("a"), "ZYZ", return_global_phase=True)
+        _run_assertions(U, expected_gates, expected_params, obtained_gates)
 
     @pytest.mark.torch
     @pytest.mark.parametrize("U,expected_gates,expected_params", single_qubit_decomps_zyz)
     def test_zyz_decomposition_torch(self, U, expected_gates, expected_params):
         """Test that a one-qubit operation in Torch is correctly decomposed."""
+
         import torch
 
         U = torch.tensor(U, dtype=torch.complex128)
-
         obtained_gates = one_qubit_decomposition(U, Wires("a"), "ZYZ", return_global_phase=True)
-
-        self._run_assertions(U, expected_gates, expected_params, obtained_gates)
+        _run_assertions(U, expected_gates, expected_params, obtained_gates)
 
     @pytest.mark.tf
     @pytest.mark.parametrize("U,expected_gates,expected_params", single_qubit_decomps_zyz)
     def test_zyz_decomposition_tf(self, U, expected_gates, expected_params):
         """Test that a one-qubit operation in Tensorflow is correctly decomposed."""
+
         import tensorflow as tf
 
         U = tf.Variable(U, dtype=tf.complex128)
-
         obtained_gates = one_qubit_decomposition(U, Wires("a"), "ZYZ", return_global_phase=True)
-
-        self._run_assertions(U, expected_gates, expected_params, obtained_gates)
+        _run_assertions(U, expected_gates, expected_params, obtained_gates)
 
     @pytest.mark.jax
     @pytest.mark.parametrize("U,expected_gates,expected_params", single_qubit_decomps_zyz)
     def test_zyz_decomposition_jax(self, U, expected_gates, expected_params):
         """Test that a one-qubit operation in JAX is correctly decomposed."""
+
         import jax
 
         # Enable float64 support
@@ -175,10 +179,11 @@ class TestQubitUnitaryZYZDecomposition:
         config.update("jax_enable_x64", True)
 
         U = jax.numpy.array(U, dtype=jax.numpy.complex128)
-
         obtained_gates = one_qubit_decomposition(U, Wires("a"), "ZYZ", return_global_phase=True)
+        _run_assertions(U, expected_gates, expected_params, obtained_gates)
 
-        self._run_assertions(U, expected_gates, expected_params, obtained_gates)
+        # Reset the configuration
+        config.update("jax_enable_x64", remember)
 
 
 typeof_gates_xyx = (qml.RX, qml.RY, qml.RX, qml.GlobalPhase)
@@ -227,72 +232,40 @@ single_qubit_decomps_xyx = [
 class TestQubitUnitaryXYXDecomposition:
     """Test that the XYX decomposition is correct."""
 
-    def _run_assertions(self, U, expected_gates, expected_params, obtained_gates):
-        assert len(obtained_gates) == 4, "Incorrect number of gates"
-        # check first 3 gates
-        for i in range(3):
-            assert isinstance(obtained_gates[i], expected_gates[i]), "Incorrect type of gate"
-            assert obtained_gates[i].wires == Wires("a"), "Incorrect wire"
-        # Check the global phase
-        assert isinstance(obtained_gates[3], expected_gates[3]), "Incorrect type of gate"
-        assert qml.math.allclose(
-            qml.math.unwrap(obtained_gates[3].parameters[:1]), expected_params[3]
-        ), "Incorrect global phase"
-        # Now we check the XYX rotation angles
-        assert qml.math.allclose(
-            [qml.math.unwrap(o.parameters)[0] for o in obtained_gates[:3]],
-            expected_params[:3],
-            atol=1e-7,
-        ), "Incorrect XYX rotation angles"
-
-        obtained_mat = reduce(
-            np.matmul, [op.matrix(wire_order=["a"]) for op in reversed(obtained_gates)]
-        )
-
-        if len(obtained_mat.shape) == 2:
-            U = [U]
-            obtained_mat = [obtained_mat]
-
-        assert all(
-            check_matrix_equivalence(curr_obtained_mat, curr_U, atol=1e-7)
-            for curr_obtained_mat, curr_U in zip(obtained_mat, qml.math.unwrap(U))
-        )
-
     @pytest.mark.parametrize("U,expected_gates,expected_params", single_qubit_decomps_xyx)
     def test_xyx_decomposition(self, U, expected_gates, expected_params):
         """Test that a one-qubit matrix in isolation is correctly decomposed."""
-        obtained_gates = one_qubit_decomposition(U, Wires("a"), "XYX", return_global_phase=True)
 
-        self._run_assertions(U, expected_gates, expected_params, obtained_gates)
+        obtained_gates = one_qubit_decomposition(U, Wires("a"), "XYX", return_global_phase=True)
+        _run_assertions(U, expected_gates, expected_params, obtained_gates)
 
     @pytest.mark.torch
     @pytest.mark.parametrize("U,expected_gates,expected_params", single_qubit_decomps_xyx)
     def test_xyx_decomposition_torch(self, U, expected_gates, expected_params):
         """Test that a one-qubit operation in Torch is correctly decomposed."""
+
         import torch
 
         U = torch.tensor(U, dtype=torch.complex128)
-
         obtained_gates = one_qubit_decomposition(U, Wires("a"), "XYX", return_global_phase=True)
-
-        self._run_assertions(U, expected_gates, expected_params, obtained_gates)
+        _run_assertions(U, expected_gates, expected_params, obtained_gates)
 
     @pytest.mark.tf
     @pytest.mark.parametrize("U,expected_gates,expected_params", single_qubit_decomps_xyx)
     def test_xyx_decomposition_tf(self, U, expected_gates, expected_params):
         """Test that a one-qubit operation in Tensorflow is correctly decomposed."""
+
         import tensorflow as tf
 
         U = tf.Variable(U, dtype=tf.complex128)
-
         obtained_gates = one_qubit_decomposition(U, Wires("a"), "XYX", return_global_phase=True)
-
-        self._run_assertions(U, expected_gates, expected_params, obtained_gates)
+        _run_assertions(U, expected_gates, expected_params, obtained_gates)
 
     @pytest.mark.jax
     @pytest.mark.parametrize("U,expected_gates,expected_params", single_qubit_decomps_xyx)
     def test_xyx_decomposition_jax(self, U, expected_gates, expected_params):
         """Test that a one-qubit operation in JAX is correctly decomposed."""
+
         import jax
 
         # Enable float64 support
@@ -302,10 +275,99 @@ class TestQubitUnitaryXYXDecomposition:
         config.update("jax_enable_x64", True)
 
         U = jax.numpy.array(U, dtype=jax.numpy.complex128)
-
         obtained_gates = one_qubit_decomposition(U, Wires("a"), "XYX", return_global_phase=True)
+        _run_assertions(U, expected_gates, expected_params, obtained_gates)
+
+        # Reset the configuration
+        config.update("jax_enable_x64", remember)
+
+
+typeof_gates_xzx = (qml.RX, qml.RZ, qml.RX, qml.GlobalPhase)
+single_qubit_decomps_xzx = [
+    (
+        np.array(
+            [
+                [-0.28829348 - 0.78829734j, 0.30364367 + 0.45085995j],
+                [0.53396245 - 0.10177564j, 0.76279558 - 0.35024096j],
+            ]
+        ),
+        typeof_gates_xzx,
+        [12.416147693665032, 1.3974974090935608, 11.448040119199066, 1.1759220332464762],
+    ),
+    (I, typeof_gates_xzx, [0, 0, 0, 0]),
+    (X, typeof_gates_xzx, [np.pi / 2, 0, np.pi / 2, -np.pi / 2]),
+    (Y, typeof_gates_xzx, [np.pi / 2, np.pi, 7 / 2 * np.pi, -np.pi / 2]),
+    (Z, typeof_gates_xzx, [0, np.pi, 0, -np.pi / 2]),
+    (H, typeof_gates_xzx, [np.pi / 2, np.pi / 2, np.pi / 2, -np.pi / 2]),
+    (
+        qml.QubitUnitary(qml.RZ.compute_matrix(np.array([np.pi, np.pi / 2])), wires=0).matrix(),
+        typeof_gates_xzx,
+        [[0, 0], [np.pi, np.pi / 2], [0, 0], [0, 0]],
+    ),
+    (
+        qml.Rot(np.array([1.2, 1.5]), np.array([1.2, 1.5]), np.array([1.2, 1.5]), wires=0).matrix(),
+        typeof_gates_xzx,
+        [
+            [0.63319625, 0.75125105],
+            [2.53416365, 3.03803113],
+            [11.93317436, 11.81511957],
+            [0, 0],
+        ],
+    ),
+]
+
+
+class TestQubitUnitaryXZXDecomposition:
+    """Test that the XZX decomposition is correct."""
+
+    @pytest.mark.parametrize("U,expected_gates,expected_params", single_qubit_decomps_xzx)
+    def test_xzx_decomposition(self, U, expected_gates, expected_params):
+        """Test that a one-qubit matrix in isolation is correctly decomposed."""
+
+        obtained_gates = one_qubit_decomposition(U, Wires("a"), "XZX", return_global_phase=True)
+        _run_assertions(U, expected_gates, expected_params, obtained_gates)
+
+    @pytest.mark.torch
+    @pytest.mark.parametrize("U,expected_gates,expected_params", single_qubit_decomps_xzx)
+    def test_xzx_decomposition_torch(self, U, expected_gates, expected_params):
+        """Test that a one-qubit operation in Torch is correctly decomposed."""
+
+        import torch
+
+        U = torch.tensor(U, dtype=torch.complex128)
+        obtained_gates = one_qubit_decomposition(U, Wires("a"), "XZX", return_global_phase=True)
+        _run_assertions(U, expected_gates, expected_params, obtained_gates)
+
+    @pytest.mark.tf
+    @pytest.mark.parametrize("U,expected_gates,expected_params", single_qubit_decomps_xzx)
+    def test_xzx_decomposition_tf(self, U, expected_gates, expected_params):
+        """Test that a one-qubit operation in Tensorflow is correctly decomposed."""
+
+        import tensorflow as tf
+
+        U = tf.Variable(U, dtype=tf.complex128)
+        obtained_gates = one_qubit_decomposition(U, Wires("a"), "XZX", return_global_phase=True)
+        _run_assertions(U, expected_gates, expected_params, obtained_gates)
 
-        self._run_assertions(U, expected_gates, expected_params, obtained_gates)
+    @pytest.mark.jax
+    @pytest.mark.parametrize("U,expected_gates,expected_params", single_qubit_decomps_xzx)
+    def test_xzx_decomposition_jax(self, U, expected_gates, expected_params):
+        """Test that a one-qubit operation in JAX is correctly decomposed."""
+
+        import 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)
+        obtained_gates = one_qubit_decomposition(U, Wires("a"), "XZX", return_global_phase=True)
+        _run_assertions(U, expected_gates, expected_params, obtained_gates)
+
+        # Reset the configuration
+        config.update("jax_enable_x64", remember)
 
 
 typeof_gates_zxz = (qml.RZ, qml.RX, qml.RZ, qml.GlobalPhase)
@@ -377,73 +439,40 @@ single_qubit_decomps_zxz = [
 class TestQubitUnitaryZXZDecomposition:
     """Test that the ZXZ decomposition is correct."""
 
-    def _run_assertions(self, U, expected_gates, expected_params, obtained_gates):
-        assert len(obtained_gates) == 4, "Incorrect number of gates"
-        # Check first 3 gates
-        for i in range(3):
-            assert isinstance(obtained_gates[i], expected_gates[i]), "Incorrect type of gate"
-            assert obtained_gates[i].wires == Wires("a"), "Incorrect wire"
-
-        # Check the global phase
-        assert isinstance(obtained_gates[3], expected_gates[3]), "Incorrect type of gate"
-        assert qml.math.allclose(
-            qml.math.unwrap(obtained_gates[3].parameters[:1]), expected_params[3]
-        ), "Incorrect global phase"
-        # Now we check the ZXZ rotation angles
-        assert qml.math.allclose(
-            [qml.math.unwrap(o.parameters)[0] for o in obtained_gates[:3]],
-            expected_params[:3],
-            atol=1e-7,
-        ), "Incorrect ZXZ rotation angles"
-
-        obtained_mat = reduce(
-            np.matmul, [op.matrix(wire_order=["a"]) for op in reversed(obtained_gates)]
-        )
-
-        if len(obtained_mat.shape) == 2:
-            U = [U]
-            obtained_mat = [obtained_mat]
-
-        assert all(
-            check_matrix_equivalence(curr_obtained_mat, curr_U, atol=1e-7)
-            for curr_obtained_mat, curr_U in zip(obtained_mat, qml.math.unwrap(U))
-        )
-
     @pytest.mark.parametrize("U,expected_gates,expected_params", single_qubit_decomps_zxz)
     def test_zxz_decomposition(self, U, expected_gates, expected_params):
         """Test that a one-qubit matrix in isolation is correctly decomposed."""
-        obtained_gates = one_qubit_decomposition(U, Wires("a"), "ZXZ", return_global_phase=True)
 
-        self._run_assertions(U, expected_gates, expected_params, obtained_gates)
+        obtained_gates = one_qubit_decomposition(U, Wires("a"), "ZXZ", return_global_phase=True)
+        _run_assertions(U, expected_gates, expected_params, obtained_gates)
 
     @pytest.mark.torch
     @pytest.mark.parametrize("U,expected_gates,expected_params", single_qubit_decomps_zxz)
     def test_zxz_decomposition_torch(self, U, expected_gates, expected_params):
         """Test that a one-qubit operation in Torch is correctly decomposed."""
+
         import torch
 
         U = torch.tensor(U, dtype=torch.complex128)
-
         obtained_gates = one_qubit_decomposition(U, Wires("a"), "ZXZ", return_global_phase=True)
-
-        self._run_assertions(U, expected_gates, expected_params, obtained_gates)
+        _run_assertions(U, expected_gates, expected_params, obtained_gates)
 
     @pytest.mark.tf
     @pytest.mark.parametrize("U,expected_gates,expected_params", single_qubit_decomps_zxz)
     def test_zxz_decomposition_tf(self, U, expected_gates, expected_params):
         """Test that a one-qubit operation in Tensorflow is correctly decomposed."""
+
         import tensorflow as tf
 
         U = tf.Variable(U, dtype=tf.complex128)
-
         obtained_gates = one_qubit_decomposition(U, Wires("a"), "ZXZ", return_global_phase=True)
-
-        self._run_assertions(U, expected_gates, expected_params, obtained_gates)
+        _run_assertions(U, expected_gates, expected_params, obtained_gates)
 
     @pytest.mark.jax
     @pytest.mark.parametrize("U,expected_gates,expected_params", single_qubit_decomps_zxz)
     def test_zxz_decomposition_jax(self, U, expected_gates, expected_params):
         """Test that a one-qubit operation in JAX is correctly decomposed."""
+
         import jax
 
         # Enable float64 support
@@ -453,153 +482,115 @@ class TestQubitUnitaryZXZDecomposition:
         config.update("jax_enable_x64", True)
 
         U = jax.numpy.array(U, dtype=jax.numpy.complex128)
-
         obtained_gates = one_qubit_decomposition(U, Wires("a"), "ZXZ", return_global_phase=True)
+        _run_assertions(U, expected_gates, expected_params, obtained_gates)
 
-        self._run_assertions(U, expected_gates, expected_params, obtained_gates)
+        # Restore the configuration
+        config.update("jax_enable_x64", remember)
 
 
-single_qubit_decomps = [
-    (
-        np.array(
-            [
-                [0, -9.831019270939975e-01 + 0.1830590094588862j],
-                [9.831019270939975e-01 + 0.1830590094588862j, 0],
-            ]
-        ),
-        "ZYZ",
-        typeof_gates_zyz,
-        [12.382273469673908, np.pi, 0.18409714468526372, 0],
-    ),
-    (
-        np.array(
-            [
-                [-0.28829348 - 0.78829734j, 0.30364367 + 0.45085995j],
-                [0.53396245 - 0.10177564j, 0.76279558 - 0.35024096j],
-            ]
-        ),
-        "XYX",
-        typeof_gates_xyx,
-        (
-            10.845351366405708,
-            1.3974974118006183,
-            0.45246583660683803,
-            1.1759220332464762,
-        ),
-    ),
-    (
-        np.array(
-            [
-                [0, -9.831019270939975e-01 + 0.1830590094588862j],
-                [9.831019270939975e-01 + 0.1830590094588862j, 0],
-            ]
-        ),
-        "ZXZ",
-        typeof_gates_zxz,
-        [10.811477142879012, np.pi, 1.7548934714801607, 0],
-    ),
-]
-
-
-rot_decomps = [
+single_qubit_decomps_rot = [
     # These will be decomposed to RZ
-    (I, "rot", qml.RZ, [0.0]),
-    (Z, "rot", qml.RZ, [np.pi]),
-    (S, "rot", qml.RZ, [np.pi / 2]),
-    (T, "rot", qml.RZ, [np.pi / 4]),
-    (qml.RZ(0.3, wires=0).matrix(), "rot", qml.RZ, [0.3]),
-    (qml.RZ(-0.5, wires=0).matrix(), "rot", qml.RZ, [-0.5]),
+    (I, qml.RZ, [0.0]),
+    (Z, qml.RZ, [np.pi]),
+    (S, qml.RZ, [np.pi / 2]),
+    (T, qml.RZ, [np.pi / 4]),
+    (qml.RZ(0.3, wires=0).matrix(), qml.RZ, [0.3]),
+    (qml.RZ(-0.5, wires=0).matrix(), qml.RZ, [-0.5]),
     # This will be decomposed to Rot
     (
         qml.Rot(0.2, 0.5, -0.3, wires=0).matrix(),
-        "rot",
         qml.Rot,
         [0.2, 0.5, -0.3],
     ),
 ]
 
 
-class TestOneQubitUnitaryDecomposition:
-    """Test that one qubit unitaries are correctly decomposed."""
+def _run_assertions_rot(U, expected_gate, expected_params, obtained_gates):
+    """Runs assertions for the rot decomposition."""
 
-    def _run_assertions(self, U, expected_gates, expected_params, obtained_gates):
-        assert len(obtained_gates) == 4, "Incorrect number of gates"
-        # Check first 3 gates
-        for i in range(3):
-            assert isinstance(obtained_gates[i], expected_gates[i]), "Incorrect type of gate"
-            assert obtained_gates[i].wires == Wires("a"), "Incorrect wire"
-
-        # Check the global phase
-        assert isinstance(obtained_gates[3], expected_gates[3]), "Incorrect type of gate"
-        assert qml.math.allclose(
-            qml.math.unwrap(obtained_gates[3].parameters[:1]), expected_params[3]
-        ), "Incorrect global phase"
-        # Now we check the rotation angles
-        assert qml.math.allclose(
-            [qml.math.unwrap(o.parameters)[0] for o in obtained_gates[:3]],
-            expected_params[:3],
-            atol=1e-7,
-        ), "Incorrect rotation angles"
-
-        obtained_mat = reduce(
-            np.matmul, [op.matrix(wire_order=["a"]) for op in reversed(obtained_gates)]
-        )
-
-        if len(obtained_mat.shape) == 2:
-            U = [U]
-            obtained_mat = [obtained_mat]
-
-        assert all(
-            check_matrix_equivalence(curr_obtained_mat, curr_U, atol=1e-7)
-            for curr_obtained_mat, curr_U in zip(obtained_mat, qml.math.unwrap(U))
-        )
-
-    @pytest.mark.parametrize("U,decomp_type,expected_gates,expected_params", single_qubit_decomps)
-    def test_one_qubit_decomposition(self, U, decomp_type, expected_gates, expected_params):
-        """Test that a one-qubit matrix in isolation is correctly decomposed."""
-        obtained_gates = one_qubit_decomposition(
-            U, Wires("a"), decomp_type, return_global_phase=True
-        )
+    assert len(obtained_gates) == 1, "Incorrect number of gates"
+    assert isinstance(obtained_gates[0], expected_gate), "Incorrect type of gate"
+    assert obtained_gates[0].wires == Wires("a"), "Incorrect wire"
+
+    # Now we check the rotation angles
+    assert qml.math.allclose(
+        qml.math.unwrap(obtained_gates[0].parameters),
+        expected_params,
+        atol=1e-7,
+    ), "Incorrect rotation angles"
+
+    obtained_mat = obtained_gates[0].matrix()
 
-        self._run_assertions(U, expected_gates, expected_params, obtained_gates)
+    if len(obtained_mat.shape) == 2:
+        U = [U]
+        obtained_mat = [obtained_mat]
+
+    assert all(
+        check_matrix_equivalence(curr_obtained_mat, curr_U, atol=1e-7)
+        for curr_obtained_mat, curr_U in zip(obtained_mat, qml.math.unwrap(U))
+    )
+
+
+class TestOneQubitRotDecomposition:
+    """Test that one qubit unitaries are correctly decomposed."""
 
     def test_one_qubit_decomposition_exception(self):
         """Test that exception for incorrect rotations argument is raised."""
+
         U = qml.Rot(1.2, 1.2, 1.2, wires=0).matrix()
         with pytest.raises(ValueError):
             one_qubit_decomposition(U, Wires("a"), "nonsense string", return_global_phase=True)
 
-    def _run_assertions_rot(self, U, expected_gate, expected_params, obtained_gates):
-        assert len(obtained_gates) == 1, "Incorrect number of gates"
-        assert isinstance(obtained_gates[0], expected_gate), "Incorrect type of gate"
-        assert obtained_gates[0].wires == Wires("a"), "Incorrect wire"
+    @pytest.mark.parametrize("U,expected_gate,expected_params", single_qubit_decomps_rot)
+    def test_one_qubit_decomposition_rot(self, U, expected_gate, expected_params):
+        """Test that a one-qubit matrix in isolation is correctly decomposed to RZ or Rot gate."""
 
-        # Now we check the rotation angles
-        assert qml.math.allclose(
-            qml.math.unwrap(obtained_gates[0].parameters),
-            expected_params,
-            atol=1e-7,
-        ), "Incorrect rotation angles"
+        obtained_gates = one_qubit_decomposition(U, Wires("a"), "rot", return_global_phase=True)
 
-        obtained_mat = obtained_gates[0].matrix()
+        _run_assertions_rot(U, expected_gate, expected_params, obtained_gates)
 
-        if len(obtained_mat.shape) == 2:
-            U = [U]
-            obtained_mat = [obtained_mat]
+    @pytest.mark.torch
+    @pytest.mark.parametrize("U,expected_gate,expected_params", single_qubit_decomps_rot)
+    def test_rot_decomposition_torch(self, U, expected_gate, expected_params):
+        """Test that a one-qubit operation in Torch is correctly decomposed."""
+
+        import torch
 
-        assert all(
-            check_matrix_equivalence(curr_obtained_mat, curr_U, atol=1e-7)
-            for curr_obtained_mat, curr_U in zip(obtained_mat, qml.math.unwrap(U))
-        )
+        U = torch.tensor(U, dtype=torch.complex128)
+        obtained_gates = one_qubit_decomposition(U, Wires("a"), "rot", return_global_phase=True)
+        _run_assertions_rot(U, expected_gate, expected_params, obtained_gates)
 
-    @pytest.mark.parametrize("U,decomp_type,expected_gate,expected_params", rot_decomps)
-    def test_one_qubit_decomposition_rot(self, U, decomp_type, expected_gate, expected_params):
-        """Test that a one-qubit matrix in isolation is correctly decomposed to RZ or Rot gate."""
-        obtained_gates = one_qubit_decomposition(
-            U, Wires("a"), decomp_type, return_global_phase=True
-        )
+    @pytest.mark.tf
+    @pytest.mark.parametrize("U,expected_gate,expected_params", single_qubit_decomps_rot)
+    def test_rot_decomposition_tf(self, U, expected_gate, expected_params):
+        """Test that a one-qubit operation in Tensorflow is correctly decomposed."""
+
+        import tensorflow as tf
+
+        U = tf.Variable(U, dtype=tf.complex128)
+        obtained_gates = one_qubit_decomposition(U, Wires("a"), "rot", return_global_phase=True)
+        _run_assertions_rot(U, expected_gate, expected_params, obtained_gates)
 
-        self._run_assertions_rot(U, expected_gate, expected_params, obtained_gates)
+    @pytest.mark.jax
+    @pytest.mark.parametrize("U,expected_gate,expected_params", single_qubit_decomps_rot)
+    def test_rot_decomposition_jax(self, U, expected_gate, expected_params):
+        """Test that a one-qubit operation in JAX is correctly decomposed."""
+
+        import 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)
+        obtained_gates = one_qubit_decomposition(U, Wires("a"), "rot", return_global_phase=True)
+        _run_assertions_rot(U, expected_gate, expected_params, obtained_gates)
+
+        # Restore the configuration
+        config.update("jax_enable_x64", remember)
 
 
 # Randomly generated set (scipy.unitary_group) of five U(4) operations.
@@ -1073,6 +1064,7 @@ class TestTwoQubitUnitaryDecomposition:
     @pytest.mark.parametrize("U", samples_3_cnots)
     def test_convert_to_su4(self, U):
         """Test a matrix in U(4) is correct converted to SU(4)."""
+
         U_su4 = _convert_to_su4(np.array(U))
 
         # Ensure the determinant is correct and the mats are equivalent up to a phase
@@ -1082,16 +1074,16 @@ class TestTwoQubitUnitaryDecomposition:
     @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."""
-        true_matrix = qml.math.kron(np.array(U_pair[0]), np.array(U_pair[1]))
 
+        true_matrix = qml.math.kron(np.array(U_pair[0]), np.array(U_pair[1]))
         A, B = _su2su2_to_tensor_products(true_matrix)
-
         assert check_matrix_equivalence(qml.math.kron(A, B), true_matrix)
 
     @pytest.mark.parametrize("wires", [[0, 1], ["a", "b"], [3, 2], ["c", 0]])
     @pytest.mark.parametrize("U", samples_3_cnots)
     def test_two_qubit_decomposition_3_cnots(self, U, wires):
         """Test that a two-qubit matrix using 3 CNOTs is correctly decomposed."""
+
         U = _convert_to_su4(np.array(U))
 
         assert _compute_num_cnots(U) == 3
@@ -1114,6 +1106,7 @@ class TestTwoQubitUnitaryDecomposition:
     @pytest.mark.parametrize("U", samples_2_cnots)
     def test_two_qubit_decomposition_2_cnots(self, U, wires):
         """Test that a two-qubit matrix using 2 CNOTs isolation is correctly decomposed."""
+
         U = _convert_to_su4(np.array(U))
 
         assert _compute_num_cnots(U) == 2
@@ -1134,6 +1127,7 @@ class TestTwoQubitUnitaryDecomposition:
     @pytest.mark.parametrize("U", samples_1_cnot)
     def test_two_qubit_decomposition_1_cnot(self, U, wires):
         """Test that a two-qubit matrix using one CNOT is correctly decomposed."""
+
         U = _convert_to_su4(np.array(U))
 
         assert _compute_num_cnots(U) == 1
@@ -1154,6 +1148,7 @@ class TestTwoQubitUnitaryDecomposition:
     @pytest.mark.parametrize("U_pair", samples_su2_su2)
     def test_two_qubit_decomposition_tensor_products(self, U_pair, wires):
         """Test that a two-qubit tensor product matrix is correctly decomposed."""
+
         U = _convert_to_su4(qml.math.kron(np.array(U_pair[0]), np.array(U_pair[1])))
 
         assert _compute_num_cnots(U) == 0
@@ -1179,6 +1174,7 @@ class TestTwoQubitUnitaryDecompositionInterfaces:
     @pytest.mark.parametrize("U", samples_3_cnots + samples_2_cnots + samples_1_cnot)
     def test_two_qubit_decomposition_torch(self, U, wires):
         """Test that a two-qubit operation in Torch is correctly decomposed."""
+
         import torch
 
         U = torch.tensor(U, dtype=torch.complex128)
@@ -1199,6 +1195,7 @@ class TestTwoQubitUnitaryDecompositionInterfaces:
     @pytest.mark.parametrize("U_pair", samples_su2_su2)
     def test_two_qubit_decomposition_tensor_products_torch(self, U_pair, wires):
         """Test that a two-qubit tensor product in Torch is correctly decomposed."""
+
         import torch
 
         U1 = torch.tensor(U_pair[0], dtype=torch.complex128)
@@ -1221,6 +1218,7 @@ class TestTwoQubitUnitaryDecompositionInterfaces:
     @pytest.mark.parametrize("U", samples_3_cnots + samples_2_cnots + samples_1_cnot)
     def test_two_qubit_decomposition_tf(self, U, wires):
         """Test that a two-qubit operation in Tensorflow is correctly decomposed."""
+
         import tensorflow as tf
 
         U = tf.Variable(U, dtype=tf.complex128)
@@ -1241,6 +1239,7 @@ class TestTwoQubitUnitaryDecompositionInterfaces:
     @pytest.mark.parametrize("U_pair", samples_su2_su2)
     def test_two_qubit_decomposition_tensor_products_tf(self, U_pair, wires):
         """Test that a two-qubit tensor product in Tensorflow is correctly decomposed."""
+
         import tensorflow as tf
 
         U1 = tf.Variable(U_pair[0], dtype=tf.complex128)
@@ -1263,6 +1262,7 @@ class TestTwoQubitUnitaryDecompositionInterfaces:
     @pytest.mark.parametrize("U", samples_3_cnots + samples_2_cnots + samples_1_cnot)
     def test_two_qubit_decomposition_jax(self, U, wires):
         """Test that a two-qubit operation in JAX is correctly decomposed."""
+
         import jax
         from jax.config import config
 
@@ -1282,11 +1282,15 @@ class TestTwoQubitUnitaryDecompositionInterfaces:
 
         assert check_matrix_equivalence(U, obtained_matrix, atol=1e-7)
 
+        # Restore config
+        config.update("jax_enable_x64", remember)
+
     @pytest.mark.jax
     @pytest.mark.parametrize("wires", [[0, 1], ["a", "b"], [3, 2], ["c", 0]])
     @pytest.mark.parametrize("U_pair", samples_su2_su2)
     def test_two_qubit_decomposition_tensor_products_jax(self, U_pair, wires):
         """Test that a two-qubit tensor product in JAX is correctly decomposed."""
+
         import jax
         from jax.config import config
 
@@ -1308,14 +1312,19 @@ class TestTwoQubitUnitaryDecompositionInterfaces:
 
         assert check_matrix_equivalence(U, obtained_matrix, atol=1e-7)
 
+        # Restore config
+        config.update("jax_enable_x64", remember)
+
     @pytest.mark.jax
     @pytest.mark.parametrize("wires", [[0, 1], ["a", "b"], [3, 2], ["c", 0]])
     @pytest.mark.parametrize("U", samples_3_cnots + samples_2_cnots + samples_1_cnot)
     def test_two_qubit_decomposition_jax_jit(self, U, wires):
         """Test that a two-qubit operation is correctly decomposed with JAX-JIT ."""
+
         import jax
         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)
@@ -1338,14 +1347,19 @@ class TestTwoQubitUnitaryDecompositionInterfaces:
 
         assert check_matrix_equivalence(U, jitted_matrix, atol=1e-7)
 
+        # Restore config
+        config.update("jax_enable_x64", remember)
+
     @pytest.mark.jax
     @pytest.mark.parametrize("wires", [[0, 1], ["a", "b"], [3, 2], ["c", 0]])
     @pytest.mark.parametrize("U_pair", samples_su2_su2)
     def test_two_qubit_decomposition_tensor_products_jax_jit(self, U_pair, wires):
         """Test that a two-qubit tensor product is correctly decomposed with JAX-JIT."""
+
         import jax
         from jax.config import config
 
+        remember = config.read("jax_enable_x64")
         config.update("jax_enable_x64", True)
 
         U1 = jax.numpy.array(U_pair[0], dtype=jax.numpy.complex128)
@@ -1369,3 +1383,6 @@ class TestTwoQubitUnitaryDecompositionInterfaces:
         jitted_matrix = jax.jit(wrapped_decomposition)(U)
 
         assert check_matrix_equivalence(U, jitted_matrix, atol=1e-7)
+
+        # Restore config
+        config.update("jax_enable_x64", remember)