Fix Polygraphy to-input multi-iteration merge

tools/Polygraphy/polygraphy/tools/data/subtool/to_input.py

@@ -58,7 +58,9 @@ def update_inputs(new_inputs, path):
                 )
 
             # Pad to appropriate length
-            inputs += [OrderedDict()] * (len(new_inputs) - len(inputs))
+            inputs += [
+                OrderedDict() for _ in range(len(new_inputs) - len(inputs))
+            ]
 
             for inp, new_inp in zip(inputs, new_inputs):
                 inp.update(new_inp)

tools/Polygraphy/tests/tools/test_data.py

@@ -16,7 +16,7 @@
 #
 import numpy as np
 from polygraphy import util
-from polygraphy.json import load_json
+from polygraphy.json import load_json, save_json
 from tests.models.meta import ONNX_MODELS
 
 
@@ -40,3 +40,17 @@ def test_merge_inputs_outputs(self, poly_run, poly_data):
             assert len(merged_data) == 1
             assert list(merged_data[0].keys()) == ["x", "y"]
             assert all(isinstance(val, np.ndarray) for val in merged_data[0].values())
+
+    def test_preserves_multi_iteration_inputs(self, poly_data):
+        iter0 = {"x": np.array([0], dtype=np.float32)}
+        iter1 = {"x": np.array([1], dtype=np.float32)}
+
+        with util.NamedTemporaryFile() as inps, util.NamedTemporaryFile() as merged:
+            save_json([iter0, iter1], inps.name)
+
+            poly_data(["to-input", inps.name, "-o", merged.name])
+
+            merged_data = load_json(merged.name)
+            assert len(merged_data) == 2
+            assert np.array_equal(merged_data[0]["x"], iter0["x"])
+            assert np.array_equal(merged_data[1]["x"], iter1["x"])