Merge pull request #633 from tensorflow/line_deser

Add cirq.LineQubit support to op_deserializer.py.

Author: jaeyoo

tensorflow_quantum/core/serialize/op_deserializer.py

@@ -19,6 +19,7 @@
 import cirq
 
 GRID_QUBIT_ID_PATTERN = r'^q?(-?\d+)_(-?\d+)$'
+LINE_QUBIT_ID_PATTERN = r'^q?(-?\d+)$'
 SUPPORTED_FUNCTIONS_FOR_LANGUAGE = {
     '': frozenset(),
     'linear': frozenset({'add', 'mul'}),
@@ -42,11 +43,17 @@ def qubit_from_proto(proto_id):
     """
 
     match = re.match(GRID_QUBIT_ID_PATTERN, proto_id)
-    if match is None:
-        raise ValueError(
-            f'Expected GridQubit proto w/ form [q]<int>_<int>, got {proto_id}')
-    row, col = match.groups()
-    return cirq.GridQubit(row=int(row), col=int(col))
+    if match is not None:
+        row, col = match.groups()
+        return cirq.GridQubit(row=int(row), col=int(col))
+
+    match = re.match(LINE_QUBIT_ID_PATTERN, proto_id)
+    if match is not None:
+        x, = match.groups()
+        return cirq.LineQubit(int(x))
+
+    raise ValueError('Expected GridQubit proto w/ form [q]<int>_<int>,'
+                     f' or LineQubit w/ form [q]<int> got {proto_id}')
 
 
 def _arg_from_proto(

tensorflow_quantum/core/serialize/op_deserializer_test.py

@@ -96,8 +96,12 @@ def _value_equality_values_(self):
 class OpDeserializerTest(tf.test.TestCase, parameterized.TestCase):
     """Test OpDeserializer functionality."""
 
-    @parameterized.parameters(TEST_CASES)
-    def test_from_proto(self, val_type, val, arg_value):
+    @parameterized.parameters([
+        CASE + (x,)
+        for CASE in TEST_CASES
+        for x in [cirq.GridQubit(1, 2), cirq.LineQubit(4)]
+    ])
+    def test_from_proto(self, val_type, val, arg_value, q):
         """Test from proto under many cases."""
         deserializer = op_deserializer.GateOpDeserializer(
             serialized_gate_id='my_gate',
@@ -116,10 +120,9 @@ def test_from_proto(self, val_type, val, arg_value):
                 'my_val': arg_value
             },
             'qubits': [{
-                'id': '1_2'
+                'id': '1_2' if isinstance(q, cirq.GridQubit) else '4'
             }]
         })
-        q = cirq.GridQubit(1, 2)
         result = deserializer.from_proto(serialized,
                                          arg_function_language='linear')
         self.assertEqual(result, GateWithAttribute(val)(q))
@@ -260,7 +263,8 @@ def test_from_proto_function_argument_not_set(self):
             _ = deserializer.from_proto(serialized,
                                         arg_function_language='linear')
 
-    def test_from_proto_value_func(self):
+    @parameterized.parameters([cirq.GridQubit(1, 2), cirq.LineQubit(4)])
+    def test_from_proto_value_func(self, q):
         """Test value func deserialization in simple case."""
         deserializer = op_deserializer.GateOpDeserializer(
             serialized_gate_id='my_gate',
@@ -282,14 +286,14 @@ def test_from_proto_value_func(self):
                 }
             },
             'qubits': [{
-                'id': '1_2'
+                'id': '1_2' if isinstance(q, cirq.GridQubit) else '4'
             }]
         })
-        q = cirq.GridQubit(1, 2)
         result = deserializer.from_proto(serialized)
         self.assertEqual(result, GateWithAttribute(1.125)(q))
 
-    def test_from_proto_not_required_ok(self):
+    @parameterized.parameters([cirq.GridQubit(1, 2), cirq.LineQubit(4)])
+    def test_from_proto_not_required_ok(self, q):
         """Deserialization succeeds for missing not required fields."""
         deserializer = op_deserializer.GateOpDeserializer(
             serialized_gate_id='my_gate',
@@ -315,13 +319,13 @@ def test_from_proto_not_required_ok(self):
                 }
             },
             'qubits': [{
-                'id': '1_2'
+                'id': '1_2' if isinstance(q, cirq.GridQubit) else '4'
             }]
         })
-        q = cirq.GridQubit(1, 2)
         result = deserializer.from_proto(serialized)
         self.assertEqual(result, GateWithAttribute(0.125)(q))
 
+
     def test_from_proto_missing_required_arg(self):
         """Error raised when required field is missing."""
         deserializer = op_deserializer.GateOpDeserializer(
@@ -382,7 +386,8 @@ def test_from_proto_required_arg_not_assigned(self):
         with self.assertRaises(ValueError):
             deserializer.from_proto(serialized)
 
-    def test_defaults(self):
+    @parameterized.parameters([cirq.GridQubit(1, 2), cirq.LineQubit(4)])
+    def test_defaults(self, q):
         """Ensure default values still deserialize."""
         deserializer = op_deserializer.GateOpDeserializer(
             serialized_gate_id='my_gate',
@@ -402,13 +407,12 @@ def test_defaults(self):
             },
             'args': {},
             'qubits': [{
-                'id': '1_2'
+                'id': '1_2' if isinstance(q, cirq.GridQubit) else '4'
             }]
         })
         g = GateWithAttribute(1.0)
         g.not_req = 'hello'
-        self.assertEqual(deserializer.from_proto(serialized),
-                         g(cirq.GridQubit(1, 2)))
+        self.assertEqual(deserializer.from_proto(serialized), g(q))
 
 
 if __name__ == "__main__":