Repository: PennyLaneAI/pennylane

Issue:
[BUG] `single_qubit_fusion` returns wrong angles for specific case
### Expected behavior

The tape transform `single_qubit_fusion` always performs valid transforms, not changing the action of the tape.

### Actual behavior

For the following specific case, the transform reports wrong angles. Having tracked this down to `optimization_utils.fuse_rot_angles` already, I just report the minimal example here:
```pycon
>>> angles1 = [np.pi, np.pi/2, 0]
>>> angles2 = [0, -np.pi/2, 0]
>>> fuse_rot_angles(angles1, angles2)
array([0., 0., 0.])
```
As we can see, although only the `RY` rotation angles should cancel out, the `RZ` rotation caused by the first angle of `angles1` also disappeared.
As far as my experiments went, this bug seems to appear for a whole range of values, but I was not able to track down for which exactly. Tested examples include:
```pycon
#works
angles1 = [np.pi, np.pi/2, 0.]
angles2 = [0., -np.pi/2, 0.23]

#doesn't work
angles1 = [np.pi, np.pi/2, 0.]
angles2 = [0., -np.pi/2, 0.]

#works 
angles1 = [np.pi, np.pi/2, 0.]
angles2 = [1., -np.pi/2, 0.3]

#doesn't work at all, produces `nan` angle
angles1 = [0.9, np.pi/2, 0.]
angles2 = [0., -np.pi/2, 0.]

#works
angles1 = [0.9, np.pi/2, 1.]
angles2 = [0., -np.pi/2, 0.]
```

### Additional information

_No response_

### Source code

_No response_

### Tracebacks

_No response_

### System information

```shell
Name: PennyLane
Version: 0.26.0.dev0
Summary: PennyLane is a Python quantum machine learning library by Xanadu Inc.
Home-page: https://github.com/XanaduAI/pennylane
Author: 
Author-email: 
License: Apache License 2.0
Location: /home/david/venvs/xanadu3.10/lib/python3.10/site-packages/PennyLane-0.26.0.dev0-py3.10.egg
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, pennylane-lightning, retworkx, scipy, semantic-version, toml
Required-by: amazon-braket-pennylane-plugin, PennyLane-Cirq, PennyLane-Lightning, PennyLane-PQ, PennyLane-qiskit, pennylane-qulacs

Platform info:           Linux-5.4.0-125-generic-x86_64-with-glibc2.31
Python version:          3.10.6
Numpy version:           1.23.1
Scipy version:           1.9.0rc2
Installed devices:
- braket.aws.qubit (amazon-braket-pennylane-plugin-1.6.9)
- braket.local.qubit (amazon-braket-pennylane-plugin-1.6.9)
- lightning.qubit (PennyLane-Lightning-0.25.0)
- projectq.classical (PennyLane-PQ-0.17.0)
- projectq.ibm (PennyLane-PQ-0.17.0)
- projectq.simulator (PennyLane-PQ-0.17.0)
- qulacs.simulator (pennylane-qulacs-0.13.0.dev0)
- cirq.mixedsimulator (PennyLane-Cirq-0.25.0.dev0)
- cirq.pasqal (PennyLane-Cirq-0.25.0.dev0)
- cirq.qsim (PennyLane-Cirq-0.25.0.dev0)
- cirq.qsimh (PennyLane-Cirq-0.25.0.dev0)
- cirq.simulator (PennyLane-Cirq-0.25.0.dev0)
- qiskit.aer (PennyLane-qiskit-0.25.0.dev0)
- qiskit.basicaer (PennyLane-qiskit-0.25.0.dev0)
- qiskit.ibmq (PennyLane-qiskit-0.25.0.dev0)
- qiskit.ibmq.circuit_runner (PennyLane-qiskit-0.25.0.dev0)
- qiskit.ibmq.sampler (PennyLane-qiskit-0.25.0.dev0)
- default.gaussian (PennyLane-0.26.0.dev0)
- default.mixed (PennyLane-0.26.0.dev0)
- default.qubit (PennyLane-0.26.0.dev0)
- default.qubit.autograd (PennyLane-0.26.0.dev0)
- default.qubit.jax (PennyLane-0.26.0.dev0)
- default.qubit.tf (PennyLane-0.26.0.dev0)
- default.qubit.torch (PennyLane-0.26.0.dev0)
- default.qutrit (PennyLane-0.26.0.dev0)
```


### Existing GitHub issues

- [X] I have searched existing GitHub issues to make sure the issue does not already exist.


Discussion:

Resolving patch:
diff --git a/pennylane/transforms/optimization/optimization_utils.py b/pennylane/transforms/optimization/optimization_utils.py
--- 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 allclose, sin, cos, arccos, arctan2, stack, _multi_dispatch, is_abstract
+from functools import partial
+import numpy as np
+
+from pennylane.math import (
+    allclose,
+    sin,
+    cos,
+    arccos,
+    arctan2,
+    stack,
+    _multi_dispatch,
+    is_abstract,
+    abs as math_abs,
+)
 from pennylane.wires import Wires
 
 
@@ -60,17 +73,40 @@ def _quaternion_product(q1, q2):
     return stack([qw, qx, qy, qz])
 
 
-def _fuse(angles_1, angles_2):
-    """Perform fusion of two angle sets. Separated out so we can do JIT with conditionals."""
-    # Compute the product of the quaternions
-    qw, qx, qy, qz = _quaternion_product(_zyz_to_quat(angles_1), _zyz_to_quat(angles_2))
+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])
 
-    # Convert the product back into the angles fed to Rot
+
+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(qw**2 - qx**2 - qy**2 + qz**2)
+    y = arccos(y_arg)
 
     z2_arg1 = 2 * (qy * qz + qw * qx)
     z2_arg2 = 2 * (qw * qy - qx * qz)
@@ -79,6 +115,36 @@ def _fuse(angles_1, angles_2):
     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)
@@ -116,7 +182,7 @@ def fuse_rot_angles(angles_1, angles_2):
             return cond(
                 allclose(angles_1[1], 0.0) * allclose(angles_2[1], 0.0),
                 _no_fuse,
-                _fuse,
+                partial(_fuse, abstract_jax=True),
                 angles_1,
                 angles_2,
             )


Test patch:
diff --git a/tests/transforms/test_optimization/test_optimization_utils.py b/tests/transforms/test_optimization/test_optimization_utils.py
--- a/tests/transforms/test_optimization/test_optimization_utils.py
+++ b/tests/transforms/test_optimization/test_optimization_utils.py
@@ -104,9 +104,16 @@ def test_quaternion_product(self, angles_1, angles_2, expected_quat):
             ([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_autograd(self, angles_1, angles_2):