change: Ignore qubit_count parameter for JAQCD (#269)

amazon-braket · Jun 27, 2024 · dc68323 · dc68323
1 parent 3378eaa
commit dc68323
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 13 deletions.
diff --git a/src/braket/default_simulator/simulator.py b/src/braket/default_simulator/simulator.py
@@ -642,7 +642,7 @@ def run_jaqcd(
         Args:
             circuit_ir (Program): ir representation of a braket circuit specifying the
                 instructions to execute.
-            qubit_count (int): The number of qubits to simulate.
+            qubit_count (int): Unused parameter; in signature for backwards-compatibility
             shots (int): The number of times to run the circuit.
             batch_size (int): The size of the circuit partitions to contract,
                 if applying multiple gates at a time is desired; see `StateVectorSimulation`.
@@ -665,8 +665,9 @@ def run_jaqcd(
             circuit_ir,
             device_action_type=DeviceActionType.JAQCD,
         )
+        qubit_map = BaseLocalSimulator._map_circuit_to_contiguous_qubits(circuit_ir)
+        qubit_count = len(qubit_map)
         BaseLocalSimulator._validate_shots_and_ir_results(shots, circuit_ir.results, qubit_count)
-        BaseLocalSimulator._map_circuit_to_contiguous_qubits(circuit_ir)
 
         operations = [
             from_braket_instruction(instruction) for instruction in circuit_ir.instructions

diff --git a/test/unit_tests/braket/default_simulator/test_state_vector_simulator.py b/test/unit_tests/braket/default_simulator/test_state_vector_simulator.py
@@ -101,7 +101,8 @@ def test_simulator_run_bell_pair(bell_ir, batch_size, caplog):
     simulator = StateVectorSimulator()
     shots_count = 10000
     if isinstance(bell_ir, JaqcdProgram):
-        result = simulator.run(bell_ir, qubit_count=2, shots=shots_count, batch_size=batch_size)
+        # Ignore qubit_count
+        result = simulator.run(bell_ir, qubit_count=10, shots=shots_count, batch_size=batch_size)
     else:
         result = simulator.run(bell_ir, shots=shots_count, batch_size=batch_size)
 
@@ -749,14 +750,6 @@ def test_simulator_run_amplitude_shots():
 
 def test_simulator_run_amplitude_no_shots_invalid_states():
     simulator = StateVectorSimulator()
-    jaqcd = JaqcdProgram.parse_raw(
-        json.dumps(
-            {
-                "instructions": [{"type": "h", "target": 0}],
-                "results": [{"type": "amplitude", "states": ["0"]}],
-            }
-        )
-    )
     qasm = OpenQASMProgram(
         source="""
         qubit[2] q;
@@ -765,10 +758,18 @@ def test_simulator_run_amplitude_no_shots_invalid_states():
         #pragma braket result amplitude "0"
         """
     )
-    with pytest.raises(ValueError):
-        simulator.run(jaqcd, qubit_count=2, shots=0)
     with pytest.raises(ValueError):
         simulator.run(qasm, shots=0)
+    jaqcd = JaqcdProgram.parse_raw(
+        json.dumps(
+            {
+                "instructions": [{"type": "h", "target": 0}, {"type": "i", "target": 1}],
+                "results": [{"type": "amplitude", "states": ["0"]}],
+            }
+        )
+    )
+    with pytest.raises(ValueError):
+        simulator.run(jaqcd, qubit_count=2, shots=0)
 
 
 def test_simulator_run_statevector_shots():