From d0126df2b0ee688d1b386a7aa97daa5064c13a0f Mon Sep 17 00:00:00 2001
From: dwierichs <david.wierichs@xanadu.ai>
Date: Tue, 23 Jul 2024 21:45:06 +0200
Subject: [PATCH 01/30] new fuse_rot_angles function, incl test for evaluation
 and Jacobian, and test for singular points of Jacobian

---
 .../optimization/merge_rotations.py           |  10 +-
 .../optimization/optimization_utils.py        | 161 +++-------------
 .../test_optimization_utils.py                | 180 +++++++++++++-----
 3 files changed, 171 insertions(+), 180 deletions(-)

diff --git a/pennylane/transforms/optimization/merge_rotations.py b/pennylane/transforms/optimization/merge_rotations.py
index 2c8178c1997..30313156425 100644
--- a/pennylane/transforms/optimization/merge_rotations.py
+++ b/pennylane/transforms/optimization/merge_rotations.py
@@ -205,11 +205,11 @@ def stop_at(obj):
                     current_gate.__class__(*cumulative_angles, wires=current_gate.wires)
                 )
         else:
-            if not allclose(cumulative_angles, zeros(len(cumulative_angles)), atol=atol, rtol=0):
-                with QueuingManager.stop_recording():
-                    new_operations.append(
-                        current_gate.__class__(*cumulative_angles, wires=current_gate.wires)
-                    )
+            # if not allclose(cumulative_angles, zeros(len(cumulative_angles)), atol=atol, rtol=0):
+            with QueuingManager.stop_recording():
+                new_operations.append(
+                    current_gate.__class__(*cumulative_angles, wires=current_gate.wires)
+                )
 
         # Remove the first 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 cf390d8aab3..28966b2d391 100644
--- a/pennylane/transforms/optimization/optimization_utils.py
+++ b/pennylane/transforms/optimization/optimization_utils.py
@@ -13,13 +13,7 @@
 # limitations under the License.
 """Utility functions for circuit optimization."""
 # pylint: disable=too-many-return-statements,import-outside-toplevel
-from functools import partial
-
-import numpy as np
-
-from pennylane.math import abs as math_abs
-from pennylane.math import allclose, arccos, arctan2, cos, get_interface, is_abstract, sin, stack
-from pennylane.math import sum as math_sum
+from pennylane.math import arccos, arctan2, asarray, cos, sin, sqrt, stack
 from pennylane.ops.identity import GlobalPhase
 from pennylane.wires import Wires
 
@@ -47,109 +41,9 @@ 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]))
-
-    return stack([qw, qx, qy, qz])
-
-
-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]
-
-    return stack([qw, qx, qy, qz])
-
-
-def _singular_quat_to_zyz(qw, qx, qy, qz, y_arg, abstract_jax=False):
-    """Compute the ZYZ angles for the singular case of qx = qy = 0"""
-    # pylint: disable=too-many-arguments
-    z1_arg1 = 2 * (qx * qy + qz * qw)
-    z1_arg2 = 1 - 2 * (qx**2 + qz**2)
-
-    if abstract_jax:
-        from jax.lax import cond
-
-        return cond(
-            y_arg > 0,
-            lambda z1_arg1, z1_arg2: stack([arctan2(z1_arg1, z1_arg2), 0.0, 0.0]),
-            lambda z1_arg1, z1_arg2: stack([-arctan2(z1_arg1, z1_arg2), np.pi, 0.0]),
-            z1_arg1,
-            z1_arg2,
-        )
-
-    if y_arg > 0:
-        z1 = arctan2(z1_arg1, z1_arg2)
-        y = z2 = 0.0
-    else:
-        z1 = -arctan2(z1_arg1, z1_arg2)
-        y = np.pi
-        z2 = 0.0
-    return stack([z1, y, z2])
-
-
-def _regular_quat_to_zyz(qw, qx, qy, qz, y_arg):
-    """Compute the ZYZ angles for the regular case (qx != 0 or qy != 0)"""
-    z1_arg1 = 2 * (qy * qz - qw * qx)
-    z1_arg2 = 2 * (qx * qz + qw * qy)
-    z1 = arctan2(z1_arg1, z1_arg2)
-
-    y = arccos(y_arg)
-
-    z2_arg1 = 2 * (qy * qz + qw * qx)
-    z2_arg2 = 2 * (qw * qy - qx * qz)
-    z2 = arctan2(z2_arg1, z2_arg2)
-
-    return stack([z1, y, z2])
-
-
-def _fuse(angles_1, angles_2, abstract_jax=False):
-    """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))
-
-    # Convert the product back into the angles fed to Rot
-    y_arg = 1 - 2 * (qx**2 + qy**2)
-
-    # Require special treatment of the case qx = qy = 0. Note that we have to check
-    # for "greater than" as well, because of imprecisions
-    if abstract_jax:
-        from jax.lax import cond
-
-        return cond(
-            math_abs(y_arg) >= 1,
-            partial(_singular_quat_to_zyz, abstract_jax=True),
-            _regular_quat_to_zyz,
-            qw,
-            qx,
-            qy,
-            qz,
-            y_arg,
-        )
-
-    # Require special treatment of the case qx = qy = 0
-    if abs(y_arg) >= 1:  # Have to check for "greater than" as well, because of imprecisions
-        return _singular_quat_to_zyz(qw, qx, qy, qz, y_arg)
-    return _regular_quat_to_zyz(qw, qx, qy, qz, y_arg)
-
-
-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.
+def fuse_rot_angles(angles1, angles2):
+    """Compute the set of rotation angles that is equivalent to performing
+    two successive ``qml.Rot`` operations.
 
     The ``qml.Rot`` operation represents the most general single-qubit operation.
     Two such operations can be fused into a new operation, however the angular dependence
@@ -162,31 +56,38 @@ def fuse_rot_angles(angles_1, angles_2):
     Returns:
         array[float]: Rotation angles for a single ``qml.Rot`` operation that
         implements the same operation as the two sets of input angles.
+
+    .. note::
+
+        The output angles are not always defined uniquely because Euler angles are not
+        defined uniquely for some rotations. ``fuse_rot_angles`` makes a particular
+        choice in this case.
+
+    .. warning::
+
+        This function is not differentiable everywhere. It has singularities for specific
+        input values, where the derivative will be ``nan``.
+
     """
+    phi1, theta1, omega1 = asarray(angles1).T
+    phi2, theta2, omega2 = asarray(angles2).T
+    c1, c2, s1, s2 = cos(theta1 / 2), cos(theta2 / 2), sin(theta1 / 2), sin(theta2 / 2)
 
-    # Check if we are tracing; if so, use the special conditionals
-    if is_abstract(angles_1) or is_abstract(angles_2):
-        interface = get_interface(angles_1, angles_2)
+    mag = sqrt(c1**2 * c2**2 + s1**2 * s2**2 - 2 * c1 * c2 * s1 * s2 * cos(omega1 + phi2))
+    theta_f = 2 * arccos(mag)
 
-        # 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
+    alpha1, beta1 = (phi1 + omega1) / 2, (phi1 - omega1) / 2
+    alpha2, beta2 = (phi2 + omega2) / 2, (phi2 - omega2) / 2
 
-            return cond(
-                allclose(angles_1[1], 0.0) * allclose(angles_2[1], 0.0),
-                _no_fuse,
-                partial(_fuse, abstract_jax=True),
-                angles_1,
-                angles_2,
-            )
+    alpha_arg1 = -c1 * c2 * sin(alpha1 + alpha2) - s1 * s2 * sin(beta2 - beta1)
+    alpha_arg2 = c1 * c2 * cos(alpha1 + alpha2) - s1 * s2 * cos(beta2 - beta1)
+    alpha_f = -1 * arctan2(alpha_arg1, alpha_arg2)
 
-    # 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)
+    beta_arg1 = -c1 * s2 * sin(alpha1 + beta2) + s1 * c2 * sin(alpha2 - beta1)
+    beta_arg2 = c1 * s2 * cos(alpha1 + beta2) + s1 * c2 * cos(alpha2 - beta1)
+    beta_f = -1 * arctan2(beta_arg1, beta_arg2)
 
-    return _fuse(angles_1, angles_2)
+    return stack([alpha_f + beta_f, theta_f, alpha_f - beta_f], axis=-1)
 
 
 def _fuse_global_phases(operations):
@@ -206,5 +107,5 @@ def _fuse_global_phases(operations):
         else:
             fused_ops.append(op)
 
-    fused_ops.append(GlobalPhase(math_sum(op.data[0] for op in global_ops)))
+    fused_ops.append(GlobalPhase(sum(op.data[0] for op in global_ops)))
     return fused_ops
diff --git a/tests/transforms/test_optimization/test_optimization_utils.py b/tests/transforms/test_optimization/test_optimization_utils.py
index 693b7561e5d..dbe20f503ef 100644
--- a/tests/transforms/test_optimization/test_optimization_utils.py
+++ b/tests/transforms/test_optimization/test_optimization_utils.py
@@ -16,17 +16,13 @@
 """
 # pylint: disable=too-few-public-methods
 
+from itertools import product
+
 import pytest
-from utils import check_matrix_equivalence
 
 import pennylane as qml
 from pennylane import numpy as np
-from pennylane.transforms.optimization.optimization_utils import (
-    _quaternion_product,
-    _zyz_to_quat,
-    find_next_gate,
-    fuse_rot_angles,
-)
+from pennylane.transforms.optimization.optimization_utils import find_next_gate, fuse_rot_angles
 
 sample_op_list = [
     qml.Hadamard(wires="a"),
@@ -56,42 +52,6 @@ def test_find_next_gate(self, wires, op_list, next_gate_idx):
 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]),
-        ],
-    )
-    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)
-
-    @pytest.mark.parametrize(
-        ("angles_1", "angles_2", "expected_quat"),
-        # Examples generated at https://www.vcalc.com/wiki/vCalc/Quaternion+Multiplication
-        [
-            ([0.0, 0.0, 0.0, 0.0], [0.1, 0.2, 0.3, 0.4], [0.0, 0.0, 0.0, 0.0]),
-            ([1.0, 0.0, 0.0, 0.0], [1.0, 0.0, 0.0, 0.0], [1.0, 0.0, 0.0, 0.0]),
-            ([1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0], [-60.0, 12.0, 30.0, 24.0]),
-            ([0.1, 0.0, -0.2, 0.15], [1.0, 0.05, 1.65, -0.25], [0.4675, -0.1925, -0.0275, 0.135]),
-            ([0.1, 0.0, 0.0, 0.15], [1.0, 0.0, 0.0, -0.25], [0.1375, 0.0, 0.0, 0.125]),
-        ],
-    )
-    def test_quaternion_product(self, angles_1, angles_2, expected_quat):
-        """Test that products of quaternions produce expected results."""
-        obtained_quat = _quaternion_product(angles_1, angles_2)
-        assert qml.math.allclose(obtained_quat, expected_quat)
-
-    @pytest.mark.autograd
     @pytest.mark.parametrize(
         ("angles_1", "angles_2"),
         [
@@ -117,7 +77,7 @@ def test_quaternion_product(self, angles_1, angles_2, expected_quat):
             ([0.9, np.pi / 2, np.pi / 2], [-np.pi / 2, -np.pi / 2, 0.0]),
         ],
     )
-    def test_full_rot_fusion_autograd(self, angles_1, angles_2):
+    def test_full_rot_fusion(self, angles_1, angles_2):
         """Test that the fusion of two Rot gates has the same effect as
         applying the Rots sequentially."""
 
@@ -130,4 +90,134 @@ def original_ops():
         fused_angles = fuse_rot_angles(angles_1, angles_2)
         matrix_obtained = qml.Rot(*fused_angles, wires=0).matrix()
 
-        assert check_matrix_equivalence(matrix_expected, matrix_obtained)
+        assert qml.math.allclose(matrix_expected, matrix_obtained)
+
+    @pytest.mark.slow
+    def test_full_rot_fusion_special_angles(self):
+        """Test the rotation angle fusion on special multiples of pi/2.
+        Also tests that fuse_rot_angles works with batching/broadcasting.
+        Do not change the test to non-broadcasted evaluation, as this will
+        increase the runtime significantly."""
+
+        special_points = np.array([3 / 2, 1, 1 / 2, 0, -1 / 2, -1, -3 / 2]) * np.pi
+        special_angles = np.array(list(product(special_points, repeat=6))).reshape((-1, 2, 3))
+        angles_1, angles_2 = np.transpose(special_angles, (1, 0, 2))
+
+        def original_ops():
+            qml.Rot(*angles_1.T, wires=0)  # Transpose to bring size-3 axis to front
+            qml.Rot(*angles_2.T, wires=0)  # Transpose to bring size-3 axis to front
+
+        matrix_expected = qml.matrix(original_ops, [0])()  # pylint:disable=too-many-function-args
+
+        fused_angles = fuse_rot_angles(angles_1, angles_2)
+        matrix_obtained = qml.Rot(
+            *fused_angles.T, wires=0
+        ).matrix()  # Transpose to bring size-3 axis to front
+
+        assert qml.math.allclose(matrix_expected, matrix_obtained)
+
+    @pytest.mark.slow
+    @pytest.mark.jax
+    def test_full_rot_fusion_jacobian(self):
+        """Test the Jacobian of the rotation angle fusion. Uses batching for performance reasons.
+        For known sources of singularities, the Jacobian is checked to indeed return NaN.
+        These sources are related to the absolute value of the upper left entry of the matrix product:
+         - If it is 1, the derivative of arccos becomes infinite (evaluated at 1), and
+         - if its square is 0, the derivative of sqrt becomes infinite (evaluated at 0).
+        """
+        import jax
+
+        special_points = np.array([3 / 2, 1, 1 / 2, 0, -1 / 2, -1, -3 / 2]) * np.pi
+        special_angles = np.array(list(product(special_points, repeat=6))).reshape((-1, 2, 3))
+        random_angles = np.random.random((1000, 2, 3))
+        all_angles = jax.numpy.concatenate([special_angles, random_angles], dtype=complex)
+
+        def mat_from_prod(angles):
+            def original_ops():
+                angles1, angles2 = angles[..., 0, :], angles[..., 1, :]
+                qml.Rot(angles1[..., 0], angles1[..., 1], angles1[..., 2], wires=0)
+                qml.Rot(angles2[..., 0], angles2[..., 1], angles2[..., 2], wires=0)
+
+            return qml.matrix(original_ops, [0])()  # pylint:disable=too-many-function-args
+
+        def mat_from_fuse(angles):
+            angles1, angles2 = angles[..., 0, :], angles[..., 1, :]
+            fused_angles = fuse_rot_angles(angles1, angles2)
+            return qml.Rot(*fused_angles.T, wires=0).matrix()
+
+        # Need holomorphic derivatives because the output matrices are complex-valued
+        jac_from_prod = jax.vmap(jax.jacobian(mat_from_prod, holomorphic=True))(all_angles)
+        jac_from_fuse = jax.vmap(jax.jacobian(mat_from_fuse, holomorphic=True))(all_angles)
+
+        # expected failures based on the sources mentioned in the docstring above.
+        thetas = all_angles[..., 1].T
+        (c1, c2), (s1, s2) = np.cos(thetas / 2), np.sin(thetas / 2)
+        omega1 = all_angles[:, 0, 2]
+        phi2 = all_angles[:, 1, 0]
+        # squared absolute value of the relevant entry of the product of the two rotation matrices
+        pre_mag = c1**2 * c2**2 + s1**2 * s2**2 - 2 * c1 * c2 * s1 * s2 * np.cos(omega1 + phi2)
+        # Compute condition for the two error sources combined
+        error_sources = (np.abs(pre_mag - 1) < 1e-12) + (pre_mag == 0j)
+
+        assert qml.math.allclose(jac_from_prod[~error_sources], jac_from_fuse[~error_sources])
+        assert qml.math.all(
+            qml.math.any(qml.math.isnan(jac_from_fuse[error_sources]), axis=[1, 2, 3, 4])
+        )
+
+        """
+        mag = np.sqrt(pre_mag)
+
+        alpha1, beta1 = (phi1 + omega1) / 2, (phi1 - omega1) / 2
+        alpha2, beta2 = (phi2 + omega2) / 2, (phi2 - omega2) / 2
+
+        alpha_arg1 = -c1 * c2 * np.sin(alpha1 + alpha2) - s1 * s2 * np.sin(beta2 - beta1)
+        alpha_arg2 = c1 * c2 * np.cos(alpha1 + alpha2) - s1 * s2 * np.cos(beta2 - beta1)
+        beta_arg1 = -c1 * s2 * np.sin(alpha1 + beta2) + s1 * c2 * np.sin(alpha2 - beta1)
+        beta_arg2 = c1 * s2 * np.cos(alpha1 + beta2) + s1 * c2 * np.cos(alpha2 - beta1)
+
+        def partial(all_angles):
+            phi1, theta1, omega1 = all_angles[..., 0, :].T
+            phi2, theta2, omega2 = all_angles[..., 1, :].T
+            c1, c2, s1, s2 = qml.math.cos(theta1 / 2), qml.math.cos(theta2 / 2), qml.math.sin(theta1 / 2), qml.math.sin(theta2 / 2)
+            mag = qml.math.sqrt(c1 ** 2 * c2 ** 2 + s1 ** 2 * s2 ** 2 - 2 * c1 * c2 * s1 * s2 * qml.math.cos(omega1 + phi2))
+
+            #alpha1, beta1 = (phi1 + omega1) / 2, (phi1 - omega1) / 2
+            #alpha2, beta2 = (phi2 + omega2) / 2, (phi2 - omega2) / 2
+
+            #alpha_arg1 = -c1 * c2 * np.sin(alpha1 + alpha2) - s1 * s2 * np.sin(beta2 - beta1)
+            #alpha_arg2 = c1 * c2 * np.cos(alpha1 + alpha2) - s1 * s2 * np.cos(beta2 - beta1)
+            #beta_arg1 = -c1 * s2 * np.sin(alpha1 + beta2) + s1 * c2 * np.sin(alpha2 - beta1)
+            #beta_arg2 = c1 * s2 * np.cos(alpha1 + beta2) + s1 * c2 * np.cos(alpha2 - beta1)
+            #print(mag)
+            return mag
+
+        jac_partial = jax.vmap(jax.jacobian(partial))(all_angles.real)
+
+        expected_singular_ids = np.where(
+            np.any([
+                (np.abs(pre_mag - 1) < 1e-12),
+                (pre_mag == 0j),
+                #np.isnan(mag),
+                #(np.abs(alpha_arg1) + np.abs(alpha_arg1) < 1e-10),
+                #(np.abs(beta_arg1) + np.abs(beta_arg1) < 1e-10)
+            ], axis=0)
+        )
+        pre_mag2 = c1 ** 2 * s2 ** 2 + s1 ** 2 * c2 ** 2 + 2 * c1 * c2 * s1 * s2 * qml.math.cos(omega1 + phi2)
+        print(pre_mag2[np.where(pre_mag==0j)[0]])
+        for idx in np.where(qml.math.any(qml.math.isnan(jac_from_fuse), axis=[1, 2, 3, 4]))[0]:
+            if idx not in expected_singular_ids[0]:
+                weird_idx = idx
+                break
+        print(weird_idx)
+        print(jac_from_fuse[weird_idx])
+        print(all_angles[weird_idx])
+        print(f"{pre_mag[weird_idx]:.20f}")
+        print(jax.jacobian(qml.math.sqrt,holomorphic=True)(pre_mag[weird_idx]))
+        print(jax.jacobian(qml.math.sqrt)(pre_mag.real[weird_idx]))
+        print(len(all_angles))
+        print(len(expected_singular_ids[0]))
+        print(len(np.where(qml.math.any(qml.math.isnan(jac_from_fuse), axis=[1, 2, 3, 4]))[0]))
+        print(jac_partial.shape)
+        print(len(np.where(qml.math.any(qml.math.isnan(jac_partial), axis=[1, 2]))[0]))
+        print(np.unique(mag[np.array(np.where(qml.math.any(qml.math.isnan(jac_partial), axis=[1, 2]))[0])]))
+        """

