Repository: PennyLaneAI/pennylane

Issue:
[BUG] various problems with QubitUnitary
### Expected behavior

QubitUnitary should be correctly executed for all unitary input matrices, should be (at least) finite difference differentiable, and parameter-shift differentiable via decomposition.

### Actual behavior

The minimal non working example code reveals several problems:

- For the given `in_mat` input matrix the matrix returned by `qml.matrix()` on a QNode that contains a single `QubitUnitary` operation is wrong and it is off by more than just a global phase. (For other input matrices such as, i.e., `in_mat = make_unitary(np.pi/2 + 1e-6)` the matrices match)
- Also executing the QNode on `default.qubit` produces the wrong state, even though the operations on the tape look correct and when drawing the tape all looks good.
- Finally, computing gradients also does not work. Differentiating at `np.pi/2` yields a `NotImplementedError` in `autograd` (autograd==1.6.2). Differentiating at random parameters yields `nan`, even with `diff_method="finite-diff"`.

I provide the output of the minimal non-working code and its output below.

### Additional information

_No response_

### Source code

```shell
import pennylane as qml
from pennylane import numpy as np

def make_unitary(theta1):
    generator = theta1 / 2 * (
        np.cos(0.2)
        / 2
        * (np.array([[0, 0, 0, 0], [0, 0, 2, 0], [0, 2, 0, 0], [0, 0, 0, 0]]))
        + np.sin(0.2)
        / 2
        * (np.array([[0, 0, 0, 0], [0, 0, 2j, 0], [0, -2j, 0, 0], [0, 0, 0, 0]]))
    )

    def expm(X):
        d, U = np.linalg.eigh(-1.0j * X)
        V = np.dot(U, np.dot(np.diag(np.exp(1.0j * d)), np.conj(U).T))
        return V

    mat = expm(-1j * generator)

    assert np.allclose(np.dot(np.transpose(np.conj(mat)), mat), np.eye(len(mat))), "mat is not unitary"

    return mat

in_mat = make_unitary(np.pi/2)
print("in_mat")
print(in_mat.round(3))

dev0 = qml.device("default.qubit", 2)
@qml.qnode(dev0, diff_method="finite-diff")
def foo():
    qml.QubitUnitary(in_mat, wires=[0, 1])
    return qml.state()

out_mat = qml.matrix(foo, wire_order=[0, 1])()
print("out_mat")
print(out_mat.round(3))
global_phase = np.angle(out_mat[0, 0] / in_mat[0, 0])
print("global_phase", global_phase)
out_mat *= np.exp(-1j * global_phase)
print("out_mat with phase corrected")
print(out_mat.round(3))
print("difference should be zero but it is not:")
print(np.round(in_mat - out_mat, 3))
print("circuit:")
print(qml.draw(foo, wire_order=[0, 1], expansion_strategy="device")())
print("tape looks correct:")
print(foo.tape.operations)
print("state returned by qnode should be |0> but it is not:")
print(foo())

@qml.qnode(dev0, diff_method="finite-diff", expansion_strategy="device")
def bar(theta):
    qml.QubitUnitary(make_unitary(theta), wires=[0, 1])
    return qml.expval(qml.PauliZ(0))

try:
    print("jac1", qml.jacobian(bar)(np.array(np.pi/2, requires_grad=True)))
except Exception as exception:
    jac2 = qml.jacobian(bar)(1e-3 * np.array(np.random.randn(1), requires_grad=True))
    print("jac2", jac2)
    if np.isnan(jac2):
        raise Exception("jac2 should not be nan") from exception
```


### Tracebacks

