From 564ca4001f1a88186ce0be54b15a74a5870cbc2d Mon Sep 17 00:00:00 2001 From: Lingavasan Date: Fri, 15 May 2026 23:30:30 -0700 Subject: [PATCH] Fix Polygraphy to-input multi-iteration merge Signed-off-by: Lingavasan --- .../polygraphy/tools/data/subtool/to_input.py | 4 +++- tools/Polygraphy/tests/tools/test_data.py | 16 +++++++++++++++- 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/tools/Polygraphy/polygraphy/tools/data/subtool/to_input.py b/tools/Polygraphy/polygraphy/tools/data/subtool/to_input.py index e2f92dde1..8b7a2e5fd 100644 --- a/tools/Polygraphy/polygraphy/tools/data/subtool/to_input.py +++ b/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) diff --git a/tools/Polygraphy/tests/tools/test_data.py b/tools/Polygraphy/tests/tools/test_data.py index e31f60dd3..cf0899bb9 100644 --- a/tools/Polygraphy/tests/tools/test_data.py +++ b/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"])