fix: support single circuits with noise in run_batch

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):

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