Standardize serialization of CalibrationTag (#7480)

- Serialize CalibrationTag as a tag rather than use a custom token index
field.
- This tag was not converted over to the new Tag message since I didn't
think anyone was still using it.

Author: dstrain115

cirq-google/cirq_google/api/v2/program.proto

@@ -81,9 +81,9 @@ message Circuit {
   // repeated moments.
   repeated int32 moment_indices = 3;
 
-  // Token that can be used to specify a version of a gate.
-  // For instance, a gate that has been calibrated for a circuit.
-  optional int32 token_constant_index = 4;
+  // Deprecated field, do not use.
+  reserved 4;
+  
   // Indices in the constant table for tags associated with the circuit
   repeated int32 tag_indices = 5;
 }
@@ -289,8 +289,8 @@ message Operation {
   // The token can be specified as a string or as a reference to
   // the constant table of the circuit.
   oneof token {
-    string token_value = 4;
-    int32 token_constant_index = 5;
+    string token_value = 4 [deprecated = true];
+    int32 token_constant_index = 5 [deprecated = true];
   }
 
   // To be deprecated
@@ -342,6 +342,9 @@ message Tag {
     // Uses parameter model to interpolate FSim gate.
     FSimViaModelTag fsim_via_model = 7;
 
+    // Calibration Tag
+    CalibrationTag calibration_tag = 9;
+
     // Catch-all for all gates that do not fit into the
     // above tags.
     InternalTag internal_tag = 8;
@@ -380,6 +383,12 @@ message NoSyncTag {
   }
 }
 
+// Tag to specify specific override tokens for operations or circuits.
+message CalibrationTag {
+  // Token to serialize
+  string token = 1;
+}
+
 // Tag to represent any internal tags or tags not yet
 // implemented in the proto.
 message InternalTag {

cirq-google/cirq_google/api/v2/program_pb2.py

@@ -17,6 +17,7 @@
 from typing import Any
 
 import cirq
+from cirq_google.api.v2 import program_pb2
 
 
 class CalibrationTag:
@@ -48,3 +49,15 @@ def __eq__(self, other) -> bool:
 
     def __hash__(self) -> int:
         return hash(self.token)
+
+    def to_proto(self, msg: program_pb2.Tag | None = None) -> program_pb2.Tag:
+        if msg is None:
+            msg = program_pb2.Tag()
+        msg.calibration_tag.token = self.token
+        return msg
+
+    @staticmethod
+    def from_proto(msg: program_pb2.Tag) -> CalibrationTag:
+        if msg.WhichOneof("tag") != "calibration_tag":
+            raise ValueError(f"Message is not a CalibrationTag, {msg}")
+        return CalibrationTag(token=msg.calibration_tag.token)

cirq-google/cirq_google/api/v2/program_pb2.pyi

@@ -14,8 +14,11 @@
 
 from __future__ import annotations
 
+import pytest
+
 import cirq
 import cirq_google
+from cirq_google.api.v2 import program_pb2
 
 
 def test_equality():
@@ -39,3 +42,14 @@ def test_str_repr():
     assert str(example_tag) == 'CalibrationTag(\'foo\')'
     assert repr(example_tag) == 'cirq_google.CalibrationTag(\'foo\')'
     cirq.testing.assert_equivalent_repr(example_tag, setup_code=('import cirq\nimport cirq_google'))
+
+
+def test_proto_serialization():
+    tag = cirq_google.CalibrationTag('foo')
+    msg = tag.to_proto()
+    assert tag == cirq_google.CalibrationTag.from_proto(msg)
+
+    with pytest.raises(ValueError, match="Message is not a CalibrationTag"):
+        msg = program_pb2.Tag()
+        msg.fsim_via_model.SetInParent()
+        cirq_google.CalibrationTag.from_proto(msg)

cirq-google/cirq_google/ops/calibration_tag.py

@@ -320,39 +320,22 @@ def _serialize_tag(
         raw_constants: dict[Any, int],
     ):
         constant = v2.program_pb2.Constant()
-        if isinstance(tag, CalibrationTag):
-            constant.string_value = tag.token
-            if tag.token in raw_constants:
-                msg.token_constant_index = raw_constants[tag.token]
+        if (tag_index := raw_constants.get(tag, None)) is None:
+            if self.tag_serializer and self.tag_serializer.can_serialize_tag(tag):
+                self.tag_serializer.to_proto(
+                    tag, msg=constant.tag_value, constants=constants, raw_constants=raw_constants
+                )
+            elif getattr(tag, 'to_proto', None) is not None:
+                tag.to_proto(constant.tag_value)  # type: ignore
             else:
-                # Token not found, add it to the list
-                msg.token_constant_index = len(constants)
+                warnings.warn(f'Unrecognized Tag {tag}, not serializing.')
+            if constant.WhichOneof('const_value'):
                 constants.append(constant)
                 if raw_constants is not None:
-                    raw_constants[tag.token] = msg.token_constant_index
+                    raw_constants[tag] = len(constants) - 1
+                msg.tag_indices.append(len(constants) - 1)
         else:
-            if isinstance(tag, DynamicalDecouplingTag):
-                # TODO(dstrain): Remove this once we are deserializing tag indices everywhere.
-                tag.to_proto(msg=msg.tags.add())
-            if (tag_index := raw_constants.get(tag, None)) is None:
-                if self.tag_serializer and self.tag_serializer.can_serialize_tag(tag):
-                    self.tag_serializer.to_proto(
-                        tag,
-                        msg=constant.tag_value,
-                        constants=constants,
-                        raw_constants=raw_constants,
-                    )
-                elif getattr(tag, 'to_proto', None) is not None:
-                    tag.to_proto(constant.tag_value)  # type: ignore
-                else:
-                    warnings.warn(f'Unrecognized Tag {tag}, not serializing.')
-                if constant.WhichOneof('const_value'):
-                    constants.append(constant)
-                    if raw_constants is not None:
-                        raw_constants[tag] = len(constants) - 1
-                    msg.tag_indices.append(len(constants) - 1)
-            else:
-                msg.tag_indices.append(tag_index)
+            msg.tag_indices.append(tag_index)
 
     def _serialize_circuit_op(
         self,
@@ -507,8 +490,6 @@ def _deserialize_circuit(
         for moment_index in circuit_proto.moment_indices:
             moments.append(deserialized_constants[moment_index])
 
-        if circuit_proto.HasField('token_constant_index'):
-            tags.append(CalibrationTag(constants[circuit_proto.token_constant_index].string_value))
         for tag_index in circuit_proto.tag_indices:
             tags.append(deserialized_constants[tag_index])
         return cirq.Circuit(moments, tags=tags)
@@ -879,6 +860,8 @@ def _deserialize_tag(self, msg: v2.program_pb2.Tag):
             return PhysicalZTag()
         elif which == 'fsim_via_model':
             return FSimViaModelTag()
+        elif which == 'calibration_tag':
+            return CalibrationTag.from_proto(msg)
         elif which == 'internal_tag':
             return InternalTag.from_proto(msg)
         else:

cirq-google/cirq_google/ops/calibration_tag_test.py

@@ -582,18 +582,18 @@ def test_serialize_deserialize_circuit_with_tokens():
     op_q0_tag1 = v2.program_pb2.Operation()
     op_q0_tag1.xpowgate.exponent.float_value = 1.0
     op_q0_tag1.qubit_constant_index.append(0)
-    op_q0_tag1.token_constant_index = 1
+    op_q0_tag1.tag_indices.append(1)
 
     op_q1_tag2 = v2.program_pb2.Operation()
     op_q1_tag2.xpowgate.exponent.float_value = 1.0
     op_q1_tag2.qubit_constant_index.append(3)
-    op_q1_tag2.token_constant_index = 4
+    op_q1_tag2.tag_indices.append(4)
 
     # Test repeated tag uses existing constant entey
     op_q0_tag2 = v2.program_pb2.Operation()
     op_q0_tag2.xpowgate.exponent.float_value = 1.0
     op_q0_tag2.qubit_constant_index.append(0)
-    op_q0_tag2.token_constant_index = 4
+    op_q0_tag2.tag_indices.append(4)
 
     proto = v2.program_pb2.Program(
         language=v2.program_pb2.Language(arg_function_language='exp', gate_set=_SERIALIZER_NAME),
@@ -602,10 +602,18 @@ def test_serialize_deserialize_circuit_with_tokens():
         ),
         constants=[
             v2.program_pb2.Constant(qubit=v2.program_pb2.Qubit(id='2_4')),
-            v2.program_pb2.Constant(string_value='abc123'),
+            v2.program_pb2.Constant(
+                tag_value=v2.program_pb2.Tag(
+                    calibration_tag=v2.program_pb2.CalibrationTag(token='abc123')
+                )
+            ),
             v2.program_pb2.Constant(operation_value=op_q0_tag1),
             v2.program_pb2.Constant(qubit=v2.program_pb2.Qubit(id='2_5')),
-            v2.program_pb2.Constant(string_value='def456'),
+            v2.program_pb2.Constant(
+                tag_value=v2.program_pb2.Tag(
+                    calibration_tag=v2.program_pb2.CalibrationTag(token='def456')
+                )
+            ),
             v2.program_pb2.Constant(operation_value=op_q1_tag2),
             v2.program_pb2.Constant(moment_value=v2.program_pb2.Moment(operation_indices=[2, 5])),
             v2.program_pb2.Constant(operation_value=op_q0_tag2),
@@ -628,12 +636,14 @@ def test_serialize_deserialize_circuit_tags():
     proto = v2.program_pb2.Program(
         language=v2.program_pb2.Language(arg_function_language='exp', gate_set=_SERIALIZER_NAME),
         circuit=v2.program_pb2.Circuit(
-            scheduling_strategy=v2.program_pb2.Circuit.MOMENT_BY_MOMENT,
-            tag_indices=[1],
-            token_constant_index=0,
+            scheduling_strategy=v2.program_pb2.Circuit.MOMENT_BY_MOMENT, tag_indices=[0, 1]
         ),
         constants=[
-            v2.program_pb2.Constant(string_value="abc123"),
+            v2.program_pb2.Constant(
+                tag_value=v2.program_pb2.Tag(
+                    calibration_tag=v2.program_pb2.CalibrationTag(token="abc123")
+                )
+            ),
             v2.program_pb2.Constant(
                 tag_value=v2.program_pb2.Tag(
                     internal_tag=v2.program_pb2.InternalTag(
@@ -697,7 +707,7 @@ def test_serialize_deserialize_circuit_with_subcircuit():
     op_tag = v2.program_pb2.Operation()
     op_tag.xpowgate.exponent.float_value = 1.0
     op_tag.qubit_constant_index.append(0)
-    op_tag.token_constant_index = 1
+    op_tag.tag_indices.append(1)
     op_symbol = v2.program_pb2.Operation()
     op_symbol.xpowgate.exponent.func.type = 'mul'
     op_symbol.xpowgate.exponent.func.args.add().arg_value.float_value = 2.0
@@ -726,7 +736,11 @@ def test_serialize_deserialize_circuit_with_subcircuit():
         ),
         constants=[
             v2.program_pb2.Constant(qubit=v2.program_pb2.Qubit(id='2_5')),
-            v2.program_pb2.Constant(string_value='abc123'),
+            v2.program_pb2.Constant(
+                tag_value=v2.program_pb2.Tag(
+                    calibration_tag=v2.program_pb2.CalibrationTag(token='abc123')
+                )
+            ),
             v2.program_pb2.Constant(operation_value=op_tag),
             v2.program_pb2.Constant(qubit=v2.program_pb2.Qubit(id='2_4')),
             v2.program_pb2.Constant(operation_value=op_symbol),