```shell
in_mat
[[ 1.   +0.j     0.   +0.j     0.   +0.j     0.   +0.j   ]
 [ 0.   +0.j     0.707-0.j     0.14 -0.693j  0.   +0.j   ]
 [ 0.   +0.j    -0.14 -0.693j  0.707+0.j     0.   +0.j   ]
 [ 0.   +0.j     0.   +0.j     0.   +0.j     1.   +0.j   ]]
out_mat
[[ 0.056+0.135j  0.752-0.404j  0.339-0.102j  0.135+0.327j]
 [-0.246-0.817j -0.056-0.135j  0.347-0.068j -0.339+0.102j]
 [-0.312+0.167j  0.293-0.197j -0.056-0.135j -0.404-0.752j]
 [ 0.327-0.135j  0.167+0.312j -0.246-0.817j -0.135+0.056j]]
global_phase 1.1780971988894835
out_mat with phase corrected
[[ 0.146+0.j    -0.085-0.849j  0.035-0.352j  0.354+0.j   ]
 [-0.849-0.085j -0.146-0.j     0.07 -0.347j -0.035+0.352j]
 [ 0.035+0.352j -0.07 -0.347j -0.146-0.j    -0.849+0.085j]
 [-0.   -0.354j  0.352-0.035j -0.849-0.085j -0.   +0.146j]]
difference should be zero but it is not:
[[ 0.854-0.j     0.085+0.849j -0.035+0.352j -0.354-0.j   ]
 [ 0.849+0.085j  0.854+0.j     0.07 -0.347j  0.035-0.352j]
 [-0.035-0.352j -0.07 -0.347j  0.854+0.j     0.849-0.085j]
 [ 0.   +0.354j -0.352+0.035j  0.849+0.085j  1.   -0.146j]]
circuit:
0: ─╭U(M0)─┤ ╭State
1: ─╰U(M0)─┤ ╰State

M0 = 
[[ 1.        +0.00000000e+00j  0.        +0.00000000e+00j
   0.        +0.00000000e+00j  0.        +0.00000000e+00j]
 [ 0.        +0.00000000e+00j  0.70710678-2.29934717e-17j
   0.14048043-6.93011723e-01j  0.        +0.00000000e+00j]
 [ 0.        +0.00000000e+00j -0.14048043-6.93011723e-01j
   0.70710678+0.00000000e+00j  0.        +0.00000000e+00j]
 [ 0.        +0.00000000e+00j  0.        +0.00000000e+00j
   0.        +0.00000000e+00j  1.        +0.00000000e+00j]]
tape looks correct:
[QubitUnitary(tensor([[ 1.        +0.00000000e+00j,  0.        +0.00000000e+00j,
          0.        +0.00000000e+00j,  0.        +0.00000000e+00j],
        [ 0.        +0.00000000e+00j,  0.70710678-2.29934717e-17j,
          0.14048043-6.93011723e-01j,  0.        +0.00000000e+00j],
        [ 0.        +0.00000000e+00j, -0.14048043-6.93011723e-01j,
          0.70710678+0.00000000e+00j,  0.        +0.00000000e+00j],
        [ 0.        +0.00000000e+00j,  0.        +0.00000000e+00j,
          0.        +0.00000000e+00j,  1.        +0.00000000e+00j]], requires_grad=True), wires=[0, 1])]
state returned by qnode should be |0> but it is not:
[ 0.0560427 +0.13529903j -0.24628222-0.81725055j -0.31150153+0.16723276j
  0.32664072-0.13529903j]
/home/cvjjm/miniforge3/envs/hfak/lib/python3.9/site-packages/autograd/numpy/numpy_wrapper.py:156: ComplexWarning: Casting complex values to real discards the imaginary part
  return A.astype(dtype, order, casting, subok, copy)
/home/cvjjm/src/covqcstack/hfak/pennylane/pennylane/numpy/tensor.py:155: RuntimeWarning: invalid value encountered in power
  res = super().__array_ufunc__(ufunc, method, *args, **kwargs)
/home/cvjjm/miniforge3/envs/hfak/lib/python3.9/site-packages/autograd/numpy/linalg.py:180: RuntimeWarning: divide by zero encountered in divide
  F = off_diag / (T(w_repeated) - w_repeated + anp.eye(N))
jac2 [nan]
Traceback (most recent call last):
  File "/home/cvjjm/miniforge3/envs/hfak/lib/python3.9/site-packages/autograd/core.py", line 31, in __init__
    vjpmaker = primitive_vjps[fun]
KeyError: <function primitive.<locals>.f_wrapped at 0x7fbc8d724040>

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/cvjjm/src/covqcstack/hfak/covvqetools/qubit_unitary_test.py", line 59, in <module>
    print("jac1", qml.jacobian(bar)(np.array(np.pi/2, requires_grad=True)))
  File "/home/cvjjm/src/covqcstack/hfak/pennylane/pennylane/_grad.py", line 455, in _jacobian_function
    jac = tuple(_jacobian(func, arg)(*args, **kwargs) for arg in _argnum)
  File "/home/cvjjm/src/covqcstack/hfak/pennylane/pennylane/_grad.py", line 455, in <genexpr>
    jac = tuple(_jacobian(func, arg)(*args, **kwargs) for arg in _argnum)
  File "/home/cvjjm/miniforge3/envs/hfak/lib/python3.9/site-packages/autograd/wrap_util.py", line 20, in nary_f
    return unary_operator(unary_f, x, *nary_op_args, **nary_op_kwargs)
  File "/home/cvjjm/miniforge3/envs/hfak/lib/python3.9/site-packages/autograd/differential_operators.py", line 60, in jacobian
    vjp, ans = _make_vjp(fun, x)
  File "/home/cvjjm/miniforge3/envs/hfak/lib/python3.9/site-packages/autograd/core.py", line 10, in make_vjp
    end_value, end_node =  trace(start_node, fun, x)
  File "/home/cvjjm/miniforge3/envs/hfak/lib/python3.9/site-packages/autograd/tracer.py", line 10, in trace
    end_box = fun(start_box)
  File "/home/cvjjm/miniforge3/envs/hfak/lib/python3.9/site-packages/autograd/wrap_util.py", line 15, in unary_f
    return fun(*subargs, **kwargs)
  File "/home/cvjjm/src/covqcstack/hfak/pennylane/pennylane/qnode.py", line 1039, in __call__
    res = qml.execute(
  File "/home/cvjjm/src/covqcstack/hfak/pennylane/pennylane/interfaces/execution.py", line 632, in execute
    tapes, post_processing = transform_program(tapes)
  File "/home/cvjjm/src/covqcstack/hfak/pennylane/pennylane/transforms/core/transform_program.py", line 435, in __call__
    new_tapes, fn = transform(tape, *targs, **tkwargs)
  File "/home/cvjjm/src/covqcstack/hfak/pennylane/pennylane/gradients/finite_difference.py", line 191, in _expand_transform_finite_diff
    expanded_tape = expand_invalid_trainable(tape)
  File "/home/cvjjm/src/covqcstack/hfak/pennylane/pennylane/transforms/tape_expand.py", line 102, in expand_fn
    tape = tape.expand(depth=depth, stop_at=stop_at)
  File "/home/cvjjm/src/covqcstack/hfak/pennylane/pennylane/tape/qscript.py", line 946, in expand
    new_script = qml.tape.tape.expand_tape(
  File "/home/cvjjm/src/covqcstack/hfak/pennylane/pennylane/tape/tape.py", line 203, in expand_tape
    obj = QuantumScript(obj.decomposition(), _update=False)
  File "/home/cvjjm/src/covqcstack/hfak/pennylane/pennylane/operation.py", line 1261, in decomposition
    return self.compute_decomposition(
  File "/home/cvjjm/src/covqcstack/hfak/pennylane/pennylane/ops/qubit/matrix_ops.py", line 222, in compute_decomposition
    return qml.ops.two_qubit_decomposition(U, Wires(wires))
  File "/home/cvjjm/src/covqcstack/hfak/pennylane/pennylane/ops/op_math/decompositions/two_qubit_unitary.py", line 611, in two_qubit_decomposition
    num_cnots = _compute_num_cnots(U)
  File "/home/cvjjm/src/covqcstack/hfak/pennylane/pennylane/ops/op_math/decompositions/two_qubit_unitary.py", line 126, in _compute_num_cnots
    evs = math.linalg.eigvals(gammaU)
  File "/home/cvjjm/miniforge3/envs/hfak/lib/python3.9/site-packages/autoray/autoray.py", line 79, in do
    return get_lib_fn(backend, fn)(*args, **kwargs)
  File "/home/cvjjm/src/covqcstack/hfak/pennylane/pennylane/numpy/wrapper.py", line 117, in _wrapped
    res = obj(*args, **kwargs)
  File "/home/cvjjm/miniforge3/envs/hfak/lib/python3.9/site-packages/autograd/tracer.py", line 45, in f_wrapped
    node = node_constructor(ans, f_wrapped, argvals, kwargs, argnums, parents)
  File "/home/cvjjm/miniforge3/envs/hfak/lib/python3.9/site-packages/autograd/core.py", line 34, in __init__
    raise NotImplementedError("VJP of {} wrt argnums {} not defined"
NotImplementedError: VJP of eigvals wrt argnums (0,) not defined

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/cvjjm/src/covqcstack/hfak/covvqetools/qubit_unitary_test.py", line 64, in <module>
    raise Exception("jac2 should not be nan") from exception
Exception: jac2 should not be nan

Compilation exited abnormally with code 1 at Tue Mar  5 11:51:09
```


### System information

```shell
Name: PennyLane
Version: 0.34.0
Summary: PennyLane is a Python quantum machine learning library by Xanadu Inc.
Home-page: https://github.com/PennyLaneAI/pennylane
Author:
Author-email:
License: Apache License 2.0
Location: /home/cvjjm/miniforge3/envs/hfak/lib/python3.9/site-packages
Editable project location: /home/cvjjm/src/covqcstack/hfak/pennylane
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, pennylane-lightning, requests, rustworkx, scipy, semantic-version, toml, typing-extensions
Required-by: covvqetools, PennyLane-Lightning

Platform info:           Linux-5.15.133.1-microsoft-standard-WSL2-x86_64-with-glibc2.31
Python version:          3.9.16
Numpy version:           1.23.0
Scipy version:           1.10.0
Installed devices:
- cov.caching.qubit (covvqetools-0.7.0)
- cov.caching.qulacs (covvqetools-0.7.0)
- cov.parsec.casbox (covvqetools-0.7.0)
- cov.parsec.qubit (covvqetools-0.7.0)
- cov.qubit (covvqetools-0.7.0)
- cov.quicksilver.casbox (covvqetools-0.7.0)
- cov.quicksilver.qubit (covvqetools-0.7.0)
- default.gaussian (PennyLane-0.34.0)
- default.mixed (PennyLane-0.34.0)
- default.qubit (PennyLane-0.34.0)
- default.qubit.autograd (PennyLane-0.34.0)
- default.qubit.jax (PennyLane-0.34.0)
- default.qubit.legacy (PennyLane-0.34.0)
- default.qubit.tf (PennyLane-0.34.0)
- default.qubit.torch (PennyLane-0.34.0)
- default.qutrit (PennyLane-0.34.0)
- null.qubit (PennyLane-0.34.0)
- lightning.qubit (PennyLane-Lightning-0.34.0)
```


