Repository: PennyLaneAI/pennylane
Pull request #1643 (2021-09-10 17:01:21)

Fix bug in single-qubit decomposition edge case.
**Context:** Development of #1552 uncovered an edge case in the `zyz_decomposition` function when the unitary being decomposed contains only off-diagonal elements.

**Description of the Change:** Fixes the bug, and adds new test cases. I also cleaned up a few parts of the tests so that fewer warnings are spit out by JAX. 

**Benefits:** Fixes bug

**Possible Drawbacks:** None

**Related GitHub Issues:** None


github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [.github/CHANGELOG.md](/PennyLaneAI/pennylane/blob/master/.github/CHANGELOG.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1643?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#1643](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1643?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (ed6d55f) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/4b82c832c5e55f9690196cf41a9bb611657cbd3f?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (4b82c83) will **decrease** coverage by `0.01%`.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1643/graphs/tree.svg?width=650&height=150&src=pr&token=9S3YzlUgGk&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1643?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)

```diff
@@            Coverage Diff             @@
##           master    #1643      +/-   ##
==========================================
- Coverage   99.14%   99.13%   -0.02%     
==========================================
  Files         196      196              
  Lines       14247    14254       +7     
==========================================
+ Hits        14125    14130       +5     
- Misses        122      124       +2     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1643?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [.../transforms/decompositions/single\_qubit\_unitary.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1643/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvZGVjb21wb3NpdGlvbnMvc2luZ2xlX3F1Yml0X3VuaXRhcnkucHk=) | `93.93% <100.00%> (-6.07%)` | :arrow_down: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1643?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1643?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Last update [4b82c83...ed6d55f](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1643?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

glassnotes: @josh146 I was trying to respond to your question about the negative sign, and found a way to streamline the implementation while looking into it. Once the CI runs do you mind taking a second look please?
glassnotes: [ch8964]

Patch:
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
index 8e3a3b2e1..334e79a98 100644
--- a/.github/CHANGELOG.md
+++ b/.github/CHANGELOG.md
@@ -424,6 +424,10 @@ and requirements-ci.txt (unpinned). This latter would be used by the CI.
 
 <h3>Bug fixes</h3>
 
+* Fix bug in edge case of single-qubit `zyz_decomposition` when only
+  off-diagonal elements are present.
+  [(#1643)](https://github.com/PennyLaneAI/pennylane/pull/1643)
+
 * `MottonenStatepreparation` can now be run with a single wire label not in a list.
   [(#1620)](https://github.com/PennyLaneAI/pennylane/pull/1620)
 
diff --git a/pennylane/transforms/decompositions/single_qubit_unitary.py b/pennylane/transforms/decompositions/single_qubit_unitary.py
index 1ce322119..c05c881f4 100644
--- a/pennylane/transforms/decompositions/single_qubit_unitary.py
+++ b/pennylane/transforms/decompositions/single_qubit_unitary.py
@@ -16,6 +16,7 @@ operations into elementary gates.
 """
 import pennylane as qml
 from pennylane import math
+import numpy as np
 
 
 def _convert_to_su2(U):
@@ -86,17 +87,18 @@ def zyz_decomposition(U, wire):
         omega = 2 * math.angle(U[1, 1])
         return [qml.RZ(omega, wires=wire)]
 
-    # If not diagonal, compute the angle of the RY
-    cos2_theta_over_2 = math.abs(U[0, 0] * U[1, 1])
-    theta = 2 * math.arccos(math.sqrt(cos2_theta_over_2))
-
-    # If the top left element is 0, can only use the off-diagonal elements We
+    # If the top left element is 0, can only use the off-diagonal elements. We
     # have to be very careful with the math here to ensure things that get
     # multiplied together are of the correct type in the different interfaces.
     if math.allclose(U[0, 0], [0.0]):
-        phi = 1j * math.log(U[0, 1] / U[1, 0])
-        omega = -phi - math.cast_like(2 * math.angle(U[1, 0]), phi)
+        phi = 0.0
+        theta = -np.pi
+        omega = 1j * math.log(U[0, 1] / U[1, 0]) - np.pi
     else:
+        # If not diagonal, compute the angle of the RY
+        cos2_theta_over_2 = math.abs(U[0, 0] * U[1, 1])
+        theta = 2 * math.arccos(math.sqrt(cos2_theta_over_2))
+
         el_division = U[0, 0] / U[1, 0]
         tan_part = math.cast_like(math.tan(theta / 2), el_division)
         omega = 1j * math.log(tan_part * el_division)


Test patch:
diff --git a/tests/ops/test_qubit_ops.py b/tests/ops/test_qubit_ops.py
index 5f2df5830..80ca573ab 100644
--- a/tests/ops/test_qubit_ops.py
+++ b/tests/ops/test_qubit_ops.py
@@ -2820,8 +2820,13 @@ class TestQubitUnitary:
             (qml.RZ(0.3, wires=0).matrix, qml.RZ, [0.3]),
             (qml.RZ(-0.5, wires=0).matrix, qml.RZ, [-0.5]),
             # Next set of gates are non-diagonal and decomposed as Rots
+            (
+                np.array([[0, -0.98310193 + 0.18305901j], [0.98310193 + 0.18305901j, 0]]),
+                qml.Rot,
+                [0, -np.pi, -5.914991017809059],
+            ),
             (H, qml.Rot, [np.pi, np.pi / 2, 0]),
-            (X, qml.Rot, [0.0, np.pi, np.pi]),
+            (X, qml.Rot, [0.0, -np.pi, -np.pi]),
             (qml.Rot(0.2, 0.5, -0.3, wires=0).matrix, qml.Rot, [0.2, 0.5, -0.3]),
             (np.exp(1j * 0.02) * qml.Rot(-1, 2, -3, wires=0).matrix, qml.Rot, [-1, 2, -3]),
         ],
diff --git a/tests/transforms/test_decompositions.py b/tests/transforms/test_decompositions.py
index 17d0ce188..2f839e87b 100644
--- a/tests/transforms/test_decompositions.py
+++ b/tests/transforms/test_decompositions.py
@@ -23,6 +23,7 @@ from pennylane import numpy as np
 from pennylane.wires import Wires
 from pennylane.transforms.decompositions import zyz_decomposition
 
+from test_optimization.utils import check_matrix_equivalence
 from gate_data import I, Z, S, T, H, X
 
 single_qubit_decomps = [
@@ -33,9 +34,15 @@ single_qubit_decomps = [
     (T, qml.RZ, [np.pi / 4]),
     (qml.RZ(0.3, wires=0).matrix, qml.RZ, [0.3]),
     (qml.RZ(-0.5, wires=0).matrix, qml.RZ, [-0.5]),
+    # This is an off-diagonal gate that gets converted to Rot, but one RZ angle is 0
+    (
+        np.array([[0, -0.98310193 + 0.18305901j], [0.98310193 + 0.18305901j, 0]]),
+        qml.Rot,
+        [0, -np.pi, -5.914991017809059],
+    ),
     # Next set of gates are non-diagonal and decomposed as Rots
     (H, qml.Rot, [np.pi, np.pi / 2, 0.0]),
-    (X, qml.Rot, [0.0, np.pi, np.pi]),
+    (X, qml.Rot, [0.0, -np.pi, -np.pi]),
     (qml.Rot(0.2, 0.5, -0.3, wires=0).matrix, qml.Rot, [0.2, 0.5, -0.3]),
     (np.exp(1j * 0.02) * qml.Rot(-1.0, 2.0, -3.0, wires=0).matrix, qml.Rot, [-1.0, 2.0, -3.0]),
 ]
@@ -55,53 +62,86 @@ class TestQubitUnitaryZYZDecomposition:
         obtained_gates = zyz_decomposition(U, Wires("a"))
 
         assert len(obtained_gates) == 1
-
         assert isinstance(obtained_gates[0], expected_gate)
         assert obtained_gates[0].wires == Wires("a")
         assert qml.math.allclose(obtained_gates[0].parameters, expected_params)
 
+        if obtained_gates[0].num_params == 1:
+            obtained_mat = qml.RZ(obtained_gates[0].parameters[0], wires=0).matrix
+        else:
+            obtained_mat = qml.Rot(*obtained_gates[0].parameters, wires=0).matrix
+
+        assert check_matrix_equivalence(obtained_mat, U)
+
     @pytest.mark.parametrize("U,expected_gate,expected_params", single_qubit_decomps)
     def test_zyz_decomposition_torch(self, U, expected_gate, expected_params):
         """Test that a one-qubit operation in Torch is correctly decomposed."""
         torch = pytest.importorskip("torch")
 
-        U = torch.tensor(U, dtype=torch.complex64)
+        U = torch.tensor(U, dtype=torch.complex128)
 
         obtained_gates = zyz_decomposition(U, wire="a")
 
         assert len(obtained_gates) == 1
         assert isinstance(obtained_gates[0], expected_gate)
         assert obtained_gates[0].wires == Wires("a")
-        assert qml.math.allclose(
-            [x.detach() for x in obtained_gates[0].parameters], expected_params
-        )
+        assert qml.math.allclose(qml.math.unwrap(obtained_gates[0].parameters), expected_params)
+
+        if obtained_gates[0].num_params == 1:
+            obtained_mat = qml.RZ(obtained_gates[0].parameters[0], wires=0).matrix
+        else:
+            obtained_mat = qml.Rot(*obtained_gates[0].parameters, wires=0).matrix
+
+        assert check_matrix_equivalence(obtained_mat, qml.math.unwrap(U))
 
     @pytest.mark.parametrize("U,expected_gate,expected_params", single_qubit_decomps)
     def test_zyz_decomposition_tf(self, U, expected_gate, expected_params):
         """Test that a one-qubit operation in Tensorflow is correctly decomposed."""
         tf = pytest.importorskip("tensorflow")
 
-        U = tf.Variable(U, dtype=tf.complex64)
+        U = tf.Variable(U, dtype=tf.complex128)
 
         obtained_gates = zyz_decomposition(U, wire="a")
 
         assert len(obtained_gates) == 1
         assert isinstance(obtained_gates[0], expected_gate)
         assert obtained_gates[0].wires == Wires("a")
-        assert qml.math.allclose([x.numpy() for x in obtained_gates[0].parameters], expected_params)
+        assert qml.math.allclose(qml.math.unwrap(obtained_gates[0].parameters), expected_params)
+
+        print(qml.math.unwrap(obtained_gates[0].parameters))
+        print(expected_params)
+
+        if obtained_gates[0].num_params == 1:
+            # With TF and RZ, need to cast since can't just unwrap
+            obtained_mat = qml.RZ(obtained_gates[0].parameters[0].numpy(), wires=0).matrix
+        else:
+            obtained_mat = qml.Rot(*qml.math.unwrap(obtained_gates[0].parameters), wires=0).matrix
+
+        assert check_matrix_equivalence(obtained_mat, U)
 
     @pytest.mark.parametrize("U,expected_gate,expected_params", single_qubit_decomps)
     def test_zyz_decomposition_jax(self, U, expected_gate, expected_params):
         """Test that a one-qubit operation in JAX is correctly decomposed."""
         jax = pytest.importorskip("jax")
 
-        U = jax.numpy.array(U, dtype=jax.numpy.complex64)
+        # Enable float64 support
+        from jax.config import config
+
+        remember = config.read("jax_enable_x64")
+        config.update("jax_enable_x64", True)
+
+        U = jax.numpy.array(U, dtype=jax.numpy.complex128)
 
         obtained_gates = zyz_decomposition(U, wire="a")
 
         assert len(obtained_gates) == 1
         assert isinstance(obtained_gates[0], expected_gate)
         assert obtained_gates[0].wires == Wires("a")
-        assert qml.math.allclose(
-            [jax.numpy.asarray(x) for x in obtained_gates[0].parameters], expected_params
-        )
+        assert qml.math.allclose(qml.math.unwrap(obtained_gates[0].parameters), expected_params)
+
+        if obtained_gates[0].num_params == 1:
+            obtained_mat = qml.RZ(obtained_gates[0].parameters[0], wires=0).matrix
+        else:
+            obtained_mat = qml.Rot(*obtained_gates[0].parameters, wires=0).matrix
+
+        assert check_matrix_equivalence(obtained_mat, U)
diff --git a/tests/transforms/test_optimization/utils.py b/tests/transforms/test_optimization/utils.py
index b6dc9f8cf..a49215c80 100644
--- a/tests/transforms/test_optimization/utils.py
+++ b/tests/transforms/test_optimization/utils.py
@@ -51,7 +51,7 @@ def check_matrix_equivalence(matrix_expected, matrix_obtained):
     """Takes two matrices and checks if multiplying one by the conjugate
     transpose of the other gives the identity."""
 
-    mat_product = qml.math.dot(qml.math.conj(matrix_obtained.T), matrix_expected)
+    mat_product = qml.math.dot(qml.math.conj(qml.math.T(matrix_obtained)), matrix_expected)
     mat_product /= mat_product[0, 0]
     return qml.math.allclose(mat_product, qml.math.eye(matrix_expected.shape[0]))
 
diff --git a/tests/transforms/test_unitary_to_rot.py b/tests/transforms/test_unitary_to_rot.py
index f20322e6a..72088004d 100644
--- a/tests/transforms/test_unitary_to_rot.py
+++ b/tests/transforms/test_unitary_to_rot.py
@@ -36,8 +36,13 @@ single_qubit_decomps = [
     (qml.RZ(0.3, wires=0).matrix, qml.RZ, [0.3]),
     (qml.RZ(-0.5, wires=0).matrix, qml.RZ, [-0.5]),
     # Next set of gates are non-diagonal and decomposed as Rots
+    (
+        np.array([[0, -0.98310193 + 0.18305901j], [0.98310193 + 0.18305901j, 0]]),
+        qml.Rot,
+        [0, -np.pi, -5.914991017809059],
+    ),
     (H, qml.Rot, [np.pi, np.pi / 2, 0.0]),
-    (X, qml.Rot, [0.0, np.pi, np.pi]),
+    (X, qml.Rot, [0.0, -np.pi, -np.pi]),
     (qml.Rot(0.2, 0.5, -0.3, wires=0).matrix, qml.Rot, [0.2, 0.5, -0.3]),
     (np.exp(1j * 0.02) * qml.Rot(-1.0, 2.0, -3.0, wires=0).matrix, qml.Rot, [-1.0, 2.0, -3.0]),
 ]
@@ -89,7 +94,7 @@ class TestDecomposeSingleQubitUnitaryTransform:
 
         assert isinstance(ops[1], expected_gate)
         assert ops[1].wires == Wires("a")
-        assert qml.math.allclose([x.detach() for x in ops[1].parameters], expected_params)
+        assert qml.math.allclose(qml.math.unwrap(ops[1].parameters), expected_params)
 
         assert isinstance(ops[2], qml.CNOT)
         assert ops[2].wires == Wires(["b", "a"])
@@ -112,7 +117,7 @@ class TestDecomposeSingleQubitUnitaryTransform:
 
         assert isinstance(ops[1], expected_gate)
         assert ops[1].wires == Wires("a")
-        assert qml.math.allclose([x.numpy() for x in ops[1].parameters], expected_params)
+        assert qml.math.allclose(qml.math.unwrap(ops[1].parameters), expected_params)
 
         assert isinstance(ops[2], qml.CNOT)
         assert ops[2].wires == Wires(["b", "a"])
@@ -122,6 +127,12 @@ class TestDecomposeSingleQubitUnitaryTransform:
         """Test that the transform works in the JAX interface."""
         jax = pytest.importorskip("jax")
 
+        # Enable float64 support
+        from jax.config import config
+
+        remember = config.read("jax_enable_x64")
+        config.update("jax_enable_x64", True)
+
         U = jax.numpy.array(U, dtype=jax.numpy.complex64)
 
         transformed_qfunc = unitary_to_rot(qfunc)
@@ -299,6 +310,12 @@ class TestQubitUnitaryDifferentiability:
         jax = pytest.importorskip("jax")
         from jax import numpy as jnp
 
+        # Enable float64 support
+        from jax.config import config
+
+        remember = config.read("jax_enable_x64")
+        config.update("jax_enable_x64", True)
+
         def qfunc_with_qubit_unitary(angles):
             z = angles[0]
             x = angles[1]