Repository: PennyLaneAI/pennylane (python)

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.

Pull request:
Fix `single_qubit_fusion` for singular case(s)
### Before submitting

Please complete the following checklist when submitting a PR:

- [x] All new features must include a unit test.
      If you've fixed a bug or added code that should be tested, add a test to the
      test directory!

- [x] All new functions and code must be clearly commented and documented.
      If you do make documentation changes, make sure that the docs build and
      render correctly by running `make docs`.

- [x] Ensure that the test suite passes, by running `make test`.

- [x] Add a new entry to the `doc/releases/changelog-dev.md` file, summarizing the
      change, and including a link back to the PR.

- [x] The PennyLane source code conforms to
      [PEP8 standards](https://www.python.org/dev/peps/pep-0008/).
      We check all of our code against [Pylint](https://www.pylint.org/).
      To lint modified files, simply `pip install pylint`, and then
      run `pylint pennylane/path/to/file.py`.

When all the above are checked, delete everything above the dashed
line and fill in the pull request template.

------------------------------------------------------------------------------------------------------------

**Context:**
For special values of rotation angles, `fuse_rot_angles` in the transforms' optimization_utils does not produce the correct
merged rotation angles. This seems to happen at singular points of `arctan2`.

**Description of the Change:**
Includes a logical branch for two special cases of quaternion values/rotation angles that respect the special values of `arctan2` attained for the special cases.
Adds tests that previously failed.

**Benefits:**
Correct merging of rotations in more scenarios.

**Possible Drawbacks:**
A bit more code.

**Related GitHub Issues:**
fixes #3023 


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index b9b27cf98..79dd35847 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -304,6 +304,10 @@
 
 <h3>Bug fixes</h3>
 
+* Fixes a bug where the tape transform `single_qubit_fusion` computed wrong rotation angles
+  for specific combinations of rotations.
+  [(#3024)](https://github.com/PennyLaneAI/pennylane/pull/3024)
+    
 * Jax gradients now work with a QNode when the quantum function was transformed by `qml.simplify`.
   [(#3017)](https://github.com/PennyLaneAI/pennylane/pull/3017)
 
diff --git a/pennylane/transforms/optimization/optimization_utils.py b/pennylane/transforms/optimization/optimization_utils.py
index 32486cd07..7bdf0a206 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 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
index 8951ab1bb..799257712 100644
--- a/tests/transforms/test_optimization/test_optimization_utils.py
+++ b/tests/transforms/test_optimization/test_optimization_utils.py
@@ -104,9 +104,16 @@ class TestRotGateFusion:
             ([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):