Repository: PennyLaneAI/pennylane
Pull request #1439 (2021-06-29 14:48:49)

Add dimension checking to QubitUnitary.
**Context:** Development of #1427 uncovered a bug where `QubitUnitary` was not checking whether the number of wires matched with the size of the unitary matrix. 

**Description of the Change:** Adds an `__init__` method to `QubitUnitary` to verify dimensionality and unitarity. Changes up the `_matrix` method accordingly. Also, converts `QubitUnitary` functions to use `qml.math`, and separates `QubitUnitary` tests into separate integration tests for all interfaces.

**Benefits:** Fixes the bug and clears up ambiguous behaviour such as 
```
qml.QubitUnitary(qml.Hadamard(wires=0).matrix, wires=[0, 1])
```
which would previously run without warnings.

**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/1439?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#1439](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1439?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (e971c8f) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/5e1b83232e93d070753f2f890525db1860d4044a?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (5e1b832) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1439/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/1439?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)

```diff
@@           Coverage Diff           @@
##           master    #1439   +/-   ##
=======================================
  Coverage   98.24%   98.24%           
=======================================
  Files         160      160           
  Lines       11991    11996    +5     
=======================================
+ Hits        11780    11785    +5     
  Misses        211      211           
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1439?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/ops/qubit.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1439/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9xdWJpdC5weQ==) | `98.73% <100.00%> (+<0.01%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1439?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/1439?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Last update [5e1b832...e971c8f](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1439?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: > Thanks @glassnotes! I left some minor comments in the review itself, but have a bigger request here:
> 
> Rather than having a unit test for `__init__` (which seems a bit redundant), perhaps it makes more sense to have an integration test, which does the following:
> 
>     * Tests the validation catches incorrect input and allows correct input using NumPy arrays and tensors
> 
>     * Tests that `ControlledQubitUnitary` continues to _avoid_ the validation?

@josh146 probably, then, it would be best to split out the `QubitUnitary` tests into their own test class like the excitation operations, since otherwise that will be quite a big chunk of test living in the `TestOperation` class. Furthermore, the extensive decomposition tests from #1427 could be included in the new class as well. 

Patch:
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
index 7c80c3c28..16e4cc751 100644
--- a/.github/CHANGELOG.md
+++ b/.github/CHANGELOG.md
@@ -108,6 +108,10 @@
   that performed post-processing on a vector-valued QNode.
   [(#)](https://github.com/PennyLaneAI/pennylane/pull/)
 
+* Fixed a bug in the initialization of `QubitUnitary` where the size of
+  the matrix was not checked against the number of wires.
+  [(#1439)](https://github.com/PennyLaneAI/pennylane/pull/1439)
+
 <h3>Documentation</h3>
 
 <h3>Contributors</h3>
diff --git a/pennylane/ops/qubit.py b/pennylane/ops/qubit.py
index 5002aa12f..e6303effd 100644
--- a/pennylane/ops/qubit.py
+++ b/pennylane/ops/qubit.py
@@ -17,6 +17,7 @@ quantum operations supported by PennyLane, as well as their conventions.
 """
 import cmath
 import functools
+import warnings
 
 # pylint:disable=abstract-method,arguments-differ,protected-access
 import math
@@ -2147,17 +2148,37 @@ class QubitUnitary(Operation):
     par_domain = "A"
     grad_method = None
 
-    @classmethod
-    def _matrix(cls, *params):
-        U = np.asarray(params[0])
+    def __init__(self, *params, wires, do_queue=True):
+        wires = Wires(wires)
 
-        if U.ndim != 2 or U.shape[0] != U.shape[1]:
-            raise ValueError("Operator must be a square matrix.")
+        # For pure QubitUnitary operations (not controlled), check that the number
+        # of wires fits the dimensions of the matrix
+        if not isinstance(self, ControlledQubitUnitary):
+            U = params[0]
 
-        if not np.allclose(U @ U.conj().T, np.identity(U.shape[0])):
-            raise ValueError("Operator must be unitary.")
+            dim = 2 ** len(wires)
 
-        return U
+            if U.shape != (dim, dim):
+                raise ValueError(
+                    f"Input unitary must be of shape {(dim, dim)} to act on {len(wires)} wires."
+                )
+
+            # Check for unitarity; due to variable precision across the different ML frameworks,
+            # here we issue a warning to check the operation, instead of raising an error outright.
+            if not qml.math.allclose(
+                qml.math.dot(U, qml.math.T(qml.math.conj(U))), qml.math.eye(qml.math.shape(U)[0])
+            ):
+                warnings.warn(
+                    f"Operator {U}\n may not be unitary."
+                    "Verify unitarity of operation, or use a datatype with increased precision.",
+                    UserWarning,
+                )
+
+        super().__init__(*params, wires=wires, do_queue=do_queue)
+
+    @classmethod
+    def _matrix(cls, *params):
+        return params[0]
 
     @staticmethod
     def decomposition(U, wires):


Test patch:
diff --git a/tests/ops/test_qubit_ops.py b/tests/ops/test_qubit_ops.py
index d63994607..047fd52f7 100644
--- a/tests/ops/test_qubit_ops.py
+++ b/tests/ops/test_qubit_ops.py
@@ -1530,73 +1530,6 @@ class TestOperations:
 
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
-    def test_qubit_unitary(self, tol):
-        """Test that the unitary operator produces the correct output."""
-        U = np.array([[1, 1], [1, -1]]) / np.sqrt(2)
-        out = qml.QubitUnitary(U, wires=0).matrix
-
-        # verify output type
-        assert isinstance(out, np.ndarray)
-
-        # verify equivalent to input state
-        assert np.allclose(out, U, atol=tol, rtol=0)
-
-    def test_qubit_unitary_exceptions(self):
-        """Tests that the unitary operator raises the proper errors."""
-        U = np.array([[1, 1], [1, -1]]) / np.sqrt(2)
-
-        # test non-square matrix
-        with pytest.raises(ValueError, match="must be a square matrix"):
-            qml.QubitUnitary(U[1:], wires=0).matrix
-
-        # test non-unitary matrix
-        U3 = U.copy()
-        U3[0, 0] += 0.5
-        with pytest.raises(ValueError, match="must be unitary"):
-            qml.QubitUnitary(U3, wires=0).matrix
-
-    @pytest.mark.parametrize(
-        "U", [np.array([0]), np.array([1, 0, 0, 1]), np.array([[[1, 0], [0, 1]]])]
-    )
-    def test_qubit_unitary_not_matrix_exception(self, U):
-        """Tests that the unitary operator raises the proper errors for arrays
-        that are not two-dimensional."""
-
-        # test non-square matrix
-        with pytest.raises(ValueError, match="must be a square matrix"):
-            qml.QubitUnitary(U, wires=0).matrix
-
-    @pytest.mark.parametrize(
-        "U,expected_gate,expected_params",
-        [  # First set of gates are diagonal and converted to RZ
-            (I, qml.RZ, [0]),
-            (Z, qml.RZ, [np.pi]),
-            (S, qml.RZ, [np.pi / 2]),
-            (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]),
-            # Next set of gates are non-diagonal and decomposed as Rots
-            (H, qml.Rot, [np.pi, np.pi / 2, 0]),
-            (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]),
-        ],
-    )
-    def test_qubit_unitary_decomposition(self, U, expected_gate, expected_params):
-        """Tests that single-qubit QubitUnitary decompositions are performed."""
-        decomp = qml.QubitUnitary.decomposition(U, wires=0)
-
-        assert len(decomp) == 1
-        assert isinstance(decomp[0], expected_gate)
-        assert np.allclose(decomp[0].parameters, expected_params)
-
-    def test_qubit_unitary_decomposition_multiqubit_invalid(self):
-        """Test that QubitUnitary is not decomposed for more than a single qubit."""
-        U = qml.CRZ(0.3, wires=[0, 1]).matrix
-
-        with pytest.raises(NotImplementedError, match="only supported for single-qubit"):
-            qml.QubitUnitary.decomposition(U, wires=[0, 1])
-
     def test_iswap_eigenval(self):
         """Tests that the ISWAP eigenvalue matches the numpy eigenvalues of the ISWAP matrix"""
         op = qml.ISWAP(wires=[0, 1])
@@ -2701,6 +2634,154 @@ class TestMultiRZ:
         spy.assert_not_called()
 
 
+class TestQubitUnitary:
+    """Tests for the QubitUnitary class."""
+
+    @pytest.mark.parametrize("U,num_wires", [(H, 1), (np.kron(H, H), 2)])
+    def test_qubit_unitary_autograd(self, U, num_wires):
+        """Test that the unitary operator produces the correct output and
+        catches incorrect input with autograd."""
+
+        out = qml.QubitUnitary(U, wires=range(num_wires)).matrix
+
+        # verify output type
+        assert isinstance(out, np.ndarray)
+
+        # verify equivalent to input state
+        assert qml.math.allclose(out, U)
+
+        # test non-square matrix
+        with pytest.raises(ValueError, match="must be of shape"):
+            qml.QubitUnitary(U[1:], wires=range(num_wires)).matrix
+
+        # test non-unitary matrix
+        U3 = U.copy()
+        U3[0, 0] += 0.5
+        with pytest.warns(UserWarning, match="may not be unitary"):
+            qml.QubitUnitary(U3, wires=range(num_wires)).matrix
+
+        # test an error is thrown when constructed with incorrect number of wires
+        with pytest.raises(ValueError, match="must be of shape"):
+            qml.QubitUnitary(U, wires=range(num_wires + 1)).matrix
+
+    @pytest.mark.parametrize("U,num_wires", [(H, 1), (np.kron(H, H), 2)])
+    def test_qubit_unitary_torch(self, U, num_wires):
+        """Test that the unitary operator produces the correct output and
+        catches incorrect input with torch."""
+        torch = pytest.importorskip("torch")
+
+        U = torch.tensor(U)
+        out = qml.QubitUnitary(U, wires=range(num_wires)).matrix
+
+        # verify output type
+        assert isinstance(out, torch.Tensor)
+
+        # verify equivalent to input state
+        assert qml.math.allclose(out, U)
+
+        # test non-square matrix
+        with pytest.raises(ValueError, match="must be of shape"):
+            qml.QubitUnitary(U[1:], wires=range(num_wires)).matrix
+
+        # test non-unitary matrix
+        U3 = U.detach().clone()
+        U3[0, 0] += 0.5
+        with pytest.warns(UserWarning, match="may not be unitary"):
+            qml.QubitUnitary(U3, wires=range(num_wires)).matrix
+
+        # test an error is thrown when constructed with incorrect number of wires
+        with pytest.raises(ValueError, match="must be of shape"):
+            qml.QubitUnitary(U, wires=range(num_wires + 1)).matrix
+
+    @pytest.mark.parametrize("U,num_wires", [(H, 1), (np.kron(H, H), 2)])
+    def test_qubit_unitary_tf(self, U, num_wires):
+        """Test that the unitary operator produces the correct output and
+        catches incorrect input with tensorflow."""
+        tf = pytest.importorskip("tensorflow")
+
+        U = tf.Variable(U)
+        out = qml.QubitUnitary(U, wires=range(num_wires)).matrix
+
+        # verify output type
+        assert isinstance(out, tf.Variable)
+
+        # verify equivalent to input state
+        assert qml.math.allclose(out, U)
+
+        # test non-square matrix
+        with pytest.raises(ValueError, match="must be of shape"):
+            qml.QubitUnitary(U[1:], wires=range(num_wires)).matrix
+
+        # test non-unitary matrix
+        U3 = tf.Variable(U + 0.5)
+        with pytest.warns(UserWarning, match="may not be unitary"):
+            qml.QubitUnitary(U3, wires=range(num_wires)).matrix
+
+        # test an error is thrown when constructed with incorrect number of wires
+        with pytest.raises(ValueError, match="must be of shape"):
+            qml.QubitUnitary(U, wires=range(num_wires + 1)).matrix
+
+    @pytest.mark.parametrize("U,num_wires", [(H, 1), (np.kron(H, H), 2)])
+    def test_qubit_unitary_jax(self, U, num_wires):
+        """Test that the unitary operator produces the correct output and
+        catches incorrect input with autograd."""
+        jax = pytest.importorskip("jax")
+        from jax import numpy as jnp
+
+        U = jnp.array(U)
+        out = qml.QubitUnitary(U, wires=range(num_wires)).matrix
+
+        # verify output type
+        assert isinstance(out, jnp.ndarray)
+
+        # verify equivalent to input state
+        assert qml.math.allclose(out, U)
+
+        # test non-square matrix
+        with pytest.raises(ValueError, match="must be of shape"):
+            qml.QubitUnitary(U[1:], wires=range(num_wires)).matrix
+
+        # test non-unitary matrix
+        U3 = U + 0.5
+        with pytest.warns(UserWarning, match="may not be unitary"):
+            qml.QubitUnitary(U3, wires=range(num_wires)).matrix
+
+        # test an error is thrown when constructed with incorrect number of wires
+        with pytest.raises(ValueError, match="must be of shape"):
+            qml.QubitUnitary(U, wires=range(num_wires + 1)).matrix
+
+    @pytest.mark.parametrize(
+        "U,expected_gate,expected_params",
+        [  # First set of gates are diagonal and converted to RZ
+            (I, qml.RZ, [0]),
+            (Z, qml.RZ, [np.pi]),
+            (S, qml.RZ, [np.pi / 2]),
+            (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]),
+            # Next set of gates are non-diagonal and decomposed as Rots
+            (H, qml.Rot, [np.pi, np.pi / 2, 0]),
+            (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]),
+        ],
+    )
+    def test_qubit_unitary_decomposition(self, U, expected_gate, expected_params):
+        """Tests that single-qubit QubitUnitary decompositions are performed."""
+        decomp = qml.QubitUnitary.decomposition(U, wires=0)
+
+        assert len(decomp) == 1
+        assert isinstance(decomp[0], expected_gate)
+        assert np.allclose(decomp[0].parameters, expected_params)
+
+    def test_qubit_unitary_decomposition_multiqubit_invalid(self):
+        """Test that QubitUnitary is not decomposed for more than a single qubit."""
+        U = qml.CRZ(0.3, wires=[0, 1]).matrix
+
+        with pytest.raises(NotImplementedError, match="only supported for single-qubit"):
+            qml.QubitUnitary.decomposition(U, wires=[0, 1])
+
+
 class TestDiagonalQubitUnitary:
     """Test the DiagonalQubitUnitary operation."""
 
diff --git a/tests/test_operation.py b/tests/test_operation.py
index 9d5ed3424..3eeb249f7 100644
--- a/tests/test_operation.py
+++ b/tests/test_operation.py
@@ -125,6 +125,9 @@ class TestOperation:
     def test_operation_init(self, test_class, monkeypatch):
         "Operation subclass initialization."
 
+        if test_class == qml.QubitUnitary:
+            pytest.skip("QubitUnitary can act on any number of wires.")
+
         if test_class in (qml.ControlledQubitUnitary, qml.MultiControlledX):
             pytest.skip("ControlledQubitUnitary alters the input params and wires in its __init__")