From 0f4758d844c293d2c7491f9a764cc2137d282f6f Mon Sep 17 00:00:00 2001
From: Dennis Wayo <117969019+DennisWayo@users.noreply.github.com>
Date: Sat, 4 Jul 2026 00:27:06 +0500
Subject: [PATCH] Avoid stacking unused emulator coefficients

---
 .../_core/instruments/emulator/emulator.py    | 14 +++++++-----
 tests/instruments/emulator/test_emulator.py   | 22 +++++++++++++++++++
 2 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/src/qibolab/_core/instruments/emulator/emulator.py b/src/qibolab/_core/instruments/emulator/emulator.py
index e13c34e0a5..63871954b3 100644
--- a/src/qibolab/_core/instruments/emulator/emulator.py
+++ b/src/qibolab/_core/instruments/emulator/emulator.py
@@ -140,7 +140,7 @@ def _sweep(
         configs: dict[str, Config],
         sweepers: list[ParallelSweepers],
         updates: dict | None = None,
-    ) -> NDArray:
+    ) -> tuple[NDArray, NDArray | None]:
         """Sweep over sequence.
 
         This function invokes itself recursively, adding an array
@@ -155,7 +155,7 @@ def _sweep(
             return self._evolve(sequence, configs, updates)
 
         state_slices: list[NDArray] = []
-        coeff_slices = []
+        coeff_slices = [] if self.save_dir is not None else None
         parsweep = sweepers[0]
         # execute once for each parallel value
         for values in zip(*(s.values for s in parsweep)):
@@ -169,14 +169,18 @@ def _sweep(
                         updates[channel].update({sweeper.parameter.name: value})
             states, coeffs = self._sweep(sequence, configs, sweepers[1:], updates)
             state_slices.append(states)
-            coeff_slices.append(coeffs)
+            if coeff_slices is not None:
+                coeff_slices.append(coeffs)
 
         # stack all slices in a single array, along the current outermost dimension
-        return np.stack(state_slices), np.stack(coeff_slices)
+        return (
+            np.stack(state_slices),
+            np.stack(coeff_slices) if coeff_slices is not None else None,
+        )
 
     def _evolve(
         self, sequence: PulseSequence, configs: dict[str, Config], updates: dict
-    ) -> NDArray:
+    ) -> tuple[NDArray, NDArray | None]:
         """Evolve a pulse sequence on the quantum emulator.
 
         This method updates the sequence parameters, generates the time grid, constructs
diff --git a/tests/instruments/emulator/test_emulator.py b/tests/instruments/emulator/test_emulator.py
index 23b21b79f6..9b9c41ac20 100644
--- a/tests/instruments/emulator/test_emulator.py
+++ b/tests/instruments/emulator/test_emulator.py
@@ -4,6 +4,8 @@
 import pytest
 
 from qibolab._core.execution_parameters import AcquisitionType, AveragingMode
+from qibolab._core.pulses import Delay
+from qibolab._core.sequence import PulseSequence
 from qibolab._core.sweeper import Parameter, Sweeper
 
 NSHOTS = 1000
@@ -51,3 +53,23 @@ def test_sweepers(platform):
     acq_handle = list(seq.channel(platform.qubits[0].acquisition))[-1].id
     res = platform.execute([seq], [[sweeper]], nshots=NSHOTS)
     assert res[acq_handle].shape == (NSHOTS, 2)
+
+
+def test_duration_sweeper_with_variable_sequence_length(platform):
+    q0 = platform.natives.single_qubit[0]
+    delay = Delay(duration=0)
+    seq = q0.RX() | PulseSequence([(platform.qubits[0].acquisition, delay)]) | q0.MZ()
+    sweeper = Sweeper(
+        parameter=Parameter.duration, values=np.array([0, 2000]), pulses=[delay]
+    )
+    acq_handle = list(seq.channel(platform.qubits[0].acquisition))[-1].id
+
+    res = platform.execute(
+        [seq],
+        [[sweeper]],
+        nshots=NSHOTS,
+        acquisition_type=AcquisitionType.DISCRIMINATION,
+        averaging_mode=AveragingMode.SINGLESHOT,
+    )
+
+    assert res[acq_handle].shape == (NSHOTS, 2)