def map_func(op: cirq.Operation, _: int) -> cirq.OP_TREE:
    if (op.gate == cirq.CX):
        (yield cirq.Z.on_each(*op.qubits))
        (yield cirq.CX(*op.qubits))
        (yield cirq.Z.on_each(*op.qubits))
    return op
def map_func(op: cirq.Operation, _: int) -> cirq.OP_TREE:
    if (op.gate == cirq.CX):
        (yield cirq.Z.on_each(*op.qubits))
        (yield cirq.CX(*op.qubits))
        (yield cirq.Z.on_each(*op.qubits))
    return op