From 961d6ca251d6fab4e31e524a5f4b27506bfb844c Mon Sep 17 00:00:00 2001
From: dwierichs <david.wierichs@xanadu.ai>
Date: Tue, 23 Jul 2024 21:47:36 +0200
Subject: [PATCH 02/30] changelog

---
 doc/releases/changelog-dev.md | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 8559de76b3b..6b0c0f737db 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -28,6 +28,9 @@
 
 <h3>Improvements 🛠</h3>
 
+* `fuse_rot_angles` respects now the global phase of the combined rotations.
+  [(#-___)](https://github.com/PennyLaneAI/pennylane/pull/-___)
+
 * `StateMP.process_state` defines rules in `cast_to_complex` for complex casting, avoiding a superfluous state vector copy in Lightning simulations
   [(#5995)](https://github.com/PennyLaneAI/pennylane/pull/5995)
 
@@ -124,6 +127,9 @@
 
 <h3>Bug fixes 🐛</h3>
 
+* `fuse_rot_angles` no longer returns wrong derivatives at singular points but returns NaN.
+  [(#____)](https://github.com/PennyLaneAI/pennylane/pull/____)
+
 * `CircuitGraph` can now handle circuits with the same operation instance occuring multiple times.
   [(#5907)](https://github.com/PennyLaneAI/pennylane/pull/5907)
 
@@ -157,4 +163,5 @@ Christina Lee,
 William Maxwell,
 Vincent Michaud-Rioux,
 Mudit Pandey,
-Erik Schultheis.
+Erik Schultheis,
+David Wierichs.

From 5c231abfbfde2495a7ebf31882c97c58b86a84e7 Mon Sep 17 00:00:00 2001
From: dwierichs <david.wierichs@xanadu.ai>
Date: Tue, 23 Jul 2024 21:48:38 +0200
Subject: [PATCH 03/30] pr number

---
 doc/releases/changelog-dev.md | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 6b0c0f737db..40be2eb3c08 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -29,7 +29,7 @@
 <h3>Improvements 🛠</h3>
 
 * `fuse_rot_angles` respects now the global phase of the combined rotations.
-  [(#-___)](https://github.com/PennyLaneAI/pennylane/pull/-___)
+  [(#6031)](https://github.com/PennyLaneAI/pennylane/pull/6031)
 
 * `StateMP.process_state` defines rules in `cast_to_complex` for complex casting, avoiding a superfluous state vector copy in Lightning simulations
   [(#5995)](https://github.com/PennyLaneAI/pennylane/pull/5995)
@@ -128,7 +128,7 @@
 <h3>Bug fixes 🐛</h3>
 
 * `fuse_rot_angles` no longer returns wrong derivatives at singular points but returns NaN.
-  [(#____)](https://github.com/PennyLaneAI/pennylane/pull/____)
+  [(#6031)](https://github.com/PennyLaneAI/pennylane/pull/6031)
 
 * `CircuitGraph` can now handle circuits with the same operation instance occuring multiple times.
   [(#5907)](https://github.com/PennyLaneAI/pennylane/pull/5907)

From 8cb80a5938122ed0da941751693c822628997625 Mon Sep 17 00:00:00 2001
From: dwierichs <david.wierichs@xanadu.ai>
Date: Tue, 23 Jul 2024 21:59:41 +0200
Subject: [PATCH 04/30] undo merge_rotations change

---
 pennylane/transforms/optimization/merge_rotations.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/pennylane/transforms/optimization/merge_rotations.py b/pennylane/transforms/optimization/merge_rotations.py
index 30313156425..2c8178c1997 100644
--- a/pennylane/transforms/optimization/merge_rotations.py
+++ b/pennylane/transforms/optimization/merge_rotations.py
@@ -205,11 +205,11 @@ def stop_at(obj):
                     current_gate.__class__(*cumulative_angles, wires=current_gate.wires)
                 )
         else:
-            # if not allclose(cumulative_angles, zeros(len(cumulative_angles)), atol=atol, rtol=0):
-            with QueuingManager.stop_recording():
-                new_operations.append(
-                    current_gate.__class__(*cumulative_angles, wires=current_gate.wires)
-                )
+            if not allclose(cumulative_angles, zeros(len(cumulative_angles)), atol=atol, rtol=0):
+                with QueuingManager.stop_recording():
+                    new_operations.append(
+                        current_gate.__class__(*cumulative_angles, wires=current_gate.wires)
+                    )
 
         # Remove the first gate from the working list
         list_copy.pop(0)

From 102a7a4c724c3a558edc39f97ad766ef30a5277e Mon Sep 17 00:00:00 2001
From: dwierichs <david.wierichs@xanadu.ai>
Date: Tue, 23 Jul 2024 22:59:57 +0200
Subject: [PATCH 05/30] ml interfaces tests, precision comment

---
 .../optimization/optimization_utils.py        |  14 +-
 .../test_optimization_utils.py                | 191 +++++++++---------
 2 files changed, 106 insertions(+), 99 deletions(-)

diff --git a/pennylane/transforms/optimization/optimization_utils.py b/pennylane/transforms/optimization/optimization_utils.py
index 28966b2d391..7c743e4f6bc 100644
--- a/pennylane/transforms/optimization/optimization_utils.py
+++ b/pennylane/transforms/optimization/optimization_utils.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 """Utility functions for circuit optimization."""
 # pylint: disable=too-many-return-statements,import-outside-toplevel
-from pennylane.math import arccos, arctan2, asarray, cos, sin, sqrt, stack
+from pennylane.math import arccos, arctan2, asarray, cos, sin, sqrt, stack, transpose
 from pennylane.ops.identity import GlobalPhase
 from pennylane.wires import Wires
 
@@ -68,12 +68,18 @@ def fuse_rot_angles(angles1, angles2):
         This function is not differentiable everywhere. It has singularities for specific
         input values, where the derivative will be ``nan``.
 
+    .. warning::
+
+        This function is numerically unstable at singular points. It is recommended to use
+        it with 64-bit floating point precision.
+
     """
-    phi1, theta1, omega1 = asarray(angles1).T
-    phi2, theta2, omega2 = asarray(angles2).T
+    phi1, theta1, omega1 = transpose(asarray(angles1))
+    phi2, theta2, omega2 = transpose(asarray(angles2))
     c1, c2, s1, s2 = cos(theta1 / 2), cos(theta2 / 2), sin(theta1 / 2), sin(theta2 / 2)
 
     mag = sqrt(c1**2 * c2**2 + s1**2 * s2**2 - 2 * c1 * c2 * s1 * s2 * cos(omega1 + phi2))
+    print(f"{mag=}")
     theta_f = 2 * arccos(mag)
 
     alpha1, beta1 = (phi1 + omega1) / 2, (phi1 - omega1) / 2
@@ -81,10 +87,12 @@ def fuse_rot_angles(angles1, angles2):
 
     alpha_arg1 = -c1 * c2 * sin(alpha1 + alpha2) - s1 * s2 * sin(beta2 - beta1)
     alpha_arg2 = c1 * c2 * cos(alpha1 + alpha2) - s1 * s2 * cos(beta2 - beta1)
+    print(f"{alpha_arg1=}, {alpha_arg2=}")
     alpha_f = -1 * arctan2(alpha_arg1, alpha_arg2)
 
     beta_arg1 = -c1 * s2 * sin(alpha1 + beta2) + s1 * c2 * sin(alpha2 - beta1)
     beta_arg2 = c1 * s2 * cos(alpha1 + beta2) + s1 * c2 * cos(alpha2 - beta1)
+    print(f"{beta_arg1=}, {beta_arg2=}")
     beta_f = -1 * arctan2(beta_arg1, beta_arg2)
 
     return stack([alpha_f + beta_f, theta_f, alpha_f - beta_f], axis=-1)
diff --git a/tests/transforms/test_optimization/test_optimization_utils.py b/tests/transforms/test_optimization/test_optimization_utils.py
index dbe20f503ef..e3a75b7d5d1 100644
--- a/tests/transforms/test_optimization/test_optimization_utils.py
+++ b/tests/transforms/test_optimization/test_optimization_utils.py
@@ -52,41 +52,102 @@ def test_find_next_gate(self, wires, op_list, next_gate_idx):
 class TestRotGateFusion:
     """Test that utility functions for fusing two qml.Rot gates function as expected."""
 
-    @pytest.mark.parametrize(
-        ("angles_1", "angles_2"),
-        [
-            ([0.15, 0.25, -0.90], [-0.5, 0.25, 1.3]),
-            ([0.0, 0.0, 0.0], [0.0, 0.0, 0.0]),
-            ([0.1, 0.2, 0.3], [0.0, 0.0, 0.0]),
-            ([0.0, 0.0, 0.0], [0.2, 0.4, -0.5]),
-            ([0.15, 0.25, -0.90], [-0.15, -0.25, 0.9]),
-            ([0.05, -1.34, 4.12], [-0.8, 0.2, 0.12]),
-            ([0.05, 0.0, 4.12], [-0.8, 0.2, 0.12]),
-            ([0.05, -1.34, 0.0], [-0.8, 0.2, 0.12]),
-            ([0.05, 0.0, 0.1], [-0.2, 0.0, 0.12]),
-            ([0.05, 0.0, 0.0], [0.0, 0.0, 0.12]),
-            ([0.05, 0.0, 0.0], [0.0, 0.0, -0.05]),
-            ([0.05, 0.0, 0.0], [-0.05, 0.0, 0.0]),
-            ([0.0, 0.0, 0.05], [-0.05, 0.0, 0.0]),
-            ([0.0, 0.0, 0.05], [0.0, 0.0, -0.05]),
-            ([0.05, 0.2, 0.0], [0.0, -0.6, 0.12]),
-            ([0.05, -1.34, 4.12], [0.0, 0.2, 0.12]),
-            ([0.05, -1.34, 4.12], [0.3, 0.0, 0.12]),
-            ([np.pi, np.pi / 2, 0.0], [0.0, -np.pi / 2, 0.0]),
-            ([0.9, np.pi / 2, 0.0], [0.0, -np.pi / 2, 0.0]),
-            ([0.9, np.pi / 2, np.pi / 2], [-np.pi / 2, -np.pi / 2, 0.0]),
-        ],
-    )
-    def test_full_rot_fusion(self, angles_1, angles_2):
-        """Test that the fusion of two Rot gates has the same effect as
-        applying the Rots sequentially."""
-
+    generic_test_angles = [
+        ([0.15, 0.25, -0.90], [-0.5, 0.25, 1.3]),
+        ([0.0, 0.0, 0.0], [0.0, 0.0, 0.0]),
+        ([0.1, 0.2, 0.3], [0.0, 0.0, 0.0]),
+        ([0.0, 0.0, 0.0], [0.2, 0.4, -0.5]),
+        ([0.15, 0.25, -0.90], [-0.15, -0.25, 0.9]),
+        ([0.05, -1.34, 4.12], [-0.8, 0.2, 0.12]),
+        ([0.05, 0.0, 4.12], [-0.8, 0.2, 0.12]),
+        ([0.05, -1.34, 0.0], [-0.8, 0.2, 0.12]),
+        ([0.05, 0.0, 0.1], [-0.2, 0.0, 0.12]),
+        ([0.05, 0.0, 0.0], [0.0, 0.0, 0.12]),
+        ([0.05, 0.0, 0.0], [0.0, 0.0, -0.05]),
+        ([0.05, 0.0, 0.0], [-0.05, 0.0, 0.0]),
+        ([0.0, 0.0, 0.05], [-0.05, 0.0, 0.0]),
+        ([0.0, 0.0, 0.05], [0.0, 0.0, -0.05]),
+        ([0.05, 0.2, 0.0], [0.0, -0.6, 0.12]),
+        ([0.05, -1.34, 4.12], [0.0, 0.2, 0.12]),
+        ([0.05, -1.34, 4.12], [0.3, 0.0, 0.12]),
+        ([np.pi, np.pi / 2, 0.0], [0.0, -np.pi / 2, 0.0]),
+        ([0.9, np.pi / 2, 0.0], [0.0, -np.pi / 2, 0.0]),
+        ([0.9, np.pi / 2, np.pi / 2], [-np.pi / 2, -np.pi / 2, 0.0]),
+    ]
+
+    @staticmethod
+    def mat_from_prod(angles_1, angles_2):
         def original_ops():
             qml.Rot(*angles_1, wires=0)
             qml.Rot(*angles_2, wires=0)
 
-        matrix_expected = qml.matrix(original_ops, [0])()  # pylint:disable=too-many-function-args
+        return qml.matrix(original_ops, [0])()  # pylint:disable=too-many-function-args
+
+    @pytest.mark.parametrize("angles_1, angles_2", generic_test_angles)
+    def test_full_rot_fusion_numpy(self, angles_1, angles_2):
+        """Test that the fusion of two Rot gates has the same effect as
+        applying the Rots sequentially."""
+
+        matrix_expected = self.mat_from_prod(angles_1, angles_2)
+        fused_angles = fuse_rot_angles(angles_1, angles_2)
+        matrix_obtained = qml.Rot(*fused_angles, wires=0).matrix()
+
+        assert qml.math.allclose(matrix_expected, matrix_obtained)
+
+    @pytest.mark.autograd
+    @pytest.mark.parametrize("angles_1, angles_2", generic_test_angles)
+    def test_full_rot_fusion_autograd(self, angles_1, angles_2):
+        """Test that the fusion of two Rot gates has the same effect as
+        applying the Rots sequentially, in Autograd."""
 
+        angles_1, angles_2 = qml.numpy.array(angles_1), qml.numpy.array(angles_1)
+        matrix_expected = self.mat_from_prod(angles_1, angles_2)
+        fused_angles = fuse_rot_angles(angles_1, angles_2)
+        matrix_obtained = qml.Rot(*fused_angles, wires=0).matrix()
+
+        assert qml.math.allclose(matrix_expected, matrix_obtained)
+
+    @pytest.mark.tf
+    @pytest.mark.parametrize("angles_1, angles_2", generic_test_angles)
+    def test_full_rot_fusion_tensorflow(self, angles_1, angles_2):
+        """Test that the fusion of two Rot gates has the same effect as
+        applying the Rots sequentially, in Tensorflow."""
+        import tensorflow as tf
+
+        angles_1, angles_2 = tf.Variable(angles_1, dtype=tf.float64), tf.Variable(
+            angles_1, dtype=tf.float64
+        )
+        matrix_expected = self.mat_from_prod(angles_1, angles_2)
+        fused_angles = fuse_rot_angles(angles_1, angles_2)
+        matrix_obtained = qml.Rot(*fused_angles, wires=0).matrix()
+
+        assert qml.math.allclose(matrix_expected, matrix_obtained)
+
+    @pytest.mark.torch
+    @pytest.mark.parametrize("angles_1, angles_2", generic_test_angles)
+    def test_full_rot_fusion_torch(self, angles_1, angles_2):
+        """Test that the fusion of two Rot gates has the same effect as
+        applying the Rots sequentially, in torch."""
+        import torch
+
+        angles_1, angles_2 = torch.tensor(
+            angles_1, requires_grad=True, dtype=torch.float64
+        ), torch.tensor(angles_1, requires_grad=True, dtype=torch.float64)
+        matrix_expected = self.mat_from_prod(angles_1, angles_2)
+        fused_angles = fuse_rot_angles(angles_1, angles_2)
+        matrix_obtained = qml.Rot(*fused_angles, wires=0).matrix()
+
+        assert qml.math.allclose(matrix_expected, matrix_obtained)
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize("angles_1, angles_2", generic_test_angles)
+    def test_full_rot_fusion_jax(self, angles_1, angles_2):
+        """Test that the fusion of two Rot gates has the same effect as
+        applying the Rots sequentially, in JAX."""
+        import jax
+
+        angles_1, angles_2 = jax.numpy.array(angles_1), jax.numpy.array(angles_1)
+        matrix_expected = self.mat_from_prod(angles_1, angles_2)
         fused_angles = fuse_rot_angles(angles_1, angles_2)
         matrix_obtained = qml.Rot(*fused_angles, wires=0).matrix()
 
@@ -103,16 +164,12 @@ def test_full_rot_fusion_special_angles(self):
         special_angles = np.array(list(product(special_points, repeat=6))).reshape((-1, 2, 3))
         angles_1, angles_2 = np.transpose(special_angles, (1, 0, 2))
 
-        def original_ops():
-            qml.Rot(*angles_1.T, wires=0)  # Transpose to bring size-3 axis to front
-            qml.Rot(*angles_2.T, wires=0)  # Transpose to bring size-3 axis to front
-
-        matrix_expected = qml.matrix(original_ops, [0])()  # pylint:disable=too-many-function-args
+        # Transpose to bring size-3 axis to front
+        matrix_expected = self.mat_from_prod(angles_1.T, angles_2.T)
 
         fused_angles = fuse_rot_angles(angles_1, angles_2)
-        matrix_obtained = qml.Rot(
-            *fused_angles.T, wires=0
-        ).matrix()  # Transpose to bring size-3 axis to front
+        # Transpose to bring size-3 axis to front
+        matrix_obtained = qml.Rot(*fused_angles.T, wires=0).matrix()
 
         assert qml.math.allclose(matrix_expected, matrix_obtained)
 
@@ -163,61 +220,3 @@ def mat_from_fuse(angles):
         assert qml.math.all(
             qml.math.any(qml.math.isnan(jac_from_fuse[error_sources]), axis=[1, 2, 3, 4])
         )
-
-        """
-        mag = np.sqrt(pre_mag)
-
-        alpha1, beta1 = (phi1 + omega1) / 2, (phi1 - omega1) / 2
-        alpha2, beta2 = (phi2 + omega2) / 2, (phi2 - omega2) / 2
-
-        alpha_arg1 = -c1 * c2 * np.sin(alpha1 + alpha2) - s1 * s2 * np.sin(beta2 - beta1)
-        alpha_arg2 = c1 * c2 * np.cos(alpha1 + alpha2) - s1 * s2 * np.cos(beta2 - beta1)
-        beta_arg1 = -c1 * s2 * np.sin(alpha1 + beta2) + s1 * c2 * np.sin(alpha2 - beta1)
-        beta_arg2 = c1 * s2 * np.cos(alpha1 + beta2) + s1 * c2 * np.cos(alpha2 - beta1)
-
-        def partial(all_angles):
-            phi1, theta1, omega1 = all_angles[..., 0, :].T
-            phi2, theta2, omega2 = all_angles[..., 1, :].T
-            c1, c2, s1, s2 = qml.math.cos(theta1 / 2), qml.math.cos(theta2 / 2), qml.math.sin(theta1 / 2), qml.math.sin(theta2 / 2)
-            mag = qml.math.sqrt(c1 ** 2 * c2 ** 2 + s1 ** 2 * s2 ** 2 - 2 * c1 * c2 * s1 * s2 * qml.math.cos(omega1 + phi2))
-
-            #alpha1, beta1 = (phi1 + omega1) / 2, (phi1 - omega1) / 2
-            #alpha2, beta2 = (phi2 + omega2) / 2, (phi2 - omega2) / 2
-
-            #alpha_arg1 = -c1 * c2 * np.sin(alpha1 + alpha2) - s1 * s2 * np.sin(beta2 - beta1)
-            #alpha_arg2 = c1 * c2 * np.cos(alpha1 + alpha2) - s1 * s2 * np.cos(beta2 - beta1)
-            #beta_arg1 = -c1 * s2 * np.sin(alpha1 + beta2) + s1 * c2 * np.sin(alpha2 - beta1)
-            #beta_arg2 = c1 * s2 * np.cos(alpha1 + beta2) + s1 * c2 * np.cos(alpha2 - beta1)
-            #print(mag)
-            return mag
-
-        jac_partial = jax.vmap(jax.jacobian(partial))(all_angles.real)
-
-        expected_singular_ids = np.where(
-            np.any([
-                (np.abs(pre_mag - 1) < 1e-12),
-                (pre_mag == 0j),
-                #np.isnan(mag),
-                #(np.abs(alpha_arg1) + np.abs(alpha_arg1) < 1e-10),
-                #(np.abs(beta_arg1) + np.abs(beta_arg1) < 1e-10)
-            ], axis=0)
-        )
-        pre_mag2 = c1 ** 2 * s2 ** 2 + s1 ** 2 * c2 ** 2 + 2 * c1 * c2 * s1 * s2 * qml.math.cos(omega1 + phi2)
-        print(pre_mag2[np.where(pre_mag==0j)[0]])
-        for idx in np.where(qml.math.any(qml.math.isnan(jac_from_fuse), axis=[1, 2, 3, 4]))[0]:
-            if idx not in expected_singular_ids[0]:
-                weird_idx = idx
-                break
-        print(weird_idx)
-        print(jac_from_fuse[weird_idx])
-        print(all_angles[weird_idx])
-        print(f"{pre_mag[weird_idx]:.20f}")
-        print(jax.jacobian(qml.math.sqrt,holomorphic=True)(pre_mag[weird_idx]))
-        print(jax.jacobian(qml.math.sqrt)(pre_mag.real[weird_idx]))
-        print(len(all_angles))
-        print(len(expected_singular_ids[0]))
-        print(len(np.where(qml.math.any(qml.math.isnan(jac_from_fuse), axis=[1, 2, 3, 4]))[0]))
-        print(jac_partial.shape)
-        print(len(np.where(qml.math.any(qml.math.isnan(jac_partial), axis=[1, 2]))[0]))
-        print(np.unique(mag[np.array(np.where(qml.math.any(qml.math.isnan(jac_partial), axis=[1, 2]))[0])]))
-        """

From c1e918e00d655068ac100d004342f34f6198d66d Mon Sep 17 00:00:00 2001
From: dwierichs <david.wierichs@xanadu.ai>
Date: Tue, 23 Jul 2024 23:04:39 +0200
Subject: [PATCH 06/30] docstring dtype

---
 pennylane/transforms/optimization/optimization_utils.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/pennylane/transforms/optimization/optimization_utils.py b/pennylane/transforms/optimization/optimization_utils.py
index 7c743e4f6bc..966abe1b0a7 100644
--- a/pennylane/transforms/optimization/optimization_utils.py
+++ b/pennylane/transforms/optimization/optimization_utils.py
@@ -50,11 +50,11 @@ def fuse_rot_angles(angles1, angles2):
     is non-trivial.
 
     Args:
-        angles_1 (float): A set of three angles for the first ``qml.Rot`` operation.
-        angles_2 (float): A set of three angles for the second ``qml.Rot`` operation.
+        angles_1 (tensor_like): A set of three angles for the first ``qml.Rot`` operation.
+        angles_2 (tensor_like): A set of three angles for the second ``qml.Rot`` operation.
 
     Returns:
-        array[float]: Rotation angles for a single ``qml.Rot`` operation that
+        tensor_like: Rotation angles for a single ``qml.Rot`` operation that
         implements the same operation as the two sets of input angles.
 
     .. note::

From 0e4085d11f03a0dc1a514ee2383c4ad8537d1a46 Mon Sep 17 00:00:00 2001
From: dwierichs <david.wierichs@xanadu.ai>
Date: Tue, 23 Jul 2024 23:05:05 +0200
Subject: [PATCH 07/30] prints

---
 pennylane/transforms/optimization/optimization_utils.py | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/pennylane/transforms/optimization/optimization_utils.py b/pennylane/transforms/optimization/optimization_utils.py
index 966abe1b0a7..7d98146aa48 100644
--- a/pennylane/transforms/optimization/optimization_utils.py
+++ b/pennylane/transforms/optimization/optimization_utils.py
@@ -79,7 +79,6 @@ def fuse_rot_angles(angles1, angles2):
     c1, c2, s1, s2 = cos(theta1 / 2), cos(theta2 / 2), sin(theta1 / 2), sin(theta2 / 2)
 
     mag = sqrt(c1**2 * c2**2 + s1**2 * s2**2 - 2 * c1 * c2 * s1 * s2 * cos(omega1 + phi2))
-    print(f"{mag=}")
     theta_f = 2 * arccos(mag)
 
     alpha1, beta1 = (phi1 + omega1) / 2, (phi1 - omega1) / 2
@@ -87,12 +86,10 @@ def fuse_rot_angles(angles1, angles2):
 
     alpha_arg1 = -c1 * c2 * sin(alpha1 + alpha2) - s1 * s2 * sin(beta2 - beta1)
     alpha_arg2 = c1 * c2 * cos(alpha1 + alpha2) - s1 * s2 * cos(beta2 - beta1)
-    print(f"{alpha_arg1=}, {alpha_arg2=}")
     alpha_f = -1 * arctan2(alpha_arg1, alpha_arg2)
 
     beta_arg1 = -c1 * s2 * sin(alpha1 + beta2) + s1 * c2 * sin(alpha2 - beta1)
     beta_arg2 = c1 * s2 * cos(alpha1 + beta2) + s1 * c2 * cos(alpha2 - beta1)
-    print(f"{beta_arg1=}, {beta_arg2=}")
     beta_f = -1 * arctan2(beta_arg1, beta_arg2)
 
     return stack([alpha_f + beta_f, theta_f, alpha_f - beta_f], axis=-1)

From 4dc44a5553d26c35a06b0637d03fc02bd572a7e4 Mon Sep 17 00:00:00 2001
From: dwierichs <david.wierichs@xanadu.ai>
Date: Wed, 24 Jul 2024 10:21:09 +0200
Subject: [PATCH 08/30] math derivation in docstring

---
 .../optimization/optimization_utils.py        | 111 +++++++++++++++++-
 1 file changed, 110 insertions(+), 1 deletion(-)

diff --git a/pennylane/transforms/optimization/optimization_utils.py b/pennylane/transforms/optimization/optimization_utils.py
index 7d98146aa48..bd8774e5f56 100644
--- a/pennylane/transforms/optimization/optimization_utils.py
+++ b/pennylane/transforms/optimization/optimization_utils.py
@@ -42,7 +42,7 @@ def find_next_gate(wires, op_list):
 
 
 def fuse_rot_angles(angles1, angles2):
-    """Compute the set of rotation angles that is equivalent to performing
+    r"""Compute the set of rotation angles that is equivalent to performing
     two successive ``qml.Rot`` operations.
 
     The ``qml.Rot`` operation represents the most general single-qubit operation.
@@ -73,6 +73,115 @@ def fuse_rot_angles(angles1, angles2):
         This function is numerically unstable at singular points. It is recommended to use
         it with 64-bit floating point precision.
 
+    .. details::
+        :title: Derivation
+        :href: derivation
+
+        The matrices for the two individual rotations are given by
+
+        .. math::
+
+            R(\phi_j,\theta_j,\omega_j)
+            &= \begin{bmatrix}
+            e^{-i(\phi_j+\omega_j)/2}\cos(\theta_j/2) & -e^{i(\phi_j-\omega_j)/2}\sin(\theta_j/2)\\
+            e^{-i(\phi_j-\omega_j)/2}\sin(\theta_j/2) & e^{i(\phi_j+\omega_j)/2}\cos(\theta_j/2)
+            \end{bmatrix}\\
+            &= \begin{bmatrix}
+            e^{-i\alpha_j}c_j & -e^{i\beta_j}s_j \\
+            e^{-i\beta_j}s_j & e^{i\alpha_j}c_j
+            \end{bmatrix},
+
+        where we introduced abbreviations :math:`\alpha_j`, :math:`\beta_j`,
+        :math:`c_j=\cos(\theta_j / 2)` and :math:`s_j=\sin(\theta_j / 2)` for notational brevity.
+        The upper left entry of the matrix product
+        :math:`R(\phi_2,\theta_2,\omega_2)R(\phi_1,\theta_1,\omega_1)` reads
+
+        .. math::
+
+            x = e^{-i(\alpha_2+\alpha_1)} c_2 c_1 - e^{i(\beta_2-\beta_1)} s_2 s_1
+
+        and should equal :math:`e^{-i(\alpha_f)/2}c_f` for the fused rotation angles.
+        This means that we can obtain :math:`\theta_f` from the magnitude of the matrix product
+        entry above, choosing :math:`c_f=\cos(\theta_f / 2)` to be non-negative:
+
+        .. math::
+
+            c_f = |x| &=
+            \left|
+            e^{-i(\alpha_2+\alpha_1)} c_2 c_1
+            -e^{i(\beta_2-\beta_1)} s_2 s_1
+            \right| \\
+            &= \sqrt{c_1^2 c_2^2 + s_1^2 s_2^2 - 2 c_1 c_2 s_1 s_2 \cos(\omega_1 + \phi_2)}.
+
+        Now we again make a choice and pick :math:`\theta_f` to be non-negative:
+
+        .. math::
+
+            \theta_f = 2\arccos(|x|).
+
+        We can extract the angle combination :math:`\alpha_f` from :math:`x` above as well via
+        :math:`\operatorname{arg}(x)`, which can be readily computed with :math:`arctan`:
+
+        .. math::
+
+            \alpha_f = -\arctan\left(
+            \frac{-c_1c_2\sin(\alpha_1+\alpha_2)-s_1s_2\sin(\beta_2-\beta_1)}
+            {c_1c_2\cos(\alpha_1+\alpha_2)-s_1s_2\cos(\beta_2-\beta_1)}
+            \right).
+
+        We can use the standard numerical function :math:`\operatorname{arctan2}`, which
+        computes :math:`\arctan(x_1/x_2)` from :math:`x_1` and :math:`x_2` while handling
+        special points suitably to obtain the argument of the underlying complex number 
+        :math:`x_2 + x_1 i`.
+
+        Finally, to obtain :math:`\beta_f`, we need a second element of the matrix product from
+        above. We compute the lower-left entry to be
+
+        .. math::
+
+            y = e^{-i(\beta_2+\alpha_1)} s_2 c_1 + e^{i(\alpha_2-\beta_1)} c_2 s_1,
+
+        which should equal :math:`e^{-i \beta_f}s_f`. From this, we can compute
+
+        .. math::
+
+            \beta_f = -\arctan\left(
+            \frac{-c_1s_2\sin(\alpha_1+\beta_2)+s_1c_2\sin(\alpha_2-\beta_1)}
+            {c_1s_2\cos(\alpha_1+\beta_2)+s_1c_2\cos(\alpha_2-\beta_1)}
+            \right).
+
+        From this, we may extract
+
+        .. math::
+
+            \phi_f = \alpha_f + \beta_f
+            \omega_f = \alpha_f - \beta_f
+
+        and are done.
+
+        **Mathematical properties:**
+
+        All functions above are well-defined on the domain we are using them on,
+        if we handle :math:`\arctan` via standard numerical implementations such as
+        ``np.arctan2``.
+        Based on the choices we made in the derivation above, the fused angles will lie in 
+        the intervals
+
+        .. math::
+
+            \phi_f, \omega_f \in [-\pi, \pi],\quad \theta_f \in [0, \pi].
+
+        Close to the boundaries of these intervals, ``fuse_rot_angles`` exhibits
+        discontinuities, depending on the combination of input angles.
+        These discontinuities also lead to singular (non-differentiable) points as discussed below.
+
+        **Differentiability:**
+
+        The function derived above is differentiable almost everywhere.
+        In particular, there are two problematic scenarios at which the derivative is not defined:
+        First, the square root is not differentiable at :math:`0`, making all input angles with
+        :math:`|x|=0` singular. Second, :math:`\arccos` is not differentiable at :math:`1`, making
+        all input angles with :math:`|x|=1` singular.
     """
     phi1, theta1, omega1 = transpose(asarray(angles1))
     phi2, theta2, omega2 = transpose(asarray(angles2))

From b436aba0c8d1f045083ff388f7d6a5981da26fa1 Mon Sep 17 00:00:00 2001
From: dwierichs <david.wierichs@xanadu.ai>
Date: Wed, 24 Jul 2024 11:03:36 +0200
Subject: [PATCH 09/30] mixed batching test

---
 .../optimization/optimization_utils.py        | 15 +++++++----
 .../test_optimization_utils.py                | 25 +++++++++++++++++++
 2 files changed, 35 insertions(+), 5 deletions(-)

diff --git a/pennylane/transforms/optimization/optimization_utils.py b/pennylane/transforms/optimization/optimization_utils.py
index bd8774e5f56..dc61ced5e4d 100644
--- a/pennylane/transforms/optimization/optimization_utils.py
+++ b/pennylane/transforms/optimization/optimization_utils.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 """Utility functions for circuit optimization."""
 # pylint: disable=too-many-return-statements,import-outside-toplevel
-from pennylane.math import arccos, arctan2, asarray, cos, sin, sqrt, stack, transpose
+from pennylane.math import arccos, arctan2, asarray, cos, moveaxis, sin, sqrt, stack
 from pennylane.ops.identity import GlobalPhase
 from pennylane.wires import Wires
 
@@ -57,6 +57,9 @@ def fuse_rot_angles(angles1, angles2):
         tensor_like: Rotation angles for a single ``qml.Rot`` operation that
         implements the same operation as the two sets of input angles.
 
+    This function supports broadcasting/batching as long as the two inputs are standard
+    broadcast-compatible.
+
     .. note::
 
         The output angles are not always defined uniquely because Euler angles are not
@@ -131,7 +134,7 @@ def fuse_rot_angles(angles1, angles2):
 
         We can use the standard numerical function :math:`\operatorname{arctan2}`, which
         computes :math:`\arctan(x_1/x_2)` from :math:`x_1` and :math:`x_2` while handling
-        special points suitably to obtain the argument of the underlying complex number 
+        special points suitably to obtain the argument of the underlying complex number
         :math:`x_2 + x_1 i`.
 
         Finally, to obtain :math:`\beta_f`, we need a second element of the matrix product from
@@ -164,7 +167,7 @@ def fuse_rot_angles(angles1, angles2):
         All functions above are well-defined on the domain we are using them on,
         if we handle :math:`\arctan` via standard numerical implementations such as
         ``np.arctan2``.
-        Based on the choices we made in the derivation above, the fused angles will lie in 
+        Based on the choices we made in the derivation above, the fused angles will lie in
         the intervals
 
         .. math::
@@ -183,8 +186,10 @@ def fuse_rot_angles(angles1, angles2):
         :math:`|x|=0` singular. Second, :math:`\arccos` is not differentiable at :math:`1`, making
         all input angles with :math:`|x|=1` singular.
     """
-    phi1, theta1, omega1 = transpose(asarray(angles1))
-    phi2, theta2, omega2 = transpose(asarray(angles2))
+
+    # moveaxis required for batched inputs
+    phi1, theta1, omega1 = moveaxis(asarray(angles1), -1, 0)
+    phi2, theta2, omega2 = moveaxis(asarray(angles2), -1, 0)
     c1, c2, s1, s2 = cos(theta1 / 2), cos(theta2 / 2), sin(theta1 / 2), sin(theta2 / 2)
 
     mag = sqrt(c1**2 * c2**2 + s1**2 * s2**2 - 2 * c1 * c2 * s1 * s2 * cos(omega1 + phi2))
diff --git a/tests/transforms/test_optimization/test_optimization_utils.py b/tests/transforms/test_optimization/test_optimization_utils.py
index e3a75b7d5d1..cdbd8a2be42 100644
--- a/tests/transforms/test_optimization/test_optimization_utils.py
+++ b/tests/transforms/test_optimization/test_optimization_utils.py
@@ -94,6 +94,31 @@ def test_full_rot_fusion_numpy(self, angles_1, angles_2):
 
         assert qml.math.allclose(matrix_expected, matrix_obtained)
 
+    mixed_batched_angles = [
+        ([[0.4, 0.1, 0.0], [0.7, 0.2, 0.1]], [-0.9, 1.2, 0.6]),  # (2, None)
+        ([-0.9, 1.2, 0.6], [[0.4, 0.1, 0.0], [0.7, 0.2, 0.1]]),  # (None, 2)
+        ([-0.9, 1.2, 0.6], [[[0.4, 0.1, 0.0], [0.7, 0.2, 0.1]]] * 4),  # (None, (4, 2))
+        (
+            [[[-0.9, 1.2, 0.6]] * 2] * 4,
+            [[[0.4, 0.1, 0.0], [0.7, 0.2, 0.1]]] * 4,
+        ),  # ((4, 2), (4, 2))
+    ]
+
+    @pytest.mark.parametrize("angles_1, angles_2", mixed_batched_angles)
+    def test_full_rot_fusion_mixed_batching(self, angles_1, angles_2):
+        """Test that the fusion of two Rot gates has the same effect as
+        applying the Rots sequentially when the input angles are batched
+        with mixed batching shapes."""
+
+        fused_angles = fuse_rot_angles(angles_1, angles_2)
+
+        reshaped_angles_1 = np.reshape(angles_1, (-1, 3) if np.ndim(angles_1) > 1 else (3,))
+        reshaped_angles_2 = np.reshape(angles_2, (-1, 3) if np.ndim(angles_2) > 1 else (3,))
+        matrix_expected = self.mat_from_prod(reshaped_angles_1.T, reshaped_angles_2.T)
+        matrix_obtained = qml.Rot(*fused_angles.reshape((-1, 3)).T, wires=0).matrix()
+
+        assert qml.math.allclose(matrix_expected, matrix_obtained)
+
     @pytest.mark.autograd
     @pytest.mark.parametrize("angles_1, angles_2", generic_test_angles)
     def test_full_rot_fusion_autograd(self, angles_1, angles_2):

From 3148643eb0b6bf6c8f859e37ff327be974206458 Mon Sep 17 00:00:00 2001
From: dwierichs <david.wierichs@xanadu.ai>
Date: Wed, 24 Jul 2024 11:16:59 +0200
Subject: [PATCH 10/30] documentation

---
 .../optimization/merge_rotations.py           |  23 +++-
 .../optimization/optimization_utils.py        |   2 +-
 .../optimization/single_qubit_fusion.py       | 126 ++++++++++++++++++
 3 files changed, 148 insertions(+), 3 deletions(-)

diff --git a/pennylane/transforms/optimization/merge_rotations.py b/pennylane/transforms/optimization/merge_rotations.py
index 38cfaefd7ba..ace24b71138 100644
--- a/pennylane/transforms/optimization/merge_rotations.py
+++ b/pennylane/transforms/optimization/merge_rotations.py
@@ -70,10 +70,29 @@ def circuit(x, y, z):
     >>> circuit(0.1, 0.2, 0.3)
     0.9553364891256055
 
+    .. details::
+        :title: Details on merging ``Rot`` gates
+        :href: details-on-rot
+
+        When merging two :class:`~.pennylane.Rot` gates, there are a number of details to consider:
+
+        First, the output angles are not always uniquely defined, because Euler angles are not
+        defined uniquely for some rotations. ``merge_rotations`` makes a particular choice in
+        this case.
+
+        Second, ``merge_rotations`` is not differentiable everywhere when used on ``Rot``.
+        It has singularities for specific rotation angles where the derivative will be NaN.
+
+        Finally, this function can be numerically unstable at singular points.
+        It is therefore recommended to use it with 64-bit floating point precision angles.
+
+        For a mathematical derivation of the fusion of two ``Rot`` gates, see the documentation
+        of :func:`~.pennylane.transforms.single_qubit_fusion`.
+
     .. details::
         :title: Usage Details
 
-        You can also apply it on quantum function.
+        You can also apply ``merge_rotations`` to a quantum function.
 
         .. code-block:: python
 
@@ -106,7 +125,7 @@ def qfunc(x, y, z):
         2: ─╰X─────────H─╰●────────┤
 
         It is also possible to explicitly specify which rotations ``merge_rotations`` should
-        be merged using the ``include_gates`` argument. For example, if in the above
+        merge using the ``include_gates`` argument. For example, if in the above
         circuit we wanted only to merge the "RX" gates, we could do so as follows:
 
         >>> optimized_qfunc = merge_rotations(include_gates=["RX"])(qfunc)
diff --git a/pennylane/transforms/optimization/optimization_utils.py b/pennylane/transforms/optimization/optimization_utils.py
index dc61ced5e4d..20bc2d061a6 100644
--- a/pennylane/transforms/optimization/optimization_utils.py
+++ b/pennylane/transforms/optimization/optimization_utils.py
@@ -69,7 +69,7 @@ def fuse_rot_angles(angles1, angles2):
     .. warning::
 
         This function is not differentiable everywhere. It has singularities for specific
-        input values, where the derivative will be ``nan``.
+        input values where the derivative will be NaN.
 
     .. warning::
 
diff --git a/pennylane/transforms/optimization/single_qubit_fusion.py b/pennylane/transforms/optimization/single_qubit_fusion.py
index 9bb56dbd500..fa79b330c32 100644
--- a/pennylane/transforms/optimization/single_qubit_fusion.py
+++ b/pennylane/transforms/optimization/single_qubit_fusion.py
@@ -67,6 +67,22 @@ def qfunc(r1, r2):
 
     The single qubit gates are fused before execution.
 
+    .. note::
+
+        The fused angles between two sets of rotation angles are not always defined uniquely
+        because Euler angles are not defined uniquely for some rotations. ``single_qubit_fusion``
+        makes a particular choice in this case.
+
+    .. warning::
+
+        This function is not differentiable everywhere. It has singularities for specific
+        input rotation angles, where the derivative will be NaN.
+
+    .. warning::
+
+        This function is numerically unstable at its singular points. It is recommended to use
+        it with 64-bit floating point precision.
+
     .. details::
         :title: Usage Details
 
@@ -96,6 +112,116 @@ def qfunc(r1, r2):
         >>> print(qml.draw(optimized_qnode)([0.1, 0.2, 0.3], [0.4, 0.5, 0.6]))
         0: ──Rot(3.57, 2.09, 2.05)──┤ ⟨X⟩
 
+    .. details::
+        :title: Derivation
+        :href: derivation
+
+        The matrices for the two individual rotations are given by
+
+        .. math::
+
+            R(\phi_j,\theta_j,\omega_j)
+            &= \begin{bmatrix}
+            e^{-i(\phi_j+\omega_j)/2}\cos(\theta_j/2) & -e^{i(\phi_j-\omega_j)/2}\sin(\theta_j/2)\\
+            e^{-i(\phi_j-\omega_j)/2}\sin(\theta_j/2) & e^{i(\phi_j+\omega_j)/2}\cos(\theta_j/2)
+            \end{bmatrix}\\
+            &= \begin{bmatrix}
+            e^{-i\alpha_j}c_j & -e^{i\beta_j}s_j \\
+            e^{-i\beta_j}s_j & e^{i\alpha_j}c_j
+            \end{bmatrix},
+
+        where we introduced abbreviations :math:`\alpha_j`, :math:`\beta_j`,
+        :math:`c_j=\cos(\theta_j / 2)` and :math:`s_j=\sin(\theta_j / 2)` for notational brevity.
+        The upper left entry of the matrix product
+        :math:`R(\phi_2,\theta_2,\omega_2)R(\phi_1,\theta_1,\omega_1)` reads
+
+        .. math::
+
+            x = e^{-i(\alpha_2+\alpha_1)} c_2 c_1 - e^{i(\beta_2-\beta_1)} s_2 s_1
+
+        and should equal :math:`e^{-i(\alpha_f)/2}c_f` for the fused rotation angles.
+        This means that we can obtain :math:`\theta_f` from the magnitude of the matrix product
+        entry above, choosing :math:`c_f=\cos(\theta_f / 2)` to be non-negative:
+
+        .. math::
+
+            c_f = |x| &=
+            \left|
+            e^{-i(\alpha_2+\alpha_1)} c_2 c_1
+            -e^{i(\beta_2-\beta_1)} s_2 s_1
+            \right| \\
+            &= \sqrt{c_1^2 c_2^2 + s_1^2 s_2^2 - 2 c_1 c_2 s_1 s_2 \cos(\omega_1 + \phi_2)}.
+
+        Now we again make a choice and pick :math:`\theta_f` to be non-negative:
+
+        .. math::
+
+            \theta_f = 2\arccos(|x|).
+
+        We can extract the angle combination :math:`\alpha_f` from :math:`x` above as well via
+        :math:`\operatorname{arg}(x)`, which can be readily computed with :math:`arctan`:
+
+        .. math::
+
+            \alpha_f = -\arctan\left(
+            \frac{-c_1c_2\sin(\alpha_1+\alpha_2)-s_1s_2\sin(\beta_2-\beta_1)}
+            {c_1c_2\cos(\alpha_1+\alpha_2)-s_1s_2\cos(\beta_2-\beta_1)}
+            \right).
+
+        We can use the standard numerical function :math:`\operatorname{arctan2}`, which
+        computes :math:`\arctan(x_1/x_2)` from :math:`x_1` and :math:`x_2` while handling
+        special points suitably to obtain the argument of the underlying complex number
+        :math:`x_2 + x_1 i`.
+
+        Finally, to obtain :math:`\beta_f`, we need a second element of the matrix product from
+        above. We compute the lower-left entry to be
+
+        .. math::
+
+            y = e^{-i(\beta_2+\alpha_1)} s_2 c_1 + e^{i(\alpha_2-\beta_1)} c_2 s_1,
+
+        which should equal :math:`e^{-i \beta_f}s_f`. From this, we can compute
+
+        .. math::
+
+            \beta_f = -\arctan\left(
+            \frac{-c_1s_2\sin(\alpha_1+\beta_2)+s_1c_2\sin(\alpha_2-\beta_1)}
+            {c_1s_2\cos(\alpha_1+\beta_2)+s_1c_2\cos(\alpha_2-\beta_1)}
+            \right).
+
+        From this, we may extract
+
+        .. math::
+
+            \phi_f = \alpha_f + \beta_f
+            \omega_f = \alpha_f - \beta_f
+
+        and are done.
+
+        **Mathematical properties:**
+
+        All functions above are well-defined on the domain we are using them on,
+        if we handle :math:`\arctan` via standard numerical implementations such as
+        ``np.arctan2``.
+        Based on the choices we made in the derivation above, the fused angles will lie in
+        the intervals
+
+        .. math::
+
+            \phi_f, \omega_f \in [-\pi, \pi],\quad \theta_f \in [0, \pi].
+
+        Close to the boundaries of these intervals, ``single_qubit_fusion`` exhibits
+        discontinuities, depending on the combination of input angles.
+        These discontinuities also lead to singular (non-differentiable) points as discussed below.
+
+        **Differentiability:**
+
+        The function derived above is differentiable almost everywhere.
+        In particular, there are two problematic scenarios at which the derivative is not defined:
+        First, the square root is not differentiable at :math:`0`, making all input angles with
+        :math:`|x|=0` singular. Second, :math:`\arccos` is not differentiable at :math:`1`, making
+        all input angles with :math:`|x|=1` singular.
+
     """
     # Make a working copy of the list to traverse
     list_copy = tape.operations.copy()

From 976a59e24584c15e20b0fe5a68bde1954d0828d4 Mon Sep 17 00:00:00 2001
From: dwierichs <david.wierichs@xanadu.ai>
Date: Wed, 24 Jul 2024 11:17:50 +0200
Subject: [PATCH 11/30] documentation

---
 .../optimization/optimization_utils.py        | 111 +-----------------
 1 file changed, 2 insertions(+), 109 deletions(-)

diff --git a/pennylane/transforms/optimization/optimization_utils.py b/pennylane/transforms/optimization/optimization_utils.py
index 20bc2d061a6..1b2c61298c7 100644
--- a/pennylane/transforms/optimization/optimization_utils.py
+++ b/pennylane/transforms/optimization/optimization_utils.py
@@ -76,115 +76,8 @@ def fuse_rot_angles(angles1, angles2):
         This function is numerically unstable at singular points. It is recommended to use
         it with 64-bit floating point precision.
 
-    .. details::
-        :title: Derivation
-        :href: derivation
-
-        The matrices for the two individual rotations are given by
-
-        .. math::
-
-            R(\phi_j,\theta_j,\omega_j)
-            &= \begin{bmatrix}
-            e^{-i(\phi_j+\omega_j)/2}\cos(\theta_j/2) & -e^{i(\phi_j-\omega_j)/2}\sin(\theta_j/2)\\
-            e^{-i(\phi_j-\omega_j)/2}\sin(\theta_j/2) & e^{i(\phi_j+\omega_j)/2}\cos(\theta_j/2)
-            \end{bmatrix}\\
-            &= \begin{bmatrix}
-            e^{-i\alpha_j}c_j & -e^{i\beta_j}s_j \\
-            e^{-i\beta_j}s_j & e^{i\alpha_j}c_j
-            \end{bmatrix},
-
-        where we introduced abbreviations :math:`\alpha_j`, :math:`\beta_j`,
-        :math:`c_j=\cos(\theta_j / 2)` and :math:`s_j=\sin(\theta_j / 2)` for notational brevity.
-        The upper left entry of the matrix product
-        :math:`R(\phi_2,\theta_2,\omega_2)R(\phi_1,\theta_1,\omega_1)` reads
-
-        .. math::
-
-            x = e^{-i(\alpha_2+\alpha_1)} c_2 c_1 - e^{i(\beta_2-\beta_1)} s_2 s_1
-
-        and should equal :math:`e^{-i(\alpha_f)/2}c_f` for the fused rotation angles.
-        This means that we can obtain :math:`\theta_f` from the magnitude of the matrix product
-        entry above, choosing :math:`c_f=\cos(\theta_f / 2)` to be non-negative:
-
-        .. math::
-
-            c_f = |x| &=
-            \left|
-            e^{-i(\alpha_2+\alpha_1)} c_2 c_1
-            -e^{i(\beta_2-\beta_1)} s_2 s_1
-            \right| \\
-            &= \sqrt{c_1^2 c_2^2 + s_1^2 s_2^2 - 2 c_1 c_2 s_1 s_2 \cos(\omega_1 + \phi_2)}.
-
-        Now we again make a choice and pick :math:`\theta_f` to be non-negative:
-
-        .. math::
-
-            \theta_f = 2\arccos(|x|).
-
-        We can extract the angle combination :math:`\alpha_f` from :math:`x` above as well via
-        :math:`\operatorname{arg}(x)`, which can be readily computed with :math:`arctan`:
-
-        .. math::
-
-            \alpha_f = -\arctan\left(
-            \frac{-c_1c_2\sin(\alpha_1+\alpha_2)-s_1s_2\sin(\beta_2-\beta_1)}
-            {c_1c_2\cos(\alpha_1+\alpha_2)-s_1s_2\cos(\beta_2-\beta_1)}
-            \right).
-
-        We can use the standard numerical function :math:`\operatorname{arctan2}`, which
-        computes :math:`\arctan(x_1/x_2)` from :math:`x_1` and :math:`x_2` while handling
-        special points suitably to obtain the argument of the underlying complex number
-        :math:`x_2 + x_1 i`.
-
-        Finally, to obtain :math:`\beta_f`, we need a second element of the matrix product from
-        above. We compute the lower-left entry to be
-
-        .. math::
-
-            y = e^{-i(\beta_2+\alpha_1)} s_2 c_1 + e^{i(\alpha_2-\beta_1)} c_2 s_1,
-
-        which should equal :math:`e^{-i \beta_f}s_f`. From this, we can compute
-
-        .. math::
-
-            \beta_f = -\arctan\left(
-            \frac{-c_1s_2\sin(\alpha_1+\beta_2)+s_1c_2\sin(\alpha_2-\beta_1)}
-            {c_1s_2\cos(\alpha_1+\beta_2)+s_1c_2\cos(\alpha_2-\beta_1)}
-            \right).
-
-        From this, we may extract
-
-        .. math::
-
-            \phi_f = \alpha_f + \beta_f
-            \omega_f = \alpha_f - \beta_f
-
-        and are done.
-
-        **Mathematical properties:**
-
-        All functions above are well-defined on the domain we are using them on,
-        if we handle :math:`\arctan` via standard numerical implementations such as
-        ``np.arctan2``.
-        Based on the choices we made in the derivation above, the fused angles will lie in
-        the intervals
-
-        .. math::
-
-            \phi_f, \omega_f \in [-\pi, \pi],\quad \theta_f \in [0, \pi].
-
-        Close to the boundaries of these intervals, ``fuse_rot_angles`` exhibits
-        discontinuities, depending on the combination of input angles.
-        These discontinuities also lead to singular (non-differentiable) points as discussed below.
-
-        **Differentiability:**
-
-        The function derived above is differentiable almost everywhere.
-        In particular, there are two problematic scenarios at which the derivative is not defined:
-        First, the square root is not differentiable at :math:`0`, making all input angles with
-        :math:`|x|=0` singular. Second, :math:`\arccos` is not differentiable at :math:`1`, making
-        all input angles with :math:`|x|=1` singular.
+    See the documentation of :func:`~.pennylane.transforms.single_qubit_fusion` for a
+    mathematical derivation of this function.
     """
 
     # moveaxis required for batched inputs

From 343d2fb3d00f354342c2a0380d4bf0ec102ccb8d Mon Sep 17 00:00:00 2001
From: dwierichs <david.wierichs@xanadu.ai>
Date: Wed, 24 Jul 2024 11:19:26 +0200
Subject: [PATCH 12/30] format

---
 pennylane/transforms/optimization/single_qubit_fusion.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/pennylane/transforms/optimization/single_qubit_fusion.py b/pennylane/transforms/optimization/single_qubit_fusion.py
index fa79b330c32..298ab29ef64 100644
--- a/pennylane/transforms/optimization/single_qubit_fusion.py
+++ b/pennylane/transforms/optimization/single_qubit_fusion.py
@@ -36,7 +36,7 @@ def single_qubit_fusion(
     (on the same qubit) with that property defined will be fused into one ``Rot``.
 
     Args:
-        tape (QNode or QuantumTape or Callable): A quantum circuit.
+        tape (QNode or QuantumTape or callable): A quantum circuit.
         atol (float): An absolute tolerance for which to apply a rotation after
             fusion. After fusion of gates, if the fused angles :math:`\theta` are such that
             :math:`|\theta|\leq \text{atol}`, no rotation gate will be applied.
@@ -45,7 +45,8 @@ def single_qubit_fusion(
             be fused will be fused.
 
     Returns:
-        qnode (QNode) or quantum function (Callable) or tuple[List[QuantumTape], function]: The transformed circuit as described in :func:`qml.transform <pennylane.transform>`.
+        qnode (QNode) or quantum function (callable) or tuple[List[QuantumTape], callable]:
+        The transformed circuit as described in :func:`qml.transform <pennylane.transform>`.
 
     **Example**
 

From 91b2f7874a7b13fee31119a73ddc772f4c503b9b Mon Sep 17 00:00:00 2001
From: dwierichs <david.wierichs@xanadu.ai>
Date: Wed, 24 Jul 2024 11:46:26 +0200
Subject: [PATCH 13/30] no_fuse simplifications. JIT test

---
 .../optimization/optimization_utils.py        | 47 ++++++++++++++++++-
 .../test_optimization_utils.py                | 11 +++--
 2 files changed, 53 insertions(+), 5 deletions(-)

diff --git a/pennylane/transforms/optimization/optimization_utils.py b/pennylane/transforms/optimization/optimization_utils.py
index 1b2c61298c7..6499eb274eb 100644
--- a/pennylane/transforms/optimization/optimization_utils.py
+++ b/pennylane/transforms/optimization/optimization_utils.py
@@ -13,7 +13,20 @@
 # limitations under the License.
 """Utility functions for circuit optimization."""
 # pylint: disable=too-many-return-statements,import-outside-toplevel
-from pennylane.math import arccos, arctan2, asarray, cos, moveaxis, sin, sqrt, stack
+from pennylane.math import (
+    allclose,
+    arccos,
+    arctan2,
+    asarray,
+    cos,
+    is_abstract,
+    moveaxis,
+    requires_grad,
+    sin,
+    sqrt,
+    stack,
+    zeros_like,
+)
 from pennylane.ops.identity import GlobalPhase
 from pennylane.wires import Wires
 
@@ -41,6 +54,28 @@ def find_next_gate(wires, op_list):
     return next_gate_idx
 
 
+def _try_no_fuse(angles1, angles2):
+    """Try to combine rotation angles without trigonometric identities
+    if some angles in the input angles vanish."""
+    # moveaxis required for batched inputs
+    phi1, theta1, omega1 = moveaxis(asarray(angles1), -1, 0)
+    phi2, theta2, omega2 = moveaxis(asarray(angles2), -1, 0)
+
+    if allclose(omega1 + phi2, 0.0):
+        return stack([phi1, theta1 + theta2, omega2])
+    zero = zeros_like(phi1) + zeros_like(phi2)
+    if allclose(theta1, 0.0):
+        # No Y rotation in first Rot
+        if allclose(theta2, 0.0):
+            # Z rotations only
+            return stack([phi1 + omega1 + phi2 + omega2, zero, zero])
+        return stack([phi1 + omega1 + phi2, theta2, omega2])
+    if allclose(theta2, 0.0):
+        # No Y rotation in second Rot
+        return stack([phi1, theta1, omega1 + phi2 + omega2])
+    return None
+
+
 def fuse_rot_angles(angles1, angles2):
     r"""Compute the set of rotation angles that is equivalent to performing
     two successive ``qml.Rot`` operations.
@@ -80,6 +115,16 @@ def fuse_rot_angles(angles1, angles2):
     mathematical derivation of this function.
     """
 
+    if not (
+        is_abstract(angles1)
+        or is_abstract(angles2)
+        or requires_grad(angles1)
+        or requires_grad(angles2)
+    ):
+        fused_angles = _try_no_fuse(angles1, angles2)
+        if fused_angles is not None:
+            return fused_angles
+
     # moveaxis required for batched inputs
     phi1, theta1, omega1 = moveaxis(asarray(angles1), -1, 0)
     phi2, theta2, omega2 = moveaxis(asarray(angles2), -1, 0)
diff --git a/tests/transforms/test_optimization/test_optimization_utils.py b/tests/transforms/test_optimization/test_optimization_utils.py
index cdbd8a2be42..945797999e1 100644
--- a/tests/transforms/test_optimization/test_optimization_utils.py
+++ b/tests/transforms/test_optimization/test_optimization_utils.py
@@ -200,10 +200,11 @@ def test_full_rot_fusion_special_angles(self):
 
     @pytest.mark.slow
     @pytest.mark.jax
-    def test_full_rot_fusion_jacobian(self):
+    @pytest.mark.parametrize("use_jit", [False, True])
+    def test_full_rot_fusion_jacobian(self, use_jit):
         """Test the Jacobian of the rotation angle fusion. Uses batching for performance reasons.
-        For known sources of singularities, the Jacobian is checked to indeed return NaN.
-        These sources are related to the absolute value of the upper left entry of the matrix product:
+        For known sources of singularities, the Jacobian is checked to indeed return NaN. These
+        sources are related to the absolute value of the upper left entry of the matrix product:
          - If it is 1, the derivative of arccos becomes infinite (evaluated at 1), and
          - if its square is 0, the derivative of sqrt becomes infinite (evaluated at 0).
         """
@@ -222,11 +223,13 @@ def original_ops():
 
             return qml.matrix(original_ops, [0])()  # pylint:disable=too-many-function-args
 
-        def mat_from_fuse(angles):
+        def _mat_from_fuse(angles):
             angles1, angles2 = angles[..., 0, :], angles[..., 1, :]
             fused_angles = fuse_rot_angles(angles1, angles2)
             return qml.Rot(*fused_angles.T, wires=0).matrix()
 
+        mat_from_fuse = jax.jit(_mat_from_fuse) if use_jit else _mat_from_fuse
+
         # Need holomorphic derivatives because the output matrices are complex-valued
         jac_from_prod = jax.vmap(jax.jacobian(mat_from_prod, holomorphic=True))(all_angles)
         jac_from_fuse = jax.vmap(jax.jacobian(mat_from_fuse, holomorphic=True))(all_angles)

From a3a752cbfa4b0a90ee4873c4e16b7dfcb99ebf44 Mon Sep 17 00:00:00 2001
From: dwierichs <david.wierichs@xanadu.ai>
Date: Wed, 24 Jul 2024 12:26:21 +0200
Subject: [PATCH 14/30]  fixed

---
 .../optimization/merge_rotations.py           | 48 +++++++++----------
 .../optimization/single_qubit_fusion.py       | 31 ++++++------
 .../test_optimization/test_merge_rotations.py | 10 ++--
 3 files changed, 42 insertions(+), 47 deletions(-)

diff --git a/pennylane/transforms/optimization/merge_rotations.py b/pennylane/transforms/optimization/merge_rotations.py
index ace24b71138..115a059f5a5 100644
--- a/pennylane/transforms/optimization/merge_rotations.py
+++ b/pennylane/transforms/optimization/merge_rotations.py
@@ -15,7 +15,15 @@
 # pylint: disable=too-many-branches
 
 import pennylane as qml
-from pennylane.math import allclose, cast_like, get_interface, is_abstract, stack, zeros
+from pennylane.math import (
+    allclose,
+    cast_like,
+    get_interface,
+    is_abstract,
+    requires_grad,
+    stack,
+    zeros,
+)
 from pennylane.ops.op_math import Adjoint
 from pennylane.ops.qubit.attributes import composable_rotations
 from pennylane.queuing import QueuingManager
@@ -187,27 +195,17 @@ def stop_at(obj):
             # If next gate is of the same type, we can merge the angles
             if current_gate.name == next_gate.name and current_gate.wires == next_gate.wires:
                 list_copy.pop(next_gate_idx + 1)
+                next_params = stack(next_gate.parameters, like=interface)
+                # jax-jit does not support cast_like
+                if not is_abstract(cumulative_angles):
+                    next_params = cast_like(next_params, cumulative_angles)
+
                 # The Rot gate must be treated separately
                 if current_gate.name == "Rot":
-                    if is_abstract(cumulative_angles):
-                        # jax-jit does not support cast_like
-                        cumulative_angles = cumulative_angles + stack(next_gate.parameters)
-                    else:
-                        cumulative_angles = fuse_rot_angles(
-                            cumulative_angles,
-                            cast_like(
-                                stack(next_gate.parameters, like=interface), cumulative_angles
-                            ),
-                        )
+                    cumulative_angles = fuse_rot_angles(cumulative_angles, next_params)
                 # Other, single-parameter rotation gates just have the angle summed
                 else:
-                    if is_abstract(cumulative_angles):
-                        # jax-jit does not support cast_like
-                        cumulative_angles = cumulative_angles + stack(next_gate.parameters)
-                    else:
-                        cumulative_angles = cumulative_angles + cast_like(
-                            stack(next_gate.parameters, like=interface), cumulative_angles
-                        )
+                    cumulative_angles = cumulative_angles + next_params
             # If it is not, we need to stop
             else:
                 break
@@ -215,20 +213,18 @@ def stop_at(obj):
             # 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
+        # If we are tracing/jitting or differentiating, 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 (
+            is_abstract(cumulative_angles)
+            or requires_grad(cumulative_angles)
+            or not allclose(cumulative_angles, 0.0, atol=atol, rtol=0)
+        ):
             with QueuingManager.stop_recording():
                 new_operations.append(
                     current_gate.__class__(*cumulative_angles, wires=current_gate.wires)
                 )
-        else:
-            if not allclose(cumulative_angles, zeros(len(cumulative_angles)), atol=atol, rtol=0):
-                with QueuingManager.stop_recording():
-                    new_operations.append(
-                        current_gate.__class__(*cumulative_angles, wires=current_gate.wires)
-                    )
 
         # Remove the first gate from the working list
         list_copy.pop(0)
diff --git a/pennylane/transforms/optimization/single_qubit_fusion.py b/pennylane/transforms/optimization/single_qubit_fusion.py
index 298ab29ef64..11cf265cee2 100644
--- a/pennylane/transforms/optimization/single_qubit_fusion.py
+++ b/pennylane/transforms/optimization/single_qubit_fusion.py
@@ -14,7 +14,7 @@
 """Transform for fusing sequences of single-qubit gates."""
 # pylint: disable=too-many-branches
 
-from pennylane.math import allclose, is_abstract, stack
+from pennylane.math import allclose, cast_like, is_abstract, requires_grad, stack
 from pennylane.ops.qubit import Rot
 from pennylane.queuing import QueuingManager
 from pennylane.tape import QuantumTape, QuantumTapeBatch
@@ -282,28 +282,29 @@ def qfunc(r1, r2):
                 next_gate_angles = stack(next_gate.single_qubit_rot_angles())
             except (NotImplementedError, AttributeError):
                 break
-
-            cumulative_angles = fuse_rot_angles(cumulative_angles, stack(next_gate_angles))
+            if not is_abstract(cumulative_angles):
+                next_gate_angles = cast_like(next_gate_angles, cumulative_angles)
+            cumulative_angles = fuse_rot_angles(cumulative_angles, next_gate_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
+        # If we are tracing/jitting or differentiating, don't perform any conditional checks and
         # apply the rotation regardless of the angles.
-        if is_abstract(cumulative_angles):
-            with QueuingManager.stop_recording():
-                new_operations.append(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(
+        # If not tracing or differentiating, check whether total rotation is trivial by checking
+        # if the RY angle and the sum of the RZ angles are close to 0
+        if (
+            is_abstract(cumulative_angles)
+            or requires_grad(cumulative_angles)
+            or not allclose(
                 stack([cumulative_angles[0] + cumulative_angles[2], cumulative_angles[1]]),
-                [0.0, 0.0],
+                0.0,
                 atol=atol,
                 rtol=0,
-            ):
-                with QueuingManager.stop_recording():
-                    new_operations.append(Rot(*cumulative_angles, wires=current_gate.wires))
+            )
+        ):
+            with QueuingManager.stop_recording():
+                new_operations.append(Rot(*cumulative_angles, wires=current_gate.wires))
 
         # Remove the starting gate from the list
         list_copy.pop(0)
diff --git a/tests/transforms/test_optimization/test_merge_rotations.py b/tests/transforms/test_optimization/test_merge_rotations.py
index bfcf264f154..c0aa9383459 100644
--- a/tests/transforms/test_optimization/test_merge_rotations.py
+++ b/tests/transforms/test_optimization/test_merge_rotations.py
@@ -424,13 +424,11 @@ def test_merge_rotations_jax_jit(self):
         @qml.qnode(qml.device("default.qubit", wires=["w1", "w2"]), interface="jax")
         @merge_rotations
         def qfunc():
-            qml.Rot(jax.numpy.array(0.1), jax.numpy.array(0.2), jax.numpy.array(0.3), wires=["w1"])
-            qml.Rot(
-                jax.numpy.array(-0.1), jax.numpy.array(-0.2), jax.numpy.array(-0.3), wires=["w1"]
-            )
+            qml.Rot(*jax.numpy.array([0.1, 0.2, 0.3]), wires=["w1"])
+            qml.Rot(*jax.numpy.array([-0.3, -0.2, -0.1]), wires=["w1"])
             qml.CRX(jax.numpy.array(0.2), wires=["w1", "w2"])
             qml.CRX(jax.numpy.array(-0.2), wires=["w1", "w2"])
-            return qml.expval(qml.PauliZ("w1"))
+            return qml.expval(qml.PauliZ("w2"))
 
         res = qfunc()
 
@@ -525,7 +523,7 @@ def test_qnode(self):
 
 @pytest.mark.xfail
 def test_merge_rotations_non_commuting_observables():
-    """Test that merge_roatations works with non-commuting observables."""
+    """Test that merge_rotations works with non-commuting observables."""
 
     @qml.transforms.merge_rotations
     def circuit(x):

From cf3e62137be00ead2c00996b6ae645479cb1a84e Mon Sep 17 00:00:00 2001
From: dwierichs <david.wierichs@xanadu.ai>
Date: Wed, 24 Jul 2024 13:25:52 +0200
Subject: [PATCH 15/30] fix precision problems in _try_no_fuse

---
 pennylane/transforms/optimization/optimization_utils.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/pennylane/transforms/optimization/optimization_utils.py b/pennylane/transforms/optimization/optimization_utils.py
index 6499eb274eb..9069f2b24c5 100644
--- a/pennylane/transforms/optimization/optimization_utils.py
+++ b/pennylane/transforms/optimization/optimization_utils.py
@@ -18,6 +18,7 @@
     arccos,
     arctan2,
     asarray,
+    cast_like,
     cos,
     is_abstract,
     moveaxis,
@@ -57,9 +58,10 @@ def find_next_gate(wires, op_list):
 def _try_no_fuse(angles1, angles2):
     """Try to combine rotation angles without trigonometric identities
     if some angles in the input angles vanish."""
+    dummy_sum = angles1 + angles2
     # moveaxis required for batched inputs
-    phi1, theta1, omega1 = moveaxis(asarray(angles1), -1, 0)
-    phi2, theta2, omega2 = moveaxis(asarray(angles2), -1, 0)
+    phi1, theta1, omega1 = moveaxis(cast_like(asarray(angles1), dummy_sum), -1, 0)
+    phi2, theta2, omega2 = moveaxis(cast_like(asarray(angles2), dummy_sum), -1, 0)
 
     if allclose(omega1 + phi2, 0.0):
         return stack([phi1, theta1 + theta2, omega2])

From 1967a3c9d6f327fb058dd465bd21facbd935fc97 Mon Sep 17 00:00:00 2001
From: dwierichs <david.wierichs@xanadu.ai>
Date: Wed, 24 Jul 2024 13:29:28 +0200
Subject: [PATCH 16/30] clean out imports

---
 .../optimization/optimization_utils.py        | 78 +++++++++----------
 1 file changed, 35 insertions(+), 43 deletions(-)

diff --git a/pennylane/transforms/optimization/optimization_utils.py b/pennylane/transforms/optimization/optimization_utils.py
index 9069f2b24c5..137a1009115 100644
--- a/pennylane/transforms/optimization/optimization_utils.py
+++ b/pennylane/transforms/optimization/optimization_utils.py
@@ -13,21 +13,6 @@
 # limitations under the License.
 """Utility functions for circuit optimization."""
 # pylint: disable=too-many-return-statements,import-outside-toplevel
-from pennylane.math import (
-    allclose,
-    arccos,
-    arctan2,
-    asarray,
-    cast_like,
-    cos,
-    is_abstract,
-    moveaxis,
-    requires_grad,
-    sin,
-    sqrt,
-    stack,
-    zeros_like,
-)
 from pennylane.ops.identity import GlobalPhase
 from pennylane.wires import Wires
 
@@ -60,21 +45,25 @@ def _try_no_fuse(angles1, angles2):
     if some angles in the input angles vanish."""
     dummy_sum = angles1 + angles2
     # moveaxis required for batched inputs
-    phi1, theta1, omega1 = moveaxis(cast_like(asarray(angles1), dummy_sum), -1, 0)
-    phi2, theta2, omega2 = moveaxis(cast_like(asarray(angles2), dummy_sum), -1, 0)
-
-    if allclose(omega1 + phi2, 0.0):
-        return stack([phi1, theta1 + theta2, omega2])
-    zero = zeros_like(phi1) + zeros_like(phi2)
-    if allclose(theta1, 0.0):
+    phi1, theta1, omega1 = qml.math.moveaxis(
+        qml.math.cast_like(qml.math.asarray(angles1), dummy_sum), -1, 0
+    )
+    phi2, theta2, omega2 = qml.math.moveaxis(
+        qml.math.cast_like(qml.math.asarray(angles2), dummy_sum), -1, 0
+    )
+
+    if qml.math.allclose(omega1 + phi2, 0.0):
+        return qml.math.stack([phi1, theta1 + theta2, omega2])
+    zero = qml.math.zeros_like(phi1) + qml.math.zeros_like(phi2)
+    if qml.math.allclose(theta1, 0.0):
         # No Y rotation in first Rot
-        if allclose(theta2, 0.0):
+        if qml.math.allclose(theta2, 0.0):
             # Z rotations only
-            return stack([phi1 + omega1 + phi2 + omega2, zero, zero])
-        return stack([phi1 + omega1 + phi2, theta2, omega2])
-    if allclose(theta2, 0.0):
+            return qml.math.stack([phi1 + omega1 + phi2 + omega2, zero, zero])
+        return qml.math.stack([phi1 + omega1 + phi2, theta2, omega2])
+    if qml.math.allclose(theta2, 0.0):
         # No Y rotation in second Rot
-        return stack([phi1, theta1, omega1 + phi2 + omega2])
+        return qml.math.stack([phi1, theta1, omega1 + phi2 + omega2])
     return None
 
 
@@ -118,35 +107,38 @@ def fuse_rot_angles(angles1, angles2):
     """
 
     if not (
-        is_abstract(angles1)
-        or is_abstract(angles2)
-        or requires_grad(angles1)
-        or requires_grad(angles2)
+        qml.math.is_abstract(angles1)
+        or qml.math.is_abstract(angles2)
+        or qml.math.requires_grad(angles1)
+        or qml.math.requires_grad(angles2)
     ):
         fused_angles = _try_no_fuse(angles1, angles2)
         if fused_angles is not None:
             return fused_angles
 
     # moveaxis required for batched inputs
-    phi1, theta1, omega1 = moveaxis(asarray(angles1), -1, 0)
-    phi2, theta2, omega2 = moveaxis(asarray(angles2), -1, 0)
-    c1, c2, s1, s2 = cos(theta1 / 2), cos(theta2 / 2), sin(theta1 / 2), sin(theta2 / 2)
+    phi1, theta1, omega1 = qml.math.moveaxis(qml.math.asarray(angles1), -1, 0)
+    phi2, theta2, omega2 = qml.math.moveaxis(qml.math.asarray(angles2), -1, 0)
+    c1, c2 = qml.math.cos(theta1 / 2), qml.math.cos(theta2 / 2)
+    s1, s2 = qml.math.sin(theta1 / 2), qml.math.sin(theta2 / 2)
 
-    mag = sqrt(c1**2 * c2**2 + s1**2 * s2**2 - 2 * c1 * c2 * s1 * s2 * cos(omega1 + phi2))
-    theta_f = 2 * arccos(mag)
+    mag = qml.math.sqrt(
+        c1**2 * c2**2 + s1**2 * s2**2 - 2 * c1 * c2 * s1 * s2 * qml.math.cos(omega1 + phi2)
+    )
+    theta_f = 2 * qml.math.arccos(mag)
 
     alpha1, beta1 = (phi1 + omega1) / 2, (phi1 - omega1) / 2
     alpha2, beta2 = (phi2 + omega2) / 2, (phi2 - omega2) / 2
 
-    alpha_arg1 = -c1 * c2 * sin(alpha1 + alpha2) - s1 * s2 * sin(beta2 - beta1)
-    alpha_arg2 = c1 * c2 * cos(alpha1 + alpha2) - s1 * s2 * cos(beta2 - beta1)
-    alpha_f = -1 * arctan2(alpha_arg1, alpha_arg2)
+    alpha_arg1 = -c1 * c2 * qml.math.sin(alpha1 + alpha2) - s1 * s2 * qml.math.sin(beta2 - beta1)
+    alpha_arg2 = c1 * c2 * qml.math.cos(alpha1 + alpha2) - s1 * s2 * qml.math.cos(beta2 - beta1)
+    alpha_f = -1 * qml.math.arctan2(alpha_arg1, alpha_arg2)
 
-    beta_arg1 = -c1 * s2 * sin(alpha1 + beta2) + s1 * c2 * sin(alpha2 - beta1)
-    beta_arg2 = c1 * s2 * cos(alpha1 + beta2) + s1 * c2 * cos(alpha2 - beta1)
-    beta_f = -1 * arctan2(beta_arg1, beta_arg2)
+    beta_arg1 = -c1 * s2 * qml.math.sin(alpha1 + beta2) + s1 * c2 * qml.math.sin(alpha2 - beta1)
+    beta_arg2 = c1 * s2 * qml.math.cos(alpha1 + beta2) + s1 * c2 * qml.math.cos(alpha2 - beta1)
+    beta_f = -1 * qml.math.arctan2(beta_arg1, beta_arg2)
 
-    return stack([alpha_f + beta_f, theta_f, alpha_f - beta_f], axis=-1)
+    return qml.math.stack([alpha_f + beta_f, theta_f, alpha_f - beta_f], axis=-1)
 
 
 def _fuse_global_phases(operations):

From 954535aabd9535ba1baa0c565ea0f0a048c3e045 Mon Sep 17 00:00:00 2001
From: dwierichs <david.wierichs@xanadu.ai>
Date: Wed, 24 Jul 2024 13:33:39 +0200
Subject: [PATCH 17/30] import fix

---
 pennylane/transforms/optimization/optimization_utils.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/pennylane/transforms/optimization/optimization_utils.py b/pennylane/transforms/optimization/optimization_utils.py
index 137a1009115..0d6781a6eb0 100644
--- a/pennylane/transforms/optimization/optimization_utils.py
+++ b/pennylane/transforms/optimization/optimization_utils.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 """Utility functions for circuit optimization."""
 # pylint: disable=too-many-return-statements,import-outside-toplevel
+import pennylane as qml
 from pennylane.ops.identity import GlobalPhase
 from pennylane.wires import Wires
 

From d23e4475b5ff530a52a750a2594ec1a04ff07f72 Mon Sep 17 00:00:00 2001
From: dwierichs <david.wierichs@xanadu.ai>
Date: Wed, 24 Jul 2024 13:49:40 +0200
Subject: [PATCH 18/30] clean out imports

---
 .../optimization/merge_rotations.py           | 25 ++++++-------------
 1 file changed, 8 insertions(+), 17 deletions(-)

diff --git a/pennylane/transforms/optimization/merge_rotations.py b/pennylane/transforms/optimization/merge_rotations.py
index 115a059f5a5..e9b01052635 100644
--- a/pennylane/transforms/optimization/merge_rotations.py
+++ b/pennylane/transforms/optimization/merge_rotations.py
@@ -15,15 +15,6 @@
 # pylint: disable=too-many-branches
 
 import pennylane as qml
-from pennylane.math import (
-    allclose,
-    cast_like,
-    get_interface,
-    is_abstract,
-    requires_grad,
-    stack,
-    zeros,
-)
 from pennylane.ops.op_math import Adjoint
 from pennylane.ops.qubit.attributes import composable_rotations
 from pennylane.queuing import QueuingManager
@@ -185,8 +176,8 @@ def stop_at(obj):
             continue
 
         # We need to use stack to get this to work and be differentiable in all interfaces
-        cumulative_angles = stack(current_gate.parameters)
-        interface = get_interface(cumulative_angles)
+        cumulative_angles = qml.math.stack(current_gate.parameters)
+        interface = qml.math.get_interface(cumulative_angles)
         # As long as there is a valid next gate, check if we can merge the angles
         while next_gate_idx is not None:
             # Get the next gate
@@ -195,10 +186,10 @@ def stop_at(obj):
             # If next gate is of the same type, we can merge the angles
             if current_gate.name == next_gate.name and current_gate.wires == next_gate.wires:
                 list_copy.pop(next_gate_idx + 1)
-                next_params = stack(next_gate.parameters, like=interface)
+                next_params = qml.math.stack(next_gate.parameters, like=interface)
                 # jax-jit does not support cast_like
-                if not is_abstract(cumulative_angles):
-                    next_params = cast_like(next_params, cumulative_angles)
+                if not qml.math.is_abstract(cumulative_angles):
+                    next_params = qml.math.cast_like(next_params, cumulative_angles)
 
                 # The Rot gate must be treated separately
                 if current_gate.name == "Rot":
@@ -217,9 +208,9 @@ def stop_at(obj):
         # apply the operation regardless of the angles. Otherwise, only apply if
         # the rotation angle is non-trivial.
         if (
-            is_abstract(cumulative_angles)
-            or requires_grad(cumulative_angles)
-            or not allclose(cumulative_angles, 0.0, atol=atol, rtol=0)
+            qml.math.is_abstract(cumulative_angles)
+            or qml.math.requires_grad(cumulative_angles)
+            or not qml.math.allclose(cumulative_angles, 0.0, atol=atol, rtol=0)
         ):
             with QueuingManager.stop_recording():
                 new_operations.append(

From c3264714c318e1c56b24ad6180d8aa3d89684aa8 Mon Sep 17 00:00:00 2001
From: dwierichs <david.wierichs@xanadu.ai>
Date: Wed, 24 Jul 2024 14:03:57 +0200
Subject: [PATCH 19/30] import stuff

---
 .../optimization/single_qubit_fusion.py          | 16 +++++++---------
 1 file changed, 7 insertions(+), 9 deletions(-)

diff --git a/pennylane/transforms/optimization/single_qubit_fusion.py b/pennylane/transforms/optimization/single_qubit_fusion.py
index 11cf265cee2..461116dcc4b 100644
--- a/pennylane/transforms/optimization/single_qubit_fusion.py
+++ b/pennylane/transforms/optimization/single_qubit_fusion.py
@@ -14,7 +14,7 @@
 """Transform for fusing sequences of single-qubit gates."""
 # pylint: disable=too-many-branches
 
-from pennylane.math import allclose, cast_like, is_abstract, requires_grad, stack
+import pennylane as qml
 from pennylane.ops.qubit import Rot
 from pennylane.queuing import QueuingManager
 from pennylane.tape import QuantumTape, QuantumTapeBatch
@@ -241,7 +241,7 @@ def qfunc(r1, r2):
         # Look for single_qubit_rot_angles; if not available, queue and move on.
         # If available, grab the angles and try to fuse.
         try:
-            cumulative_angles = stack(current_gate.single_qubit_rot_angles())
+            cumulative_angles = qml.math.stack(current_gate.single_qubit_rot_angles())
         except (NotImplementedError, AttributeError):
             new_operations.append(current_gate)
             list_copy.pop(0)
@@ -279,11 +279,9 @@ def qfunc(r1, r2):
             # 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 = qml.math.stack(next_gate.single_qubit_rot_angles())
             except (NotImplementedError, AttributeError):
                 break
-            if not is_abstract(cumulative_angles):
-                next_gate_angles = cast_like(next_gate_angles, cumulative_angles)
             cumulative_angles = fuse_rot_angles(cumulative_angles, next_gate_angles)
 
             list_copy.pop(next_gate_idx + 1)
@@ -294,10 +292,10 @@ def qfunc(r1, r2):
         # If not tracing or differentiating, check whether total rotation is trivial by checking
         # if the RY angle and the sum of the RZ angles are close to 0
         if (
-            is_abstract(cumulative_angles)
-            or requires_grad(cumulative_angles)
-            or not allclose(
-                stack([cumulative_angles[0] + cumulative_angles[2], cumulative_angles[1]]),
+            qml.math.is_abstract(cumulative_angles)
+            or qml.math.requires_grad(cumulative_angles)
+            or not qml.math.allclose(
+                qml.math.stack([cumulative_angles[0] + cumulative_angles[2], cumulative_angles[1]]),
                 0.0,
                 atol=atol,
                 rtol=0,

From 39d082787202e826052c8980b1ca16dae20bd382 Mon Sep 17 00:00:00 2001
From: dwierichs <david.wierichs@xanadu.ai>
Date: Wed, 24 Jul 2024 14:09:35 +0200
Subject: [PATCH 20/30] fix, rename angles, comment

---
 .../optimization/optimization_utils.py        | 32 +++++++++----------
 1 file changed, 16 insertions(+), 16 deletions(-)

diff --git a/pennylane/transforms/optimization/optimization_utils.py b/pennylane/transforms/optimization/optimization_utils.py
index 0d6781a6eb0..998d082cc37 100644
--- a/pennylane/transforms/optimization/optimization_utils.py
+++ b/pennylane/transforms/optimization/optimization_utils.py
@@ -41,17 +41,17 @@ def find_next_gate(wires, op_list):
     return next_gate_idx
 
 
-def _try_no_fuse(angles1, angles2):
+def _try_no_fuse(angles_1, angles_2):
     """Try to combine rotation angles without trigonometric identities
     if some angles in the input angles vanish."""
-    dummy_sum = angles1 + angles2
+    angles_1 = qml.math.asarray(angles_1)
+    angles_2 = qml.math.asarray(angles_2)
+    # This sum is only computed to obtain a dtype-coerced object that respects
+    # TensorFlow's coercion rules between Python/NumPy objects and TF objects.
+    _sum = angles_1 + angles_2
     # moveaxis required for batched inputs
-    phi1, theta1, omega1 = qml.math.moveaxis(
-        qml.math.cast_like(qml.math.asarray(angles1), dummy_sum), -1, 0
-    )
-    phi2, theta2, omega2 = qml.math.moveaxis(
-        qml.math.cast_like(qml.math.asarray(angles2), dummy_sum), -1, 0
-    )
+    phi1, theta1, omega1 = qml.math.moveaxis(qml.math.cast_like(angles_1, _sum), -1, 0)
+    phi2, theta2, omega2 = qml.math.moveaxis(qml.math.cast_like(angles_2, _sum), -1, 0)
 
     if qml.math.allclose(omega1 + phi2, 0.0):
         return qml.math.stack([phi1, theta1 + theta2, omega2])
@@ -68,7 +68,7 @@ def _try_no_fuse(angles1, angles2):
     return None
 
 
-def fuse_rot_angles(angles1, angles2):
+def fuse_rot_angles(angles_1, angles_2):
     r"""Compute the set of rotation angles that is equivalent to performing
     two successive ``qml.Rot`` operations.
 
@@ -108,18 +108,18 @@ def fuse_rot_angles(angles1, angles2):
     """
 
     if not (
-        qml.math.is_abstract(angles1)
-        or qml.math.is_abstract(angles2)
-        or qml.math.requires_grad(angles1)
-        or qml.math.requires_grad(angles2)
+        qml.math.is_abstract(angles_1)
+        or qml.math.is_abstract(angles_2)
+        or qml.math.requires_grad(angles_1)
+        or qml.math.requires_grad(angles_2)
     ):
-        fused_angles = _try_no_fuse(angles1, angles2)
+        fused_angles = _try_no_fuse(angles_1, angles_2)
         if fused_angles is not None:
             return fused_angles
 
     # moveaxis required for batched inputs
-    phi1, theta1, omega1 = qml.math.moveaxis(qml.math.asarray(angles1), -1, 0)
-    phi2, theta2, omega2 = qml.math.moveaxis(qml.math.asarray(angles2), -1, 0)
+    phi1, theta1, omega1 = qml.math.moveaxis(qml.math.asarray(angles_1), -1, 0)
+    phi2, theta2, omega2 = qml.math.moveaxis(qml.math.asarray(angles_2), -1, 0)
     c1, c2 = qml.math.cos(theta1 / 2), qml.math.cos(theta2 / 2)
     s1, s2 = qml.math.sin(theta1 / 2), qml.math.sin(theta2 / 2)
 

From 65c3e8bfde0b82de724d72662d4540e1d45a12da Mon Sep 17 00:00:00 2001
From: David Wierichs <david.wierichs@xanadu.ai>
Date: Thu, 25 Jul 2024 22:28:49 +0200
Subject: [PATCH 21/30] Apply suggestions from code review

Co-authored-by: Vincent Michaud-Rioux <vincent.michaud-rioux@xanadu.ai>
---
 doc/releases/changelog-dev.md                            | 2 +-
 pennylane/transforms/optimization/merge_rotations.py     | 2 +-
 pennylane/transforms/optimization/single_qubit_fusion.py | 8 ++++----
 3 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 516dbacd451..ea616b1fd9f 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -33,7 +33,7 @@
 
 <h3>Improvements 🛠</h3>
 
-* `fuse_rot_angles` respects now the global phase of the combined rotations.
+* `fuse_rot_angles` now respects the global phase of the combined rotations.
   [(#6031)](https://github.com/PennyLaneAI/pennylane/pull/6031)
 
 * `StateMP.process_state` defines rules in `cast_to_complex` for complex casting, avoiding a superfluous state vector copy in Lightning simulations
diff --git a/pennylane/transforms/optimization/merge_rotations.py b/pennylane/transforms/optimization/merge_rotations.py
index ace24b71138..81be5b33988 100644
--- a/pennylane/transforms/optimization/merge_rotations.py
+++ b/pennylane/transforms/optimization/merge_rotations.py
@@ -77,7 +77,7 @@ def circuit(x, y, z):
         When merging two :class:`~.pennylane.Rot` gates, there are a number of details to consider:
 
         First, the output angles are not always uniquely defined, because Euler angles are not
-        defined uniquely for some rotations. ``merge_rotations`` makes a particular choice in
+        unique for some rotations. ``merge_rotations`` makes a particular choice in
         this case.
 
         Second, ``merge_rotations`` is not differentiable everywhere when used on ``Rot``.
diff --git a/pennylane/transforms/optimization/single_qubit_fusion.py b/pennylane/transforms/optimization/single_qubit_fusion.py
index 298ab29ef64..89963a7c026 100644
--- a/pennylane/transforms/optimization/single_qubit_fusion.py
+++ b/pennylane/transforms/optimization/single_qubit_fusion.py
@@ -36,7 +36,7 @@ def single_qubit_fusion(
     (on the same qubit) with that property defined will be fused into one ``Rot``.
 
     Args:
-        tape (QNode or QuantumTape or callable): A quantum circuit.
+        tape (QNode or QuantumTape or Callable): A quantum circuit.
         atol (float): An absolute tolerance for which to apply a rotation after
             fusion. After fusion of gates, if the fused angles :math:`\theta` are such that
             :math:`|\theta|\leq \text{atol}`, no rotation gate will be applied.
@@ -45,7 +45,7 @@ def single_qubit_fusion(
             be fused will be fused.
 
     Returns:
-        qnode (QNode) or quantum function (callable) or tuple[List[QuantumTape], callable]:
+        qnode (QNode) or quantum function (Callable) or tuple[List[QuantumTape], Callable]:
         The transformed circuit as described in :func:`qml.transform <pennylane.transform>`.
 
     **Example**
@@ -160,7 +160,7 @@ def qfunc(r1, r2):
             \theta_f = 2\arccos(|x|).
 
         We can extract the angle combination :math:`\alpha_f` from :math:`x` above as well via
-        :math:`\operatorname{arg}(x)`, which can be readily computed with :math:`arctan`:
+        :math:`\operatorname{arg}(x)`, which can be readily computed with :math:`\arctan`:
 
         .. math::
 
@@ -218,7 +218,7 @@ def qfunc(r1, r2):
         **Differentiability:**
 
         The function derived above is differentiable almost everywhere.
-        In particular, there are two problematic scenarios at which the derivative is not defined:
+        In particular, there are two problematic scenarios at which the derivative is not defined.
         First, the square root is not differentiable at :math:`0`, making all input angles with
         :math:`|x|=0` singular. Second, :math:`\arccos` is not differentiable at :math:`1`, making
         all input angles with :math:`|x|=1` singular.

From ccfcb34db1c5f8179cd91cf035b5e0ba23c97629 Mon Sep 17 00:00:00 2001
From: dwierichs <david.wierichs@xanadu.ai>
Date: Thu, 25 Jul 2024 22:48:31 +0200
Subject: [PATCH 22/30] squash tests

---
 .../test_optimization_utils.py                | 78 +++++++------------
 1 file changed, 26 insertions(+), 52 deletions(-)

diff --git a/tests/transforms/test_optimization/test_optimization_utils.py b/tests/transforms/test_optimization/test_optimization_utils.py
index 945797999e1..0872bc8fc30 100644
--- a/tests/transforms/test_optimization/test_optimization_utils.py
+++ b/tests/transforms/test_optimization/test_optimization_utils.py
@@ -75,24 +75,29 @@ class TestRotGateFusion:
         ([0.9, np.pi / 2, np.pi / 2], [-np.pi / 2, -np.pi / 2, 0.0]),
     ]
 
-    @staticmethod
-    def mat_from_prod(angles_1, angles_2):
+    def run_interface_test(self, angles_1, angles_2):
+        """Execute standard test lines for different interfaces and batch tests.
+        Note that the transpose calls only are relevant for tests with batching."""
+
         def original_ops():
-            qml.Rot(*angles_1, wires=0)
-            qml.Rot(*angles_2, wires=0)
+            qml.Rot(*qml.math.transpose(angles_1), wires=0)
+            qml.Rot(*qml.math.transpose(angles_2), wires=0)
+
+        matrix_expected = qml.matrix(original_ops, [0])()  # pylint:disable=too-many-function-args
+
+        fused_angles = fuse_rot_angles(angles_1, angles_2)
+        # The reshape is only used in the _mixed_batching test. Otherwise it is irrelevant.
+        matrix_obtained = qml.Rot(
+            *qml.math.transpose(qml.math.reshape(fused_angles, (-1, 3))), wires=0
+        ).matrix()
 
-        return qml.matrix(original_ops, [0])()  # pylint:disable=too-many-function-args
+        assert qml.math.allclose(matrix_expected, matrix_obtained)
 
     @pytest.mark.parametrize("angles_1, angles_2", generic_test_angles)
     def test_full_rot_fusion_numpy(self, angles_1, angles_2):
         """Test that the fusion of two Rot gates has the same effect as
         applying the Rots sequentially."""
-
-        matrix_expected = self.mat_from_prod(angles_1, angles_2)
-        fused_angles = fuse_rot_angles(angles_1, angles_2)
-        matrix_obtained = qml.Rot(*fused_angles, wires=0).matrix()
-
-        assert qml.math.allclose(matrix_expected, matrix_obtained)
+        self.run_interface_test(angles_1, angles_2)
 
     mixed_batched_angles = [
         ([[0.4, 0.1, 0.0], [0.7, 0.2, 0.1]], [-0.9, 1.2, 0.6]),  # (2, None)
@@ -110,14 +115,9 @@ def test_full_rot_fusion_mixed_batching(self, angles_1, angles_2):
         applying the Rots sequentially when the input angles are batched
         with mixed batching shapes."""
 
-        fused_angles = fuse_rot_angles(angles_1, angles_2)
-
         reshaped_angles_1 = np.reshape(angles_1, (-1, 3) if np.ndim(angles_1) > 1 else (3,))
         reshaped_angles_2 = np.reshape(angles_2, (-1, 3) if np.ndim(angles_2) > 1 else (3,))
-        matrix_expected = self.mat_from_prod(reshaped_angles_1.T, reshaped_angles_2.T)
-        matrix_obtained = qml.Rot(*fused_angles.reshape((-1, 3)).T, wires=0).matrix()
-
-        assert qml.math.allclose(matrix_expected, matrix_obtained)
+        self.run_interface_test(reshaped_angles_1, reshaped_angles_2)
 
     @pytest.mark.autograd
     @pytest.mark.parametrize("angles_1, angles_2", generic_test_angles)
@@ -126,11 +126,7 @@ def test_full_rot_fusion_autograd(self, angles_1, angles_2):
         applying the Rots sequentially, in Autograd."""
 
         angles_1, angles_2 = qml.numpy.array(angles_1), qml.numpy.array(angles_1)
-        matrix_expected = self.mat_from_prod(angles_1, angles_2)
-        fused_angles = fuse_rot_angles(angles_1, angles_2)
-        matrix_obtained = qml.Rot(*fused_angles, wires=0).matrix()
-
-        assert qml.math.allclose(matrix_expected, matrix_obtained)
+        self.run_interface_test(angles_1, angles_2)
 
     @pytest.mark.tf
     @pytest.mark.parametrize("angles_1, angles_2", generic_test_angles)
@@ -139,14 +135,9 @@ def test_full_rot_fusion_tensorflow(self, angles_1, angles_2):
         applying the Rots sequentially, in Tensorflow."""
         import tensorflow as tf
 
-        angles_1, angles_2 = tf.Variable(angles_1, dtype=tf.float64), tf.Variable(
-            angles_1, dtype=tf.float64
-        )
-        matrix_expected = self.mat_from_prod(angles_1, angles_2)
-        fused_angles = fuse_rot_angles(angles_1, angles_2)
-        matrix_obtained = qml.Rot(*fused_angles, wires=0).matrix()
-
-        assert qml.math.allclose(matrix_expected, matrix_obtained)
+        angles_1 = tf.Variable(angles_1, dtype=tf.float64)
+        angles_2 = tf.Variable(angles_2, dtype=tf.float64)
+        self.run_interface_test(angles_1, angles_2)
 
     @pytest.mark.torch
     @pytest.mark.parametrize("angles_1, angles_2", generic_test_angles)
@@ -155,14 +146,9 @@ def test_full_rot_fusion_torch(self, angles_1, angles_2):
         applying the Rots sequentially, in torch."""
         import torch
 
-        angles_1, angles_2 = torch.tensor(
-            angles_1, requires_grad=True, dtype=torch.float64
-        ), torch.tensor(angles_1, requires_grad=True, dtype=torch.float64)
-        matrix_expected = self.mat_from_prod(angles_1, angles_2)
-        fused_angles = fuse_rot_angles(angles_1, angles_2)
-        matrix_obtained = qml.Rot(*fused_angles, wires=0).matrix()
-
-        assert qml.math.allclose(matrix_expected, matrix_obtained)
+        angles_1 = torch.tensor(angles_1, requires_grad=True, dtype=torch.float64)
+        angles_2 = torch.tensor(angles_2, requires_grad=True, dtype=torch.float64)
+        self.run_interface_test(angles_1, angles_2)
 
     @pytest.mark.jax
     @pytest.mark.parametrize("angles_1, angles_2", generic_test_angles)
@@ -172,11 +158,7 @@ def test_full_rot_fusion_jax(self, angles_1, angles_2):
         import jax
 
         angles_1, angles_2 = jax.numpy.array(angles_1), jax.numpy.array(angles_1)
-        matrix_expected = self.mat_from_prod(angles_1, angles_2)
-        fused_angles = fuse_rot_angles(angles_1, angles_2)
-        matrix_obtained = qml.Rot(*fused_angles, wires=0).matrix()
-
-        assert qml.math.allclose(matrix_expected, matrix_obtained)
+        self.run_interface_test(angles_1, angles_2)
 
     @pytest.mark.slow
     def test_full_rot_fusion_special_angles(self):
@@ -188,15 +170,7 @@ def test_full_rot_fusion_special_angles(self):
         special_points = np.array([3 / 2, 1, 1 / 2, 0, -1 / 2, -1, -3 / 2]) * np.pi
         special_angles = np.array(list(product(special_points, repeat=6))).reshape((-1, 2, 3))
         angles_1, angles_2 = np.transpose(special_angles, (1, 0, 2))
-
-        # Transpose to bring size-3 axis to front
-        matrix_expected = self.mat_from_prod(angles_1.T, angles_2.T)
-
-        fused_angles = fuse_rot_angles(angles_1, angles_2)
-        # Transpose to bring size-3 axis to front
-        matrix_obtained = qml.Rot(*fused_angles.T, wires=0).matrix()
-
-        assert qml.math.allclose(matrix_expected, matrix_obtained)
+        self.run_interface_test(angles_1, angles_2)
 
     @pytest.mark.slow
     @pytest.mark.jax

From f4f1c7f79e3957345794075516245e5365720b8d Mon Sep 17 00:00:00 2001
From: dwierichs <david.wierichs@xanadu.ai>
Date: Fri, 26 Jul 2024 09:57:33 +0200
Subject: [PATCH 23/30] tiny

---
 pennylane/transforms/optimization/merge_rotations.py     | 2 +-
 pennylane/transforms/optimization/optimization_utils.py  | 2 +-
 pennylane/transforms/optimization/single_qubit_fusion.py | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/pennylane/transforms/optimization/merge_rotations.py b/pennylane/transforms/optimization/merge_rotations.py
index 81be5b33988..7945dd23408 100644
--- a/pennylane/transforms/optimization/merge_rotations.py
+++ b/pennylane/transforms/optimization/merge_rotations.py
@@ -76,7 +76,7 @@ def circuit(x, y, z):
 
         When merging two :class:`~.pennylane.Rot` gates, there are a number of details to consider:
 
-        First, the output angles are not always uniquely defined, because Euler angles are not
+        First, the output angles are not always defined uniquely, because Euler angles are not
         unique for some rotations. ``merge_rotations`` makes a particular choice in
         this case.
 
diff --git a/pennylane/transforms/optimization/optimization_utils.py b/pennylane/transforms/optimization/optimization_utils.py
index 998d082cc37..a3dbbffef13 100644
--- a/pennylane/transforms/optimization/optimization_utils.py
+++ b/pennylane/transforms/optimization/optimization_utils.py
@@ -90,7 +90,7 @@ def fuse_rot_angles(angles_1, angles_2):
     .. note::
 
         The output angles are not always defined uniquely because Euler angles are not
-        defined uniquely for some rotations. ``fuse_rot_angles`` makes a particular
+        unique for some rotations. ``fuse_rot_angles`` makes a particular
         choice in this case.
 
     .. warning::
diff --git a/pennylane/transforms/optimization/single_qubit_fusion.py b/pennylane/transforms/optimization/single_qubit_fusion.py
index 89963a7c026..7e6a43dbe25 100644
--- a/pennylane/transforms/optimization/single_qubit_fusion.py
+++ b/pennylane/transforms/optimization/single_qubit_fusion.py
@@ -71,7 +71,7 @@ def qfunc(r1, r2):
     .. note::
 
         The fused angles between two sets of rotation angles are not always defined uniquely
-        because Euler angles are not defined uniquely for some rotations. ``single_qubit_fusion``
+        because Euler angles are not unique for some rotations. ``single_qubit_fusion``
         makes a particular choice in this case.
 
     .. warning::

From 7c37e1488c599b7e2e934ee6b56e4f45556af9f0 Mon Sep 17 00:00:00 2001
From: dwierichs <david.wierichs@xanadu.ai>
Date: Fri, 26 Jul 2024 10:12:55 +0200
Subject: [PATCH 24/30] expand doc

---
 .../optimization/optimization_utils.py        | 10 +++----
 .../optimization/single_qubit_fusion.py       | 26 +++++++++++++++----
 2 files changed, 26 insertions(+), 10 deletions(-)

diff --git a/pennylane/transforms/optimization/optimization_utils.py b/pennylane/transforms/optimization/optimization_utils.py
index a3dbbffef13..671738182fc 100644
--- a/pennylane/transforms/optimization/optimization_utils.py
+++ b/pennylane/transforms/optimization/optimization_utils.py
@@ -44,8 +44,6 @@ def find_next_gate(wires, op_list):
 def _try_no_fuse(angles_1, angles_2):
     """Try to combine rotation angles without trigonometric identities
     if some angles in the input angles vanish."""
-    angles_1 = qml.math.asarray(angles_1)
-    angles_2 = qml.math.asarray(angles_2)
     # This sum is only computed to obtain a dtype-coerced object that respects
     # TensorFlow's coercion rules between Python/NumPy objects and TF objects.
     _sum = angles_1 + angles_2
@@ -55,11 +53,11 @@ def _try_no_fuse(angles_1, angles_2):
 
     if qml.math.allclose(omega1 + phi2, 0.0):
         return qml.math.stack([phi1, theta1 + theta2, omega2])
-    zero = qml.math.zeros_like(phi1) + qml.math.zeros_like(phi2)
     if qml.math.allclose(theta1, 0.0):
         # No Y rotation in first Rot
         if qml.math.allclose(theta2, 0.0):
             # Z rotations only
+            zero = qml.math.zeros_like(phi1) + qml.math.zeros_like(phi2)
             return qml.math.stack([phi1 + omega1 + phi2 + omega2, zero, zero])
         return qml.math.stack([phi1 + omega1 + phi2, theta2, omega2])
     if qml.math.allclose(theta2, 0.0):
@@ -106,6 +104,8 @@ def fuse_rot_angles(angles_1, angles_2):
     See the documentation of :func:`~.pennylane.transforms.single_qubit_fusion` for a
     mathematical derivation of this function.
     """
+    angles_1 = qml.math.asarray(angles_1)
+    angles_2 = qml.math.asarray(angles_2)
 
     if not (
         qml.math.is_abstract(angles_1)
@@ -118,8 +118,8 @@ def fuse_rot_angles(angles_1, angles_2):
             return fused_angles
 
     # moveaxis required for batched inputs
-    phi1, theta1, omega1 = qml.math.moveaxis(qml.math.asarray(angles_1), -1, 0)
-    phi2, theta2, omega2 = qml.math.moveaxis(qml.math.asarray(angles_2), -1, 0)
+    phi1, theta1, omega1 = qml.math.moveaxis(angles_1, -1, 0)
+    phi2, theta2, omega2 = qml.math.moveaxis(angles_2, -1, 0)
     c1, c2 = qml.math.cos(theta1 / 2), qml.math.cos(theta2 / 2)
     s1, s2 = qml.math.sin(theta1 / 2), qml.math.sin(theta2 / 2)
 
diff --git a/pennylane/transforms/optimization/single_qubit_fusion.py b/pennylane/transforms/optimization/single_qubit_fusion.py
index 7e6a43dbe25..71cfe5e8571 100644
--- a/pennylane/transforms/optimization/single_qubit_fusion.py
+++ b/pennylane/transforms/optimization/single_qubit_fusion.py
@@ -131,7 +131,7 @@ def qfunc(r1, r2):
             e^{-i\beta_j}s_j & e^{i\alpha_j}c_j
             \end{bmatrix},
 
-        where we introduced abbreviations :math:`\alpha_j`, :math:`\beta_j`,
+        where we introduced abbreviations :math:`\alpha_j,\beta_j=\frac{\phi_j\pm\omega_j}{2}`,
         :math:`c_j=\cos(\theta_j / 2)` and :math:`s_j=\sin(\theta_j / 2)` for notational brevity.
         The upper left entry of the matrix product
         :math:`R(\phi_2,\theta_2,\omega_2)R(\phi_1,\theta_1,\omega_1)` reads
@@ -140,7 +140,7 @@ def qfunc(r1, r2):
 
             x = e^{-i(\alpha_2+\alpha_1)} c_2 c_1 - e^{i(\beta_2-\beta_1)} s_2 s_1
 
-        and should equal :math:`e^{-i(\alpha_f)/2}c_f` for the fused rotation angles.
+        and should equal :math:`e^{-i\alpha_f}c_f` for the fused rotation angles.
         This means that we can obtain :math:`\theta_f` from the magnitude of the matrix product
         entry above, choosing :math:`c_f=\cos(\theta_f / 2)` to be non-negative:
 
@@ -159,7 +159,7 @@ def qfunc(r1, r2):
 
             \theta_f = 2\arccos(|x|).
 
-        We can extract the angle combination :math:`\alpha_f` from :math:`x` above as well via
+        We can also extract the angle combination :math:`\alpha_f` from :math:`x` via
         :math:`\operatorname{arg}(x)`, which can be readily computed with :math:`\arctan`:
 
         .. math::
@@ -171,7 +171,7 @@ def qfunc(r1, r2):
 
         We can use the standard numerical function :math:`\operatorname{arctan2}`, which
         computes :math:`\arctan(x_1/x_2)` from :math:`x_1` and :math:`x_2` while handling
-        special points suitably to obtain the argument of the underlying complex number
+        special points suitably, to obtain the argument of the underlying complex number
         :math:`x_2 + x_1 i`.
 
         Finally, to obtain :math:`\beta_f`, we need a second element of the matrix product from
@@ -194,11 +194,27 @@ def qfunc(r1, r2):
 
         .. math::
 
-            \phi_f = \alpha_f + \beta_f
+            \phi_f = \alpha_f + \beta_f\qquad
             \omega_f = \alpha_f - \beta_f
 
         and are done.
 
+        **Special cases:**
+
+        There are a number of special cases for which we can skip the computation above and
+        can combine rotation angles directly.
+
+        1. If :math:`\omega_1=\phi_2=0`, we can simply merge the ``RY`` rotation angles
+           :math:`\theta_j` and obtain :math:`(\phi_1, \theta_1+\theta_2, \omega_2)`.
+
+        2. If :math:`\theta_j=0`, we can merge the two ``RZ`` rotations of the same ``Rot``
+           and obtain :math:`(\phi_1+\omega_1+\phi_2, \theta_2, \omega_2)` or
+           :math:`(\phi_1, \theta_1, \omega_1+\phi_2+\omega_2)`. If both ``RY`` angles vanish
+           we get :math:`(\phi_1+\omega_1+\phi_2+\omega_2, 0, 0)`.
+
+        Note that this optimization is not performed for differentiable input parameters,
+        in order to maintain differentiability.
+
         **Mathematical properties:**
 
         All functions above are well-defined on the domain we are using them on,

From d172a89b8933f3bc1389fc4ea6e2d46ac7a9fc2f Mon Sep 17 00:00:00 2001
From: David Wierichs <david.wierichs@xanadu.ai>
Date: Mon, 29 Jul 2024 21:05:31 +0200
Subject: [PATCH 25/30] Apply suggestions from code review

Co-authored-by: Korbinian Kottmann <43949391+Qottmann@users.noreply.github.com>
---
 pennylane/transforms/optimization/single_qubit_fusion.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pennylane/transforms/optimization/single_qubit_fusion.py b/pennylane/transforms/optimization/single_qubit_fusion.py
index 71cfe5e8571..ca55c88162f 100644
--- a/pennylane/transforms/optimization/single_qubit_fusion.py
+++ b/pennylane/transforms/optimization/single_qubit_fusion.py
@@ -169,7 +169,7 @@ def qfunc(r1, r2):
             {c_1c_2\cos(\alpha_1+\alpha_2)-s_1s_2\cos(\beta_2-\beta_1)}
             \right).
 
-        We can use the standard numerical function :math:`\operatorname{arctan2}`, which
+        We can use the standard numerical function ``arctan2``, which
         computes :math:`\arctan(x_1/x_2)` from :math:`x_1` and :math:`x_2` while handling
         special points suitably, to obtain the argument of the underlying complex number
         :math:`x_2 + x_1 i`.

From b3faa45aa9d3a299625ac2f7ee97097e358d244b Mon Sep 17 00:00:00 2001
From: David Wierichs <david.wierichs@xanadu.ai>
Date: Tue, 30 Jul 2024 16:03:54 +0200
Subject: [PATCH 26/30] Apply suggestions from code review

Co-authored-by: Thomas R. Bromley <49409390+trbromley@users.noreply.github.com>
---
 pennylane/transforms/optimization/merge_rotations.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pennylane/transforms/optimization/merge_rotations.py b/pennylane/transforms/optimization/merge_rotations.py
index 7945dd23408..9f5a5770d3f 100644
--- a/pennylane/transforms/optimization/merge_rotations.py
+++ b/pennylane/transforms/optimization/merge_rotations.py
@@ -83,7 +83,7 @@ def circuit(x, y, z):
         Second, ``merge_rotations`` is not differentiable everywhere when used on ``Rot``.
         It has singularities for specific rotation angles where the derivative will be NaN.
 
-        Finally, this function can be numerically unstable at singular points.
+        Finally, this function can be numerically unstable near singular points.
         It is therefore recommended to use it with 64-bit floating point precision angles.
 
         For a mathematical derivation of the fusion of two ``Rot`` gates, see the documentation

From b7354186e2adc1c3bce6b0d1a737ac88560d93de Mon Sep 17 00:00:00 2001
From: dwierichs <david.wierichs@xanadu.ai>
Date: Thu, 1 Aug 2024 21:44:50 +0200
Subject: [PATCH 27/30] interface tests

---
 .../optimization/optimization_utils.py        |   6 +-
 .../test_optimization_utils.py                | 169 ++++++++++++++----
 2 files changed, 143 insertions(+), 32 deletions(-)

diff --git a/pennylane/transforms/optimization/optimization_utils.py b/pennylane/transforms/optimization/optimization_utils.py
index 671738182fc..dbc845234ee 100644
--- a/pennylane/transforms/optimization/optimization_utils.py
+++ b/pennylane/transforms/optimization/optimization_utils.py
@@ -118,8 +118,10 @@ def fuse_rot_angles(angles_1, angles_2):
             return fused_angles
 
     # moveaxis required for batched inputs
-    phi1, theta1, omega1 = qml.math.moveaxis(angles_1, -1, 0)
-    phi2, theta2, omega2 = qml.math.moveaxis(angles_2, -1, 0)
+    angles_1 = qml.math.moveaxis(angles_1, -1, 0)
+    angles_2 = qml.math.moveaxis(angles_2, -1, 0)
+    phi1, theta1, omega1 = angles_1[0], angles_1[1], angles_1[2]
+    phi2, theta2, omega2 = angles_2[0], angles_2[1], angles_2[2]
     c1, c2 = qml.math.cos(theta1 / 2), qml.math.cos(theta2 / 2)
     s1, s2 = qml.math.sin(theta1 / 2), qml.math.sin(theta2 / 2)
 
diff --git a/tests/transforms/test_optimization/test_optimization_utils.py b/tests/transforms/test_optimization/test_optimization_utils.py
index 0872bc8fc30..0b15342c406 100644
--- a/tests/transforms/test_optimization/test_optimization_utils.py
+++ b/tests/transforms/test_optimization/test_optimization_utils.py
@@ -16,6 +16,7 @@
 """
 # pylint: disable=too-few-public-methods
 
+from functools import partial
 from itertools import product
 
 import pytest
@@ -172,22 +173,9 @@ def test_full_rot_fusion_special_angles(self):
         angles_1, angles_2 = np.transpose(special_angles, (1, 0, 2))
         self.run_interface_test(angles_1, angles_2)
 
-    @pytest.mark.slow
-    @pytest.mark.jax
-    @pytest.mark.parametrize("use_jit", [False, True])
-    def test_full_rot_fusion_jacobian(self, use_jit):
-        """Test the Jacobian of the rotation angle fusion. Uses batching for performance reasons.
-        For known sources of singularities, the Jacobian is checked to indeed return NaN. These
-        sources are related to the absolute value of the upper left entry of the matrix product:
-         - If it is 1, the derivative of arccos becomes infinite (evaluated at 1), and
-         - if its square is 0, the derivative of sqrt becomes infinite (evaluated at 0).
-        """
-        import jax
-
-        special_points = np.array([3 / 2, 1, 1 / 2, 0, -1 / 2, -1, -3 / 2]) * np.pi
-        special_angles = np.array(list(product(special_points, repeat=6))).reshape((-1, 2, 3))
-        random_angles = np.random.random((1000, 2, 3))
-        all_angles = jax.numpy.concatenate([special_angles, random_angles], dtype=complex)
+    def run_jacobian_test(self, all_angles, jac_fn, is_batched, jit_fn=None, array_fn=None):
+        """Execute standard test lines for testing Jacobians with different interfaces.
+        #Note that the transpose calls only are relevant for tests with batching."""
 
         def mat_from_prod(angles):
             def original_ops():
@@ -197,28 +185,149 @@ def original_ops():
 
             return qml.matrix(original_ops, [0])()  # pylint:disable=too-many-function-args
 
-        def _mat_from_fuse(angles):
+        def mat_from_fuse(angles):
             angles1, angles2 = angles[..., 0, :], angles[..., 1, :]
-            fused_angles = fuse_rot_angles(angles1, angles2)
-            return qml.Rot(*fused_angles.T, wires=0).matrix()
+            fused_angles = qml.math.transpose(fuse_rot_angles(angles1, angles2))
+            return qml.Rot(fused_angles[0], fused_angles[1], fused_angles[2], wires=0).matrix()
 
-        mat_from_fuse = jax.jit(_mat_from_fuse) if use_jit else _mat_from_fuse
+        if jit_fn is not None:
+            mat_from_fuse = jit_fn(mat_from_fuse)
 
-        # Need holomorphic derivatives because the output matrices are complex-valued
-        jac_from_prod = jax.vmap(jax.jacobian(mat_from_prod, holomorphic=True))(all_angles)
-        jac_from_fuse = jax.vmap(jax.jacobian(mat_from_fuse, holomorphic=True))(all_angles)
+        if is_batched:
+            jac_from_prod = jac_fn(mat_from_prod)(all_angles)
+            jac_from_fuse = jac_fn(mat_from_fuse)(all_angles)
+        else:
+            jac_from_prod = qml.math.stack([jac_fn(mat_from_prod)(a) for a in all_angles])
+            jac_from_fuse = qml.math.stack([jac_fn(mat_from_fuse)(a) for a in all_angles])
+
+        if array_fn is not None:
+            # Convert to vanilla numpy
+            all_angles = array_fn(all_angles)
 
         # expected failures based on the sources mentioned in the docstring above.
-        thetas = all_angles[..., 1].T
-        (c1, c2), (s1, s2) = np.cos(thetas / 2), np.sin(thetas / 2)
+        thetas = qml.math.transpose(all_angles[..., 1])
+        (c1, c2), (s1, s2) = qml.math.cos(thetas / 2), qml.math.sin(thetas / 2)
         omega1 = all_angles[:, 0, 2]
         phi2 = all_angles[:, 1, 0]
         # squared absolute value of the relevant entry of the product of the two rotation matrices
-        pre_mag = c1**2 * c2**2 + s1**2 * s2**2 - 2 * c1 * c2 * s1 * s2 * np.cos(omega1 + phi2)
+        pre_mag = (
+            c1**2 * c2**2 + s1**2 * s2**2 - 2 * c1 * c2 * s1 * s2 * qml.math.cos(omega1 + phi2)
+        )
         # Compute condition for the two error sources combined
-        error_sources = (np.abs(pre_mag - 1) < 1e-12) + (pre_mag == 0j)
+        error_sources = (qml.math.abs(pre_mag - 1) < 1e-12) | (pre_mag == 0)
+
+        print(jac_from_fuse)
+        print(jac_from_prod)
+        print(error_sources)
 
         assert qml.math.allclose(jac_from_prod[~error_sources], jac_from_fuse[~error_sources])
-        assert qml.math.all(
-            qml.math.any(qml.math.isnan(jac_from_fuse[error_sources]), axis=[1, 2, 3, 4])
-        )
+        nans = qml.math.isnan(jac_from_fuse[error_sources])
+        nans = qml.math.reshape(nans, (len(nans), -1))
+        assert qml.math.all(qml.math.any(nans, axis=1))
+
+    @pytest.mark.slow
+    @pytest.mark.jax
+    @pytest.mark.parametrize("use_jit", [False, True])
+    def test_jacobian_jax(self, use_jit):
+        """Test the Jacobian of the rotation angle fusion with JAX. Uses batching for performance
+        reasons. For known sources of singularities, the Jacobian is checked to indeed return NaN.
+        These sources are related to the absolute value of the upper left entry of the matrix
+        product:
+         - If it is 1, the derivative of arccos becomes infinite (evaluated at 1), and
+         - if its square is 0, the derivative of sqrt becomes infinite (evaluated at 0).
+        """
+        import jax
+
+        special_points = np.array([3 / 2, 1, 1 / 2, 0, -1 / 2, -1, -3 / 2]) * np.pi
+        special_angles = np.array(list(product(special_points, repeat=6))).reshape((-1, 2, 3))
+        random_angles = np.random.random((1000, 2, 3))
+        # Need holomorphic derivatives and complex inputs because the output matrices are complex
+        all_angles = jax.numpy.concatenate([special_angles, random_angles], dtype=complex)
+        jac_fn = lambda fn: jax.vmap(jax.jacobian(fn, holomorphic=True))
+        jit_fn = jax.jit if use_jit else None
+        self.run_jacobian_test(all_angles, jac_fn, is_batched=True, jit_fn=jit_fn)
+
+    @pytest.mark.slow
+    @pytest.mark.torch
+    def test_jacobian_torch(self):
+        """Test the Jacobian of the rotation angle fusion with torch.
+        For known sources of singularities, the Jacobian is checked to indeed return NaN.
+        These sources are related to the absolute value of the upper left entry of the matrix
+        product:
+         - If it is 1, the derivative of arccos becomes infinite (evaluated at 1), and
+         - if its square is 0, the derivative of sqrt becomes infinite (evaluated at 0).
+        """
+        import torch
+
+        # Testing fewer points than with batching to limit test runtimes
+        special_points = np.array([1, 1 / 2, 0, -1]) * np.pi
+        special_angles = np.array(list(product(special_points, repeat=6))).reshape((-1, 2, 3))
+        random_angles = np.random.random((100, 2, 3))
+        all_angles = np.concatenate([special_angles, random_angles])
+
+        # Need holomorphic derivatives and complex inputs because the output matrices are complex
+        all_angles = torch.tensor(all_angles, requires_grad=True)
+        jac_fn = lambda fn: (lambda *args: torch.autograd.functional.jacobian(fn, args))
+        array_fn = lambda x: x.detach().numpy()
+        self.run_jacobian_test(all_angles, jac_fn, is_batched=False, array_fn=array_fn)
+
+    @pytest.mark.slow
+    @pytest.mark.autograd
+    def test_jacobian_autograd(self):
+        """Test the Jacobian of the rotation angle fusion with Autograd.
+        For known sources of singularities, the Jacobian is checked to indeed return NaN.
+        These sources are related to the absolute value of the upper left entry of the matrix
+        product:
+         - If it is 1, the derivative of arccos becomes infinite (evaluated at 1), and
+         - if its square is 0, the derivative of sqrt becomes infinite (evaluated at 0).
+        """
+        special_points = np.array([1, 0, -1]) * np.pi
+        special_angles = np.array(list(product(special_points, repeat=6))).reshape((-1, 2, 3))
+        random_angles = np.random.random((100, 2, 3))
+        # Need holomorphic derivatives and complex inputs because the output matrices are complex
+        all_angles = qml.numpy.concatenate([special_angles, random_angles], requires_grad=True)
+
+        def jacobian(fn):
+            real_fn = lambda *args: qml.math.real(fn(*args))
+            imag_fn = lambda *args: qml.math.imag(fn(*args))
+            real_jac_fn = qml.jacobian(real_fn)
+            imag_jac_fn = qml.jacobian(imag_fn)
+            return lambda *args: real_jac_fn(*args) + 1j * imag_jac_fn(*args)
+
+        self.run_jacobian_test(all_angles, jacobian, is_batched=False)
+
+    @pytest.mark.skip
+    @pytest.mark.slow
+    @pytest.mark.tf
+    def test_jacobian_tf(self):
+        """Test the Jacobian of the rotation angle fusion with TensorFlow.
+        For known sources of singularities, the Jacobian is checked to indeed return NaN.
+        These sources are related to the absolute value of the upper left entry of the matrix
+        product:
+         - If it is 1, the derivative of arccos becomes infinite (evaluated at 1), and
+         - if its square is 0, the derivative of sqrt becomes infinite (evaluated at 0).
+        """
+        import tensorflow as tf
+
+        # Testing fewer points than with batching to limit test runtimes
+        special_points = np.array([0, 1]) * np.pi
+        special_angles = np.array(list(product(special_points, repeat=6))).reshape((-1, 2, 3))
+        random_angles = np.random.random((3, 2, 3))
+        all_angles = np.concatenate([special_angles, random_angles])
+
+        def jacobian(fn):
+
+            def jac_fn(arg):
+                arg = tf.Variable(arg)
+                with tf.GradientTape() as t:
+                    out = fn(arg)
+                return t.jacobian(out, arg)
+
+            def vmapped_jac_fn(arg):
+                return tf.vectorized_map(jac_fn, arg)
+
+            return jac_fn
+
+        # Need holomorphic derivatives and complex inputs because the output matrices are complex
+        all_angles = tf.Variable(all_angles)
+        self.run_jacobian_test(all_angles, jacobian, is_batched=False)

From 931da7d49d3594f6f95277df80228678102500cb Mon Sep 17 00:00:00 2001
From: dwierichs <david.wierichs@xanadu.ai>
Date: Fri, 2 Aug 2024 11:38:45 +0200
Subject: [PATCH 28/30] update torch

---
 .../test_optimization_utils.py                | 19 +++++++++++--------
 1 file changed, 11 insertions(+), 8 deletions(-)

diff --git a/tests/transforms/test_optimization/test_optimization_utils.py b/tests/transforms/test_optimization/test_optimization_utils.py
index 0b15342c406..6596364aff0 100644
--- a/tests/transforms/test_optimization/test_optimization_utils.py
+++ b/tests/transforms/test_optimization/test_optimization_utils.py
@@ -216,10 +216,6 @@ def mat_from_fuse(angles):
         # Compute condition for the two error sources combined
         error_sources = (qml.math.abs(pre_mag - 1) < 1e-12) | (pre_mag == 0)
 
-        print(jac_from_fuse)
-        print(jac_from_prod)
-        print(error_sources)
-
         assert qml.math.allclose(jac_from_prod[~error_sources], jac_from_fuse[~error_sources])
         nans = qml.math.isnan(jac_from_fuse[error_sources])
         nans = qml.math.reshape(nans, (len(nans), -1))
@@ -260,16 +256,23 @@ def test_jacobian_torch(self):
         import torch
 
         # Testing fewer points than with batching to limit test runtimes
-        special_points = np.array([1, 1 / 2, 0, -1]) * np.pi
+        special_points = np.array([1, 0, -1]) * np.pi
         special_angles = np.array(list(product(special_points, repeat=6))).reshape((-1, 2, 3))
-        random_angles = np.random.random((100, 2, 3))
+        random_angles = np.random.random((10, 2, 3))
         all_angles = np.concatenate([special_angles, random_angles])
 
         # Need holomorphic derivatives and complex inputs because the output matrices are complex
         all_angles = torch.tensor(all_angles, requires_grad=True)
-        jac_fn = lambda fn: (lambda *args: torch.autograd.functional.jacobian(fn, args))
+
+        def jacobian(fn):
+            real_fn = lambda arg: qml.math.real(fn(arg))
+            imag_fn = lambda arg: qml.math.imag(fn(arg))
+            real_jac_fn = lambda arg: torch.autograd.functional.jacobian(real_fn, (arg,))
+            imag_jac_fn = lambda arg: torch.autograd.functional.jacobian(imag_fn, (arg,))
+            return lambda arg: real_jac_fn(arg)[0] + 1j * imag_jac_fn(arg)[0]
+
         array_fn = lambda x: x.detach().numpy()
-        self.run_jacobian_test(all_angles, jac_fn, is_batched=False, array_fn=array_fn)
+        self.run_jacobian_test(all_angles, jacobian, is_batched=False, array_fn=array_fn)
 
     @pytest.mark.slow
     @pytest.mark.autograd

From cb2b20a91765052622bcc9aae14507c8b4bcc187 Mon Sep 17 00:00:00 2001
From: dwierichs <david.wierichs@xanadu.ai>
Date: Fri, 2 Aug 2024 11:40:42 +0200
Subject: [PATCH 29/30] lint

---
 .../transforms/test_optimization/test_optimization_utils.py  | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/tests/transforms/test_optimization/test_optimization_utils.py b/tests/transforms/test_optimization/test_optimization_utils.py
index 6596364aff0..ff31cd999c6 100644
--- a/tests/transforms/test_optimization/test_optimization_utils.py
+++ b/tests/transforms/test_optimization/test_optimization_utils.py
@@ -16,7 +16,6 @@
 """
 # pylint: disable=too-few-public-methods
 
-from functools import partial
 from itertools import product
 
 import pytest
@@ -173,6 +172,7 @@ def test_full_rot_fusion_special_angles(self):
         angles_1, angles_2 = np.transpose(special_angles, (1, 0, 2))
         self.run_interface_test(angles_1, angles_2)
 
+    # pylint: disable=too-many-arguments
     def run_jacobian_test(self, all_angles, jac_fn, is_batched, jit_fn=None, array_fn=None):
         """Execute standard test lines for testing Jacobians with different interfaces.
         #Note that the transpose calls only are relevant for tests with batching."""
@@ -326,9 +326,6 @@ def jac_fn(arg):
                     out = fn(arg)
                 return t.jacobian(out, arg)
 
-            def vmapped_jac_fn(arg):
-                return tf.vectorized_map(jac_fn, arg)
-
             return jac_fn
 
         # Need holomorphic derivatives and complex inputs because the output matrices are complex

From 1fc17c5c835a747eff8dade9750b02d703b69806 Mon Sep 17 00:00:00 2001
From: dwierichs <david.wierichs@xanadu.ai>
Date: Wed, 7 Aug 2024 12:32:22 +0200
Subject: [PATCH 30/30] code

---
 pennylane/transforms/optimization/merge_rotations.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/pennylane/transforms/optimization/merge_rotations.py b/pennylane/transforms/optimization/merge_rotations.py
index 4c488b1b2a8..8218e83e672 100644
--- a/pennylane/transforms/optimization/merge_rotations.py
+++ b/pennylane/transforms/optimization/merge_rotations.py
@@ -184,7 +184,7 @@ def stop_at(obj):
             next_gate = list_copy[next_gate_idx + 1]
 
             # If next gate is of the same type, we can merge the angles
-            if current_gate.name == next_gate.name and current_gate.wires == next_gate.wires:
+            if isinstance(current_gate, type(next_gate)) and current_gate.wires == next_gate.wires:
                 list_copy.pop(next_gate_idx + 1)
                 next_params = qml.math.stack(next_gate.parameters, like=interface)
                 # jax-jit does not support cast_like
@@ -192,7 +192,7 @@ def stop_at(obj):
                     next_params = qml.math.cast_like(next_params, cumulative_angles)
 
                 # The Rot gate must be treated separately
-                if current_gate.name == "Rot":
+                if isinstance(current_gate, qml.Rot):
                     cumulative_angles = fuse_rot_angles(cumulative_angles, next_params)
                 # Other, single-parameter rotation gates just have the angle summed
                 else: