fix: update type annotations and tests for Sequence[int] and QubitPermutationGate

Author: Aqil-Ahmad

cirq-core/cirq/sim/classical_simulator.py

@@ -38,7 +38,7 @@ def _is_identity(action) -> bool:
 class ClassicalBasisState(qis.QuantumStateRepresentation):
     """Represents a classical basis state for efficient state evolution."""
 
-    def __init__(self, initial_state: list[int] | np.ndarray):
+    def __init__(self, initial_state: Sequence[int] | np.ndarray):
         """Initializes the ClassicalBasisState object.
 
         Args:
@@ -85,21 +85,19 @@ def __init__(
 
         Args:
             qubits: The qubits to simulate.
-            initial_state: The initial state for the simulation. Accepts int or a sequence of int.
+            initial_state: The initial state for the simulation. Accepts int or Sequence[int].
             classical_data: The classical data container for the simulation.
 
         Raises:
             ValueError: If qubits not provided and initial_state is int.
-                If initial_state is not an int, list[int], tuple[int], or np.ndarray.
+                If initial_state is not an int or Sequence[int].
                 If initial_state is a np.ndarray and its shape is not 1-dimensional.
-                If gate is not one of X, SWAP, QubitPermutationGate, a controlled version
-                    of X or SWAP, or a measurement.
 
         An initial_state value of type integer is parsed in big endian order.
         """
         if isinstance(initial_state, int):
             if qubits is None:
-                raise ValueError('qubits must be provided if initial_state is not list[int]')
+                raise ValueError('qubits must be provided if initial_state is not Sequence[int]')
             state = ClassicalBasisState(
                 big_endian_int_to_bits(initial_state, bit_count=len(qubits))
             )
@@ -109,10 +107,10 @@ def __init__(
                     f'initial_state must be 1-dimensional, got shape {initial_state.shape}'
                 )
             state = ClassicalBasisState(list(initial_state))
-        elif isinstance(initial_state, (list, tuple)):
+        elif isinstance(initial_state, Sequence) and not isinstance(initial_state, (str, bytes)):
             state = ClassicalBasisState(list(initial_state))
         else:
-            raise ValueError('initial_state must be an int, list[int], tuple[int], or np.ndarray')
+            raise ValueError('initial_state must be an int or Sequence[int]')
         super().__init__(state=state, qubits=qubits, classical_data=classical_data)
 
     def _act_on_fallback_(self, action, qubits: Sequence[cirq.Qid], allow_decompose: bool = True):
@@ -125,6 +123,10 @@ def _act_on_fallback_(self, action, qubits: Sequence[cirq.Qid], allow_decompose:
 
         Returns:
             True if the operation was applied successfully.
+
+        Raises:
+            ValueError: If gate is not one of X, SWAP, QubitPermutationGate, a controlled version
+                of X or SWAP, or a measurement.
         """
         gate = action.gate if isinstance(action, ops.Operation) else action
         mapped_qubits = [self.qubit_map[i] for i in qubits]

cirq-core/cirq/sim/classical_simulator_test.py

@@ -60,13 +60,21 @@ def test_Swap():
     np.testing.assert_equal(results, expected_results)
 
 
-def test_qubit_permutation_gate():
-    q0, q1, q2 = cirq.LineQubit.range(3)
-    perm_gate = cirq.QubitPermutationGate([2, 0, 1])
-    circuit = cirq.Circuit(perm_gate(q0, q1, q2), cirq.measure(q0, q1, q2, key='key'))
+@pytest.mark.parametrize(
+    "n,perm,state",
+    [
+        (n, np.random.permutation(n).tolist(), np.random.choice(2, size=n))
+        for n in np.random.randint(3, 8, size=10)
+    ],
+)
+def test_qubit_permutation_gate(n, perm, state):
+    qubits = cirq.LineQubit.range(n)
+    perm_gate = cirq.QubitPermutationGate(perm)
+    circuit = cirq.Circuit(perm_gate(*qubits), cirq.measure(*qubits, key='key'))
     sim = cirq.ClassicalStateSimulator()
-    result = sim.simulate(circuit, initial_state=[1, 0, 1])
-    np.testing.assert_equal(result.measurements['key'], [1, 1, 0])
+    result = sim.simulate(circuit, initial_state=state)
+    expected = [state[perm[i]] for i in range(n)]
+    np.testing.assert_equal(result.measurements['key'], expected)
 
 
 def test_CCNOT():