Pad inhomogenous result arrays

cirq-core/cirq/sim/simulator_base.py

@@ -291,12 +291,19 @@ def _run(
             for k, r in step_result._classical_data.records.items():
                 if k not in records:
                     records[k] = []
-                records[k].append(r)
+                records[k].append(np.array(r, dtype=np.uint8))
             for k, cr in step_result._classical_data.channel_records.items():
                 if k not in records:
                     records[k] = []
-                records[k].append([cr])
-        return {str(k): np.array(v, dtype=np.uint8) for k, v in records.items()}
+                records[k].append(np.array([cr], dtype=np.uint8))
+
+        def pad_evenly(results: Sequence[np.ndarray]):
+            largest = max([result.shape[0] for result in results])
+            return np.array(
+                [np.pad(result, ((0, largest - result.shape[0]), (0, 0))) for result in results]
+            )
+
+        return {str(k): pad_evenly(v) for k, v in records.items()}
 
     def simulate_sweep_iter(
         self,

cirq-core/cirq/sim/simulator_base_test.py

@@ -394,6 +394,22 @@ def _has_unitary_(self):
     assert op2.count == 2
 
 
+def test_inhomogeneous_measurement_count_padding():
+    q = cirq.LineQubit(0)
+    key = cirq.MeasurementKey('m')
+    sim = cirq.Simulator()
+    c = cirq.Circuit(
+        cirq.CircuitOperation(
+            cirq.FrozenCircuit(cirq.X(q) ** 0.2, cirq.measure(q, key=key)),
+            use_repetition_ids=False,
+            repeat_until=cirq.KeyCondition(key),
+        )
+    )
+    results = sim.run(c, repetitions=10)
+    for i in range(10):
+        assert np.sum(results.records['m'][i, :, :]) == 1
+
+
 def test_deprecated_final_step_result():
     class OldCountingSimulator(CountingSimulator):
         def _create_simulator_trial_result(  # type: ignore