### Existing GitHub issues

- [X] I have searched existing GitHub issues to make sure the issue does not already exist.

Discussion:
albi3ro: > Is it my impression, or are tests stuck?

Sometimes they won't run if the diff is out of date.

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index b4456e597..4b450f416 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -352,6 +352,9 @@
 
 <h3>Bug fixes 🐛</h3>
 
+* `two_qubit_decomposition` no longer diverges at a special case of unitary matrix.
+  [(#5448)](https://github.com/PennyLaneAI/pennylane/pull/5448)
+
 * The `qml.QNSPSAOptimizer` now correctly handles optimization for legacy devices that do not follow the new API design.
   [(#5497)](https://github.com/PennyLaneAI/pennylane/pull/5497)
 
diff --git a/pennylane/math/single_dispatch.py b/pennylane/math/single_dispatch.py
index 1999ce69e..f1ed70589 100644
--- a/pennylane/math/single_dispatch.py
+++ b/pennylane/math/single_dispatch.py
@@ -265,6 +265,13 @@ ar.autoray._FUNC_ALIASES["tensorflow", "arctan2"] = "atan2"
 ar.autoray._FUNC_ALIASES["tensorflow", "diag"] = "diag"
 
 
+ar.register_function(
+    "tensorflow",
+    "finfo",
+    lambda *args, **kwargs: _i("tf").experimental.numpy.finfo(*args, **kwargs),
+)
+
+
 def _tf_convert_to_tensor(x, **kwargs):
     if isinstance(x, _i("tf").Tensor) and "dtype" in kwargs:
         return _i("tf").cast(x, **kwargs)
diff --git a/pennylane/ops/op_math/decompositions/two_qubit_unitary.py b/pennylane/ops/op_math/decompositions/two_qubit_unitary.py
index 3f72ff49b..9baf4b7da 100644
--- a/pennylane/ops/op_math/decompositions/two_qubit_unitary.py
+++ b/pennylane/ops/op_math/decompositions/two_qubit_unitary.py
@@ -347,6 +347,7 @@ def _decomposition_2_cnots(U, wires):
     part has the same spectrum as U, and then we can recover A, B, C, D.
     """
     # Compute the rotation angles
+
     u = math.dot(Edag, math.dot(U, E))
     gammaU = math.dot(u, math.T(u))
     evs, _ = math.linalg.eig(gammaU)
@@ -383,6 +384,7 @@ def _decomposition_2_cnots(U, wires):
     else:
         # For the non-special case, the eigenvalues come in conjugate pairs.
         # We need to find two non-conjugate eigenvalues to extract the angles.
+
         x = math.angle(evs[0])
         y = math.angle(evs[1])
 
@@ -392,7 +394,6 @@ def _decomposition_2_cnots(U, wires):
 
         delta = (x + y) / 2
         phi = (x - y) / 2
-
         interior_decomp = [
             qml.CNOT(wires=[wires[1], wires[0]]),
             qml.RZ(delta, wires=wires[0]),
@@ -400,7 +401,10 @@ def _decomposition_2_cnots(U, wires):
             qml.CNOT(wires=[wires[1], wires[0]]),
         ]
 
-        RZd = qml.RZ(math.cast_like(delta, 1j), wires=0).matrix()
+        # need to perturb x by 5 precision to avoid a discontinuity at a special case.
+        # see https://github.com/PennyLaneAI/pennylane/issues/5308
+        precision = qml.math.finfo(delta.dtype).eps
+        RZd = qml.RZ(math.cast_like(delta + 5 * precision, 1j), wires=0).matrix()
         RXp = qml.RX(phi, wires=0).matrix()
         inner_matrix = math.kron(RZd, RXp)
 


Test patch:
diff --git a/tests/ops/op_math/test_decompositions.py b/tests/ops/op_math/test_decompositions.py
index 83a27d508..6f5e7e8bd 100644
--- a/tests/ops/op_math/test_decompositions.py
+++ b/tests/ops/op_math/test_decompositions.py
@@ -1230,3 +1230,37 @@ class TestTwoQubitUnitaryDecompositionInterfaces:
         jitted_matrix = jax.jit(wrapped_decomposition)(U)
 
         assert check_matrix_equivalence(U, jitted_matrix, atol=1e-7)
+
+
+def test_two_qubit_decomposition_special_case_discontinuity():
+    """Test that two_qubit_decomposition still provides accurate numbers at a special case."""
+
+    def make_unitary(theta1):
+        generator = (
+            theta1
+            / 2
+            * (
+                np.cos(0.2)
+                / 2
+                * (np.array([[0, 0, 0, 0], [0, 0, 2, 0], [0, 2, 0, 0], [0, 0, 0, 0]]))
+                + np.sin(0.2)
+                / 2
+                * (np.array([[0, 0, 0, 0], [0, 0, 2j, 0], [0, -2j, 0, 0], [0, 0, 0, 0]]))
+            )
+        )
+
+        def expm(val):
+            d, U = np.linalg.eigh(-1.0j * val)
+            return np.dot(U, np.dot(np.diag(np.exp(1.0j * d)), np.conj(U).T))
+
+        mat = expm(-1j * generator)
+
+        assert np.allclose(
+            np.dot(np.transpose(np.conj(mat)), mat), np.eye(len(mat))
+        ), "mat is not unitary"
+
+        return mat
+
+    mat = make_unitary(np.pi / 2)
+    decomp_mat = qml.matrix(two_qubit_decomposition, wire_order=(0, 1))(mat, wires=(0, 1))
+    assert qml.math.allclose(mat, decomp_mat)