diff --git a/src/braket/devices/local_simulator.py b/src/braket/devices/local_simulator.py index f6e8ef22a..5c8e1c8ea 100644 --- a/src/braket/devices/local_simulator.py +++ b/src/braket/devices/local_simulator.py @@ -161,20 +161,22 @@ def run_batch( raise NotImplementedError("LocalSimulator.run_batch does not support per-task shots.") inputs = inputs or {} - if self._noise_model: - task_specifications = [ - self._noise_model.apply(task_specification) - for task_specification in task_specifications - ] - - if not max_parallel: - max_parallel = cpu_count() - single_task = isinstance( task_specifications, Circuit | OpenQASMProgram | Problem | AnalogHamiltonianSimulation, ) + if self._noise_model: + if single_task: + task_specifications = self._noise_model.apply(task_specifications) + else: + task_specifications = [ + self._noise_model.apply(task_specification) + for task_specification in task_specifications + ] + + max_parallel = max_parallel or cpu_count() + single_input = isinstance(inputs, dict) if not single_task and not single_input and len(task_specifications) != len(inputs): diff --git a/test/unit_tests/braket/devices/test_local_simulator.py b/test/unit_tests/braket/devices/test_local_simulator.py index 60fb36ac3..828591539 100644 --- a/test/unit_tests/braket/devices/test_local_simulator.py +++ b/test/unit_tests/braket/devices/test_local_simulator.py @@ -920,6 +920,22 @@ def test_run_batch_with_noise_model(mock_run_multiple, noise_model): assert mock_apply.call_count == 2 +@patch.object(DummyProgramDensityMatrixSimulator, "run_multiple") +def test_run_batch_single_circuit_with_noise_model(mock_run_multiple, noise_model): + mock_run_multiple.return_value = [GATE_MODEL_RESULT] + device = LocalSimulator("dummy_oq3_dm", noise_model=noise_model) + circuit = Circuit().h(0).cnot(0, 1) + + with patch.object(device._noise_model, "apply", wraps=device._noise_model.apply) as mock_apply: + results = device.run_batch(circuit, shots=4).results() + + assert len(results) == 1 + mock_apply.assert_called_once_with(circuit) + payloads = mock_run_multiple.call_args.args[0] + assert len(payloads) == 1 + assert "#pragma braket noise bit_flip(0.05) q[0]" in payloads[0].source + + @patch.object(DummyProgramDensityMatrixSimulator, "run") def test_run_noisy_circuit_with_noise_model(mock_run, noise_model): mock_run.return_value = GATE_MODEL_RESULT