New implementation of _group_interchangeable_qubits can create duplicate keys in returned Tuple

[vtomole]

Description of the issue
#4941 modified the implementation of _group_interchangeable_qubits, but in doing so inadvertently made a breaking change in how the returning Tuple is constructed.

The example code below shows how the constructed tuple would return ((0, frozenset({cirq.LineQubit(0), cirq.LineQubit(1), cirq.LineQubit(3)})), (2, frozenset({cirq.LineQubit(2)}))) with previous implementation.

Now it returns ((0, frozenset({cirq.LineQubit(0), cirq.LineQubit(1)})), (2,frozenset({cirq.LineQubit(2)})), (0, frozenset({cirq.LineQubit(3)}))).

The second tuple has 2 keys that are 0.

How to reproduce the issue

import itertools
from typing import Tuple, FrozenSet, Union, cast, Dict, List

import cirq
from cirq.ops import gate_features


class MyGate(cirq.Gate, cirq.InterchangeableQubitsGate):
    def __init__(self, num_qubits) -> None:
        self._num_qubits = num_qubits

    def num_qubits(self) -> int:
        return self._num_qubits

    def qubit_index_to_equivalence_group_key(self, index: int) -> int:
        if index % 2 == 0:
            return index
        return 0


class OldGroupInterchangeableQubitsGateOperation(cirq.GateOperation):
    def _group_interchangeable_qubits(
        self,
    ) -> Tuple[Union["cirq.Qid", Tuple[int, FrozenSet["cirq.Qid"]]], ...]:

        if not isinstance(self.gate, gate_features.InterchangeableQubitsGate):
            return self.qubits

        groups: Dict[int, List["cirq.Qid"]] = {}
        for i, q in enumerate(self.qubits):
            k = self.gate.qubit_index_to_equivalence_group_key(i)
            if k not in groups:
                groups[k] = []
            groups[k].append(q)
        returned_tuple = tuple(sorted((k, frozenset(v)) for k, v in groups.items()))
        print(returned_tuple)
        return returned_tuple


class NewGroupInterchangeableQubitsGateOperation(cirq.GateOperation):
    def _group_interchangeable_qubits(
        self,
    ) -> Tuple[Union["cirq.Qid", Tuple[int, FrozenSet["cirq.Qid"]]], ...]:
        if not isinstance(self.gate, gate_features.InterchangeableQubitsGate):
            return self.qubits
        else:

            def make_key(i_q: Tuple[int, "cirq.Qid"]) -> int:
                return cast(
                    gate_features.InterchangeableQubitsGate, self.gate
                ).qubit_index_to_equivalence_group_key(i_q[0])

            returned_tuple = tuple(
                (k, frozenset(g for _, g in kg))
                for k, kg in itertools.groupby(enumerate(self.qubits), make_key)
            )
            print(returned_tuple)
            return returned_tuple


qubits = cirq.LineQubit.range(4)
gate = MyGate(4)

print("Old group Tuples")
# prints
# ((0, frozenset({cirq.LineQubit(0), cirq.LineQubit(1), cirq.LineQubit(3)})), (2, frozenset({cirq.LineQubit(2)})))
# ((0, frozenset({cirq.LineQubit(0), cirq.LineQubit(2), cirq.LineQubit(3)})), (2, frozenset({cirq.LineQubit(1)})))
OldGroupInterchangeableQubitsGateOperation(
    gate, [qubits[0], qubits[1], qubits[2], qubits[3]]
) == OldGroupInterchangeableQubitsGateOperation(
    gate, [qubits[3], qubits[2], qubits[1], qubits[0]]
)

print("New group Tuples")
# prints
# ((0, frozenset({cirq.LineQubit(0), cirq.LineQubit(1)})), (2, frozenset({cirq.LineQubit(2)})), (0, frozenset({cirq.LineQubit(3)})))
# ((0, frozenset({cirq.LineQubit(2), cirq.LineQubit(3)})), (2, frozenset({cirq.LineQubit(1)})), (0, frozenset({cirq.LineQubit(0)})))
NewGroupInterchangeableQubitsGateOperation(
    gate, [qubits[0], qubits[1], qubits[2], qubits[3]]
) == NewGroupInterchangeableQubitsGateOperation(
    gate, [qubits[3], qubits[2], qubits[1], qubits[0]]
)

Cirq version
0.14.0

[MichaelBroughton]

Discussion from sync meeting: Lets try and get this back into a state where you can only have one key, and also add tests.

[vtomole]

Hey @Ashalynd , your pythonic implementation creates duplicate keys and I can't seem to find a way to get around that besides merging tuples with duplicate keys at the end. Doing this is rather complex so i'm thinking of reverting that code back to the non-pythonic version unless you have a better idea.

[Ashalynd]

Oh boy, that was a silly oversight on my side: https://docs.python.org/3/library/itertools.html#itertools.groupby
TL/DR for posterity: the stuff grouped by groupby had to be sorted prior to grouping. My bad for forgetting about it :-(