Support for MultiRZ and CRot in default.qubit.tf (PennyLaneAI#921)

* Adding MultiRZ default.tensor.tf func * Add tests; docstrings; update * String compare * Temporarily remove CRot support so that it gets decomposed * Formatting * Fix * Integration tests for parametrized gates taking a tf variable * Update pennylane/devices/tf_ops.py Co-authored-by: Josh Izaac <josh146@gmail.com> * Changelog Co-authored-by: Josh Izaac <josh146@gmail.com>
alejomonbar · Dec 1, 2020 · 37f8e8f · 37f8e8f
1 parent 055f2ac
commit 37f8e8f
Show file tree

Hide file tree

Showing 5 changed files with 122 additions and 2 deletions.
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
@@ -304,6 +304,11 @@
 
 <h3>Bug fixes</h3>
 
+* The `default.qubit.tf` device is updated to handle TensorFlow objects (e.g.,
+  `tf.Variable`) as gate parameters correctly when using the `MultiRZ` and
+  `CRot` operations.
+  [(#921)](https://github.com/PennyLaneAI/pennylane/pull/921)
+
 * PennyLane tensor objects are now unwrapped in BaseQNode when passed as a
   keyword argument to the quantum function.
   [(#903)](https://github.com/PennyLaneAI/pennylane/pull/903)

diff --git a/pennylane/devices/default_qubit_tf.py b/pennylane/devices/default_qubit_tf.py
@@ -135,9 +135,11 @@ class DefaultQubitTF(DefaultQubit):
         "RY": tf_ops.RY,
         "RZ": tf_ops.RZ,
         "Rot": tf_ops.Rot,
+        "MultiRZ": tf_ops.MultiRZ,
         "CRX": tf_ops.CRX,
         "CRY": tf_ops.CRY,
         "CRZ": tf_ops.CRZ,
+        "CRot": tf_ops.CRot,
     }
 
     C_DTYPE = tf.complex128
@@ -199,6 +201,8 @@ def _get_unitary_matrix(self, unitary):
             object will be returned.
         """
         if unitary.name in self.parametric_ops:
+            if unitary.name == "MultiRZ":
+                return self.parametric_ops[unitary.name](unitary.parameters, len(unitary.wires))
             return self.parametric_ops[unitary.name](*unitary.parameters)
 
         if isinstance(unitary, DiagonalOperation):

diff --git a/pennylane/devices/tf_ops.py b/pennylane/devices/tf_ops.py
@@ -16,6 +16,7 @@
 """
 import tensorflow as tf
 from numpy import kron
+from pennylane.utils import pauli_eigs
 
 C_DTYPE = tf.complex128
 R_DTYPE = tf.float64
@@ -102,6 +103,21 @@ def Rot(a, b, c):
     return tf.linalg.diag(RZ(c)) @ RY(b) @ tf.linalg.diag(RZ(a))
 
 
+def MultiRZ(theta, n):
+    r"""Arbitrary multi Z rotation.
+
+    Args:
+        theta (float): rotation angle
+        n (int): number of wires the rotation acts on
+
+    Returns:
+        tf.Tensor[complex]: diagonal part of the MultiRZ matrix
+    """
+    theta = tf.cast(theta, dtype=C_DTYPE)
+    multi_Z_rot_eigs = tf.exp(-1j * theta / 2 * pauli_eigs(n))
+    return tf.convert_to_tensor(multi_Z_rot_eigs)
+
+
 def CRX(theta):
     r"""Two-qubit controlled rotation about the x axis.
 

diff --git a/tests/devices/test_default_qubit_tf.py b/tests/devices/test_default_qubit_tf.py
@@ -47,6 +47,8 @@
     CRoty,
     CRotz,
     CRot3,
+    MultiRZ1,
+    MultiRZ2,
 )
 
 np.random.seed(42)
@@ -72,9 +74,9 @@
 #####################################################
 
 single_qubit = [(qml.S, S), (qml.T, T), (qml.PauliX, X), (qml.PauliY, Y), (qml.PauliZ, Z), (qml.Hadamard, H)]
-single_qubit_param = [(qml.PhaseShift, Rphi), (qml.RX, Rotx), (qml.RY, Roty), (qml.RZ, Rotz)]
+single_qubit_param = [(qml.PhaseShift, Rphi), (qml.RX, Rotx), (qml.RY, Roty), (qml.RZ, Rotz), (qml.MultiRZ, MultiRZ1)]
 two_qubit = [(qml.CZ, CZ), (qml.CNOT, CNOT), (qml.SWAP, SWAP)]
-two_qubit_param = [(qml.CRX, CRotx), (qml.CRY, CRoty), (qml.CRZ, CRotz)]
+two_qubit_param = [(qml.CRX, CRotx), (qml.CRY, CRoty), (qml.CRZ, CRotz), (qml.MultiRZ, MultiRZ2)]
 three_qubit = [(qml.Toffoli, Toffoli), (qml.CSWAP, CSWAP)]
 
 
@@ -926,6 +928,69 @@ def circuit():
         expected = np.array([amplitude, 0, np.conj(amplitude), 0])
         assert np.allclose(state, expected, atol=tol, rtol=0)
 
+    @pytest.mark.parametrize("theta", [0.5432, -0.232])
+    @pytest.mark.parametrize("op,func", single_qubit_param)
+    def test_one_qubit_param_gates(self, theta, op, func, init_state, tol):
+        """Test the integration of the one-qubit single parameter rotations by passing
+        a TF data structure as a parameter"""
+        dev = qml.device("default.qubit.tf", wires=1)
+        state = init_state(1)
+
+        @qml.qnode(dev, interface='tf')
+        def circuit(params):
+            qml.QubitStateVector(state, wires=[0])
+            op(params[0], wires=[0])
+            return qml.expval(qml.PauliZ(0))
+
+        # Pass a TF Variable to the qfunc
+        params = tf.Variable(np.array([theta]))
+        circuit(params)
+        res = dev.state
+        expected = func(theta) @ state
+        assert np.allclose(res.numpy(), expected, atol=tol, rtol=0)
+
+    @pytest.mark.parametrize("theta", [0.5432, 4.213])
+    @pytest.mark.parametrize("op,func", two_qubit_param)
+    def test_two_qubit_param_gates(self, theta, op, func, init_state, tol):
+        """Test the integration of the two-qubit single parameter rotations by passing
+        a TF data structure as a parameter"""
+        dev = qml.device("default.qubit.tf", wires=2)
+        state = init_state(2)
+
+        @qml.qnode(dev, interface='tf')
+        def circuit(params):
+            qml.QubitStateVector(state, wires=[0,1])
+            op(params[0], wires=[0, 1])
+            return qml.expval(qml.PauliZ(0))
+
+        # Pass a TF Variable to the qfunc
+        params = tf.Variable(np.array([theta]))
+        circuit(params)
+        res = dev.state
+        expected = func(theta) @ state
+        assert np.allclose(res.numpy(), expected, atol=tol, rtol=0)
+
+    def test_controlled_rotation_integration(self, init_state, tol):
+        """Test the integration of the two-qubit controlled rotation by passing
+        a TF data structure as a parameter"""
+        dev = qml.device("default.qubit.tf", wires=2)
+        a = 1.7
+        b = 1.3432
+        c = -0.654
+        state = init_state(2)
+
+        @qml.qnode(dev, interface='tf')
+        def circuit(params):
+            qml.QubitStateVector(state, wires=[0,1])
+            qml.CRot(params[0], params[1], params[2], wires=[0,1])
+            return qml.expval(qml.PauliZ(0))
+
+        # Pass a TF Variable to the qfunc
+        params = tf.Variable(np.array([a,b,c]))
+        circuit(params)
+        res = dev.state
+        expected = CRot3(a, b, c) @ state
+        assert np.allclose(res.numpy(), expected, atol=tol, rtol=0)
 
 class TestPassthruIntegration:
     """Tests for integration with the PassthruQNode"""

diff --git a/tests/gate_data.py b/tests/gate_data.py
@@ -177,3 +177,33 @@ def CRot3(a, b, c):
             ],
         ]
     )
+
+def MultiRZ1(theta):
+    r"""Arbitrary multi Z rotation on one wire.
+
+    Args:
+        theta (float): rotation angle
+
+    Returns:
+        array: the one-wire MultiRZ matrix
+    """
+    return np.array([[np.exp(-1j * theta / 2), 0.0 + 0.0j], [0.0 + 0.0j, np.exp(1j * theta / 2)]])
+
+
+def MultiRZ2(theta):
+    r"""Arbitrary multi Z rotation on two wires.
+
+    Args:
+        theta (float): rotation angle
+
+    Returns:
+        array: the two-wire MultiRZ matrix
+    """
+    return np.array(
+        [
+            [np.exp(-1j * theta / 2), 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
+            [0.0 + 0.0j, np.exp(1j * theta / 2), 0.0 + 0.0j, 0.0 + 0.0j],
+            [0.0 + 0.0j, 0.0 + 0.0j, np.exp(1j * theta / 2), 0.0 + 0.0j],
+            [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, np.exp(-1j * theta / 2)],
+        ]
+    )