Support for stimcirq gates and operations in cirq_google protos

cirq-core/cirq/ops/gate_operation_test.py

@@ -490,7 +490,9 @@ def all_subclasses(cls):
     gate_subclasses = {
         g
         for g in all_subclasses(cirq.Gate)
-        if "cirq." in g.__module__ and "contrib" not in g.__module__ and "test" not in g.__module__
+        if g.__module__.startswith("cirq.")
+        and "contrib" not in g.__module__
+        and "test" not in g.__module__
     }
 
     test_module_spec = cirq.testing.json.spec_for("cirq.protocols")

cirq-google/cirq_google/serialization/circuit_serializer.py

@@ -191,7 +191,6 @@ def _serialize_gate_op(
             ValueError: If the operation cannot be serialized.
         """
         gate = op.gate
-
         if isinstance(gate, InternalGate):
             arg_func_langs.internal_gate_arg_to_proto(gate, out=msg.internalgate)
         elif isinstance(gate, cirq.XPowGate):
@@ -258,6 +257,23 @@ def _serialize_gate_op(
             arg_func_langs.float_arg_to_proto(
                 gate.q1_detune_mhz, out=msg.couplerpulsegate.q1_detune_mhz
             )
+        elif op.__module__.startswith('stimcirq') or (
+            gate is not None and gate.__module__.startswith('stimcirq')
+        ):
+            # Special handling for stimcirq objects, which can be both operations and gates.
+            stimcirq_obj = op if op.__module__.startswith('stimcirq') else gate
+            if stimcirq_obj is not None and hasattr(stimcirq_obj, '_json_dict_'):
+                # All stimcirq gates currently have _json_dict_defined
+                msg.internalgate.name = type(stimcirq_obj).__name__
+                msg.internalgate.module = 'stimcirq'
+                if isinstance(stimcirq_obj, cirq.Gate):
+                    msg.internalgate.num_qubits = stimcirq_obj.num_qubits()
+                else:
+                    msg.internalgate.num_qubits = len(stimcirq_obj.qubits)
+
+                # Store json_dict objects in gate_args
+                for k, v in stimcirq_obj._json_dict_().items():
+                    arg_func_langs.arg_to_proto(value=v, out=msg.internalgate.gate_args[k])
         else:
             raise ValueError(f'Cannot serialize op {op!r} of type {type(gate)}')
 
@@ -659,7 +675,31 @@ def _deserialize_gate_op(
                 raise ValueError(f"dimensions {dimensions} for ResetChannel must be an integer!")
             op = cirq.ResetChannel(dimension=dimensions)(*qubits)
         elif which_gate_type == 'internalgate':
-            op = arg_func_langs.internal_gate_from_proto(operation_proto.internalgate)(*qubits)
+            parsed_as_stimcirq = False
+            msg = operation_proto.internalgate
+            if msg.module == 'stimcirq':
+                # special handling for stimcirq
+                try:
+                    import stimcirq
+
+                    # Use JSON resolver to instantiate the object
+                    if msg.name in stimcirq.JSON_RESOLVERS_DICT:
+                        kwargs = {}
+                        for k, v in msg.gate_args.items():
+                            arg = arg_func_langs.arg_from_proto(v)
+                            if arg is not None:
+                                kwargs[k] = arg
+                        op = stimcirq.JSON_RESOLVERS_DICT[msg.name](**kwargs)
+                        if qubits:
+                            op = op(*qubits)
+                        parsed_as_stimcirq = True
+
+                except ModuleNotFoundError:  # pragma: no cover
+                    # fall back to creating internal gates if stimcirq not installed
+                    pass
+            if not parsed_as_stimcirq:
+                # all other internal gates
+                op = arg_func_langs.internal_gate_from_proto(msg)(*qubits)
         elif which_gate_type == 'couplerpulsegate':
             gate = CouplerPulse(
                 hold_time=cirq.Duration(

cirq-google/cirq_google/serialization/circuit_serializer_test.py

@@ -1034,3 +1034,22 @@ def test_reset_gate_with_improper_argument():
 
     with pytest.raises(ValueError, match="must be an integer"):
         serializer.deserialize(circuit_proto)
+
+
+def test_stimcirq_gates():
+    try:
+        import stimcirq
+    except ModuleNotFoundError:  # pragma: no cover
+        # Stimcirq not found, these are optional tests.
+        return
+    serializer = cg.CircuitSerializer()
+    q = cirq.q(1, 2)
+    q2 = cirq.q(2, 2)
+    c = cirq.Circuit(
+        cirq.Moment(stimcirq.CXSwapGate(inverted=True)(q, q2)),
+        cirq.Moment(cirq.measure(q, key="m")),
+        cirq.Moment(stimcirq.DetAnnotation(parity_keys=["m"])),
+    )
+    msg = serializer.serialize(c)
+    deserialized_circuit = serializer.deserialize(msg)
+    assert deserialized_circuit == c

dev_tools/conf/mypy.ini

@@ -21,7 +21,7 @@ follow_imports = silent
 ignore_missing_imports = true
 
 # Non-Google
-[mypy-IPython.*,sympy.*,matplotlib.*,proto.*,pandas.*,scipy.*,freezegun.*,mpl_toolkits.*,networkx.*,ply.*,astroid.*,pytest.*,_pytest.*,pylint.*,setuptools.*,qiskit.*,quimb.*,pylatex.*,filelock.*,sortedcontainers.*,tqdm.*,ruamel.*,absl.*,tensorflow_docs.*,ipywidgets.*,cachetools.*]
+[mypy-IPython.*,sympy.*,matplotlib.*,proto.*,pandas.*,scipy.*,freezegun.*,mpl_toolkits.*,networkx.*,ply.*,astroid.*,pytest.*,_pytest.*,pylint.*,setuptools.*,qiskit.*,quimb.*,pylatex.*,filelock.*,sortedcontainers.*,tqdm.*,ruamel.*,absl.*,tensorflow_docs.*,ipywidgets.*,cachetools.*,stimcirq.*]
 follow_imports = silent
 ignore_missing_imports = true
 

dev_tools/requirements/deps/dev-tools.txt

@@ -12,3 +12,6 @@ asv
 
 # For verifying behavior of qasm output.
 qiskit-aer~=0.16.1
+
+# For testing stimcirq compatibility (cirq-google)
+stimcirq