chore: import upstream snapshot with attribution

tools/Polygraphy/tests/comparator/test_comparator.py

@@ -0,0 +1,196 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import subprocess as sp
+
+import numpy as np
+import pytest
+import tensorrt as trt
+
+from polygraphy import util, mod
+from polygraphy.backend.onnx import GsFromOnnx, OnnxFromBytes
+from polygraphy.backend.onnxrt import OnnxrtRunner, SessionFromOnnx
+from polygraphy.backend.pluginref import PluginRefRunner
+from polygraphy.backend.trt import (
+    EngineFromNetwork,
+    NetworkFromOnnxBytes,
+    TrtRunner,
+    network_from_onnx_bytes,
+)
+from polygraphy.backend.trt.util import get_all_tensors
+from polygraphy.comparator import (
+    Comparator,
+    CompareFunc,
+    DataLoader,
+    IterationResult,
+    PostprocessFunc,
+    RunResults,
+)
+from polygraphy.exception import PolygraphyException
+from tests.models.meta import ONNX_MODELS
+
+build_torch = lambda a, **kwargs: util.array.to_torch(np.array(a, **kwargs))
+
+
+class TestComparator:
+    def test_warmup_runs(self):
+        onnx_loader = ONNX_MODELS["identity"].loader
+        runner = OnnxrtRunner(SessionFromOnnx(onnx_loader))
+        run_results = Comparator.run([runner], warm_up=2)
+        assert len(run_results[runner.name]) == 1
+
+    def test_list_as_data_loader(self):
+        onnx_loader = ONNX_MODELS["identity"].loader
+        runner = OnnxrtRunner(SessionFromOnnx(onnx_loader), name="onnx_runner")
+
+        data = [{"x": np.ones((1, 1, 2, 2), dtype=np.float32)}] * 2
+        run_results = Comparator.run([runner], data_loader=data)
+        iter_results = run_results["onnx_runner"]
+        assert len(iter_results) == 2
+        for actual, expected in zip(iter_results, data):
+            assert np.all(actual["y"] == expected["x"])
+
+    def test_generator_as_data_loader(self):
+        onnx_loader = ONNX_MODELS["identity"].loader
+        runner = OnnxrtRunner(SessionFromOnnx(onnx_loader), name="onnx_runner")
+
+        def data():
+            for feed_dict in [{"x": np.ones((1, 1, 2, 2), dtype=np.float32)}] * 2:
+                yield feed_dict
+
+        run_results = Comparator.run([runner], data_loader=data())
+        iter_results = run_results["onnx_runner"]
+        assert len(iter_results) == 2
+        for actual, expected in zip(iter_results, data()):
+            assert np.all(actual["y"] == expected["x"])
+
+    def test_multiple_runners(self):
+        onnx_bytes = ONNX_MODELS["identity"].loader()
+        build_onnxrt_session = SessionFromOnnx(onnx_bytes)
+        load_engine = EngineFromNetwork(NetworkFromOnnxBytes(onnx_bytes))
+        gs_graph = GsFromOnnx(OnnxFromBytes(onnx_bytes))
+
+        runners = [
+            OnnxrtRunner(build_onnxrt_session),
+            PluginRefRunner(gs_graph),
+            TrtRunner(load_engine),
+        ]
+
+        run_results = Comparator.run(runners)
+        compare_func = CompareFunc.simple(check_shapes=True)
+        assert bool(Comparator.compare_accuracy(run_results, compare_func=compare_func))
+        assert len(list(run_results.values())[0]) == 1  # Default number of iterations
+
+    def test_postprocess(self):
+        onnx_loader = ONNX_MODELS["identity"].loader
+        run_results = Comparator.run([OnnxrtRunner(SessionFromOnnx(onnx_loader))])
+        # Output shape is (1, 1, 2, 2)
+        postprocessed = Comparator.postprocess(
+            run_results, postprocess_func=PostprocessFunc.top_k(k=(1, -1))
+        )
+        for _, results in postprocessed.items():
+            for result in results:
+                for _, output in result.items():
+                    assert output.shape == (1, 1, 2, 1)
+
+    def test_errors_do_not_hang(self):
+        # Should error because interface is not implemented correctly.
+        class FakeRunner:
+            def __init__(self):
+                self.name = "fake"
+
+        runners = [FakeRunner()]
+        with pytest.raises(PolygraphyException):
+            Comparator.run(runners, use_subprocess=True, subprocess_polling_interval=1)
+
+    def test_segfault_does_not_hang(self):
+        def raise_called_process_error():
+            class FakeSegfault(sp.CalledProcessError):
+                pass
+
+            raise FakeSegfault(-11, ["simulate", "segfault"])
+
+        runners = [TrtRunner(EngineFromNetwork(raise_called_process_error))]
+        with pytest.raises(PolygraphyException):
+            Comparator.run(runners, use_subprocess=True, subprocess_polling_interval=1)
+
+    def test_multirun_outputs_are_different(self):
+        onnx_loader = ONNX_MODELS["identity"].loader
+        runner = TrtRunner(EngineFromNetwork(NetworkFromOnnxBytes(onnx_loader)))
+        run_results = Comparator.run([runner], data_loader=DataLoader(iterations=2))
+
+        iteration0 = run_results[runner.name][0]
+        iteration1 = run_results[runner.name][1]
+        for name in iteration0.keys():
+            assert util.array.any(iteration0[name] != iteration1[name])
+
+    @pytest.mark.parametrize("array_type", [np.array, build_torch])
+    def test_validate_nan(self, array_type):
+        run_results = RunResults()
+        run_results["fake-runner"] = [
+            IterationResult(outputs={"x": array_type(np.nan)})
+        ]
+        assert not Comparator.validate(run_results)
+
+    @pytest.mark.parametrize("array_type", [np.array, build_torch])
+    def test_validate_inf(self, array_type):
+        run_results = RunResults()
+        run_results["fake-runner"] = [
+            IterationResult(outputs={"x": array_type(np.inf)})
+        ]
+        assert not Comparator.validate(run_results, check_inf=True)
+
+    def test_dim_param_trt_onnxrt(self):
+        load_onnx_bytes = ONNX_MODELS["dim_param"].loader
+        build_onnxrt_session = SessionFromOnnx(load_onnx_bytes)
+        load_engine = EngineFromNetwork(NetworkFromOnnxBytes(load_onnx_bytes))
+
+        runners = [
+            OnnxrtRunner(build_onnxrt_session),
+            TrtRunner(load_engine),
+        ]
+
+        run_results = Comparator.run(runners)
+        compare_func = CompareFunc.simple(check_shapes=True)
+        assert bool(Comparator.compare_accuracy(run_results, compare_func=compare_func))
+        assert len(list(run_results.values())[0]) == 1  # Default number of iterations
+
+    @pytest.mark.skipif(
+        mod.version(trt.__version__) < mod.version("10.0"),
+        reason="Feature not present before 10.0",
+    )
+    def test_debug_tensors(self):
+        model = ONNX_MODELS["identity"]
+        builder, network, parser = network_from_onnx_bytes(model.loader)
+        tensor_map = get_all_tensors(network)
+        network.mark_debug(tensor_map["x"])
+        load_engine = EngineFromNetwork((builder, network, parser))
+        runners = [TrtRunner(load_engine)]
+        data = [{"x": np.ones((1, 1, 2, 2), dtype=np.float32)}]
+        run_results = Comparator.run(runners, data_loader=data)
+        for iteration_list in run_results.values():
+            # There should be 2 outputs, debug tensor "x" and output "y"
+            assert len(list(iteration_list[0].items())) == 2
+        run_results["fake-runner"] = [
+            IterationResult(
+                outputs={
+                    "x": np.ones((1, 1, 2, 2), dtype=np.float32),
+                    "y": np.ones((1, 1, 2, 2), dtype=np.float32),
+                }
+            )
+        ]
+        compare_func = CompareFunc.simple(check_shapes=True)
+        assert bool(Comparator.compare_accuracy(run_results, compare_func=compare_func))

tools/Polygraphy/tests/comparator/test_compare.py

@@ -0,0 +1,432 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import numpy as np
+import pytest
+from polygraphy import util
+from polygraphy.comparator import CompareFunc, IterationResult
+from polygraphy.datatype import DataType
+from polygraphy.exception import PolygraphyException
+from polygraphy.logger import G_LOGGER
+
+build_torch = lambda a, **kwargs: util.array.to_torch(np.array(a, **kwargs))
+
+
+@pytest.mark.parametrize("array_type", [np.array, build_torch], ids=["numpy", "torch"])
+class TestSimpleCompareFunc:
+    @pytest.mark.parametrize(
+        "values0, values1, dtype, expected_max_absdiff, expected_max_reldiff",
+        [
+            # Low precision arrays should be casted to higher precisions to avoid overflows/underflows.
+            ([0], [1], DataType.UINT8, 1, 1.0),
+            ([1], [0], DataType.UINT8, 1, np.inf),
+            ([0], [1], DataType.UINT16, 1, 1.0),
+            ([1], [0], DataType.UINT16, 1, np.inf),
+            ([0], [1], DataType.UINT32, 1, 1.0),
+            ([1], [0], DataType.UINT32, 1, np.inf),
+            ([25], [30], DataType.INT8, 5, 5.0 / 30.0),
+            (
+                [25],
+                [30],
+                DataType.FLOAT16,
+                5,
+                np.array([5.0], dtype=np.float32) / np.array([30.0], dtype=np.float32),
+            ),
+            ([1], [0], DataType.FLOAT16, 1, 1 / np.finfo(float).eps),
+        ],
+    )
+    def test_comparison(
+        self,
+        values0,
+        values1,
+        dtype,
+        expected_max_absdiff,
+        expected_max_reldiff,
+        array_type,
+    ):
+        if array_type != np.array:
+            try:
+                DataType.to_dtype(dtype, "torch")
+            except:
+                pytest.skip(f"Cannot convert {dtype} to torch")
+        iter_result0 = IterationResult(
+            outputs={"output": array_type(values0, dtype=dtype.numpy())}
+        )
+        iter_result1 = IterationResult(
+            outputs={"output": array_type(values1, dtype=dtype.numpy())}
+        )
+
+        compare_func = CompareFunc.simple()
+        acc = compare_func(iter_result0, iter_result1)
+
+        comp_result = acc["output"]
+        assert np.isclose(comp_result.max_absdiff, expected_max_absdiff)
+        assert np.isclose(comp_result.max_absdiff, comp_result.mean_absdiff)
+        assert np.isclose(comp_result.max_absdiff, comp_result.median_absdiff)
+
+        assert np.isclose(comp_result.max_reldiff, expected_max_reldiff)
+        assert np.isclose(comp_result.max_reldiff, comp_result.mean_reldiff)
+        assert np.isclose(comp_result.max_reldiff, comp_result.median_reldiff)
+
+    @pytest.mark.parametrize(
+        "values0, values1, dtype, quantile, expected_abs_quantile, expected_rel_quantile",
+        [
+            ([0, 0.1, 0.5, 0.75, 1], [2, 2, 2, 2, 2], DataType.FLOAT16, 0.5, 1.5, 0.75),
+            (
+                [0, 0.1, 0.5, 0.75, 1],
+                [2, 2, 2, 2, 2],
+                DataType.FLOAT16,
+                0.75,
+                1.9,
+                0.95,
+            ),
+            (
+                [0.2, 0.2, 0.125, 0.11],
+                [0.1, 0.1, 0.1, 0.1],
+                DataType.FLOAT16,
+                0.5,
+                0.0625,
+                0.625,
+            ),
+            ([0, 1, 2, 3, 4], [2, 2, 2, 2, 2], DataType.UINT8, 0.5, 1, 0.5),
+            ([0, 1, 2, 3, 4], [2, 2, 2, 2, 2], DataType.UINT8, 0.75, 2, 1),
+        ],
+    )
+    def test_quantile(
+        self,
+        values0,
+        values1,
+        dtype,
+        quantile,
+        expected_abs_quantile,
+        expected_rel_quantile,
+        array_type,
+    ):
+        if array_type != np.array:
+            try:
+                DataType.to_dtype(dtype, "torch")
+            except:
+                pytest.skip(f"Cannot convert {dtype} to torch")
+        iter_result0 = IterationResult(
+            outputs={"output": array_type(values0, dtype=dtype.numpy())}
+        )
+        iter_result1 = IterationResult(
+            outputs={"output": array_type(values1, dtype=dtype.numpy())}
+        )
+
+        compare_func = CompareFunc.simple(
+            check_error_stat="quantile", error_quantile=quantile
+        )
+        acc = compare_func(iter_result0, iter_result1)
+
+        comp_result = acc["output"]
+        assert np.isclose(
+            comp_result.quantile_absdiff, expected_abs_quantile, atol=1e-4, rtol=1e-4
+        )
+        assert comp_result.quantile_absdiff <= comp_result.max_absdiff
+        assert (quantile >= 0.5) == (
+            comp_result.quantile_absdiff >= comp_result.median_absdiff
+        )
+
+        assert np.isclose(
+            comp_result.quantile_reldiff, expected_rel_quantile, atol=1e-4, rtol=1e-4
+        )
+        assert comp_result.quantile_reldiff <= comp_result.max_reldiff
+        assert (quantile >= 0.5) == (
+            comp_result.quantile_reldiff >= comp_result.median_reldiff
+        )
+
+    def test_can_compare_bool(self, array_type):
+        iter_result0 = IterationResult(
+            outputs={"output": array_type(np.zeros((4, 4), dtype=bool))}
+        )
+        iter_result1 = IterationResult(
+            outputs={"output": array_type(np.ones((4, 4), dtype=bool))}
+        )
+
+        compare_func = CompareFunc.simple()
+        acc = compare_func(iter_result0, iter_result1)
+
+        assert not acc["output"]
+
+    @pytest.mark.parametrize("mode", ["abs", "rel"])
+    def test_per_output_tol(self, mode, array_type):
+        OUT0_NAME = "output0"
+        OUT1_NAME = "output1"
+        OUT_VALS = array_type(np.ones((4, 4)))
+
+        iter_result0 = IterationResult(
+            outputs={OUT0_NAME: OUT_VALS, OUT1_NAME: OUT_VALS}
+        )
+        iter_result1 = IterationResult(
+            outputs={OUT0_NAME: OUT_VALS, OUT1_NAME: OUT_VALS + 1}
+        )
+
+        # With default tolerances, out1 is wrong for the second result.
+        compare_func = CompareFunc.simple()
+        acc = compare_func(iter_result0, iter_result1)
+        assert bool(acc[OUT0_NAME])
+        assert not bool(acc[OUT1_NAME])
+
+        # But with custom tolerances, it should pass.
+        tols = {
+            OUT0_NAME: 0.0,
+            OUT1_NAME: 1.0,
+        }
+
+        if mode == "abs":
+            compare_func = CompareFunc.simple(atol=tols)
+        else:
+            compare_func = CompareFunc.simple(rtol=tols)
+
+        acc = compare_func(iter_result0, iter_result1)
+        assert bool(acc[OUT0_NAME])
+        assert bool(acc[OUT1_NAME])
+
+    @pytest.mark.parametrize("mode", ["abs", "rel"])
+    def test_per_output_tol_fallback(self, mode, array_type):
+        OUT0_NAME = "output0"
+        OUT1_NAME = "output1"
+        OUT_VALS = array_type(np.ones((4, 4)))
+
+        iter_result0 = IterationResult(
+            outputs={OUT0_NAME: OUT_VALS + 1, OUT1_NAME: OUT_VALS}
+        )
+        iter_result1 = IterationResult(
+            outputs={OUT0_NAME: OUT_VALS, OUT1_NAME: OUT_VALS + 1}
+        )
+
+        acc = CompareFunc.simple()(iter_result0, iter_result1)
+        assert not bool(acc[OUT0_NAME])
+        assert not bool(acc[OUT1_NAME])
+
+        # Do not specify tolerance for OUT0_NAME - it should fail with fallback tolerance
+        tols = {
+            OUT1_NAME: 1.0,
+        }
+
+        if mode == "abs":
+            compare_func = CompareFunc.simple(atol=tols)
+        else:
+            compare_func = CompareFunc.simple(rtol=tols)
+
+        acc = compare_func(iter_result0, iter_result1)
+        assert not bool(acc[OUT0_NAME])
+        assert bool(acc[OUT1_NAME])
+
+    @pytest.mark.parametrize("mode", ["abs", "rel"])
+    def test_default_tol_in_map(self, mode, array_type):
+        # "" can be used to indicate a global tolerance
+        OUT0_NAME = "output0"
+        OUT_VALS = array_type(np.ones((4, 4)))
+
+        iter_result0 = IterationResult(outputs={OUT0_NAME: OUT_VALS})
+        iter_result1 = IterationResult(outputs={OUT0_NAME: OUT_VALS + 1})
+
+        tols = {
+            "": 1.0,
+        }
+
+        if mode == "abs":
+            compare_func = CompareFunc.simple(atol=tols)
+        else:
+            compare_func = CompareFunc.simple(rtol=tols)
+
+        acc = compare_func(iter_result0, iter_result1)
+        assert bool(acc[OUT0_NAME])
+
+    @pytest.mark.parametrize(
+        "shape",
+        [
+            tuple(),
+            (0, 2, 1, 2),
+            (1,),
+            (2, 2, 2, 2),
+        ],
+    )
+    def test_non_matching_outputs(self, shape, array_type):
+        iter_result0 = IterationResult(
+            outputs={"output": array_type(np.zeros(shape, dtype=np.float32))}
+        )
+        iter_result1 = IterationResult(
+            outputs={"output": array_type(np.ones(shape, dtype=np.float32))}
+        )
+
+        compare_func = CompareFunc.simple()
+
+        with G_LOGGER.verbosity(G_LOGGER.ULTRA_VERBOSE):
+            acc = compare_func(iter_result0, iter_result1)
+
+        assert util.is_empty_shape(shape) or not acc["output"]
+
+    @pytest.mark.parametrize("check_error_stat", ["max", "median", "mean", "elemwise"])
+    @pytest.mark.parametrize(
+        "func",
+        [
+            np.zeros,
+            np.ones,
+        ],
+    )
+    def test_check_error_stat(self, func, check_error_stat, array_type):
+        iter_result0 = IterationResult(
+            outputs={"output": array_type(func((100,), dtype=np.float32))}
+        )
+        iter_result1 = IterationResult(
+            outputs={"output": array_type(func((100,), dtype=np.float32))}
+        )
+
+        iter_result0["output"][0] += 100
+
+        # Even though the max diff is 100, atol=1 should cause this to pass since we're checking
+        # against the mean error.
+        compare_func = CompareFunc.simple(check_error_stat=check_error_stat, atol=1)
+
+        if check_error_stat in ["max", "elemwise"]:
+            assert not compare_func(iter_result0, iter_result1)["output"]
+        else:
+            assert compare_func(iter_result0, iter_result1)["output"]
+
+    @pytest.mark.parametrize("check_error_stat", ["max", "median", "mean", "elemwise"])
+    def test_atol_rtol_either_pass(self, check_error_stat, array_type):
+        # If either rtol/atol is sufficient, the compare_func should pass
+        res0 = IterationResult(outputs={"output": array_type([1, 2], dtype=np.float32)})
+        res1 = IterationResult(
+            outputs={"output": array_type((1.25, 2.5), dtype=np.float32)}
+        )
+
+        assert not CompareFunc.simple(check_error_stat=check_error_stat)(res0, res1)[
+            "output"
+        ]
+
+        assert CompareFunc.simple(check_error_stat=check_error_stat, rtol=0.25)(
+            res0, res1
+        )["output"]
+        assert CompareFunc.simple(check_error_stat=check_error_stat, atol=0.5)(
+            res0, res1
+        )["output"]
+
+    def test_atol_rtol_combined_pass(self, array_type):
+        # We should also be able to mix them - i.e. rtol might enough for some, atol for others.
+        # If they cover the entire output range, it should pass.
+        res0 = IterationResult(
+            outputs={"output": array_type([0, 1, 2, 3], dtype=np.float32)}
+        )
+        res1 = IterationResult(
+            outputs={"output": array_type((0.15, 1.25, 2.5, 3.75), dtype=np.float32)}
+        )
+
+        assert not CompareFunc.simple()(res0, res1)["output"]
+
+        assert not CompareFunc.simple(atol=0.3)(res0, res1)["output"]
+        assert not CompareFunc.simple(rtol=0.25)(res0, res1)["output"]
+
+        assert CompareFunc.simple(atol=0.3, rtol=0.25)(res0, res1)["output"]
+
+    @pytest.mark.parametrize(
+        "check_error_stat",
+        [
+            {"output0": "mean", "output1": "max"},
+            {"": "mean", "output1": "elemwise"},
+            {"output0": "mean"},
+            {"": "mean"},
+        ],
+    )
+    def test_per_output_error_stat(self, check_error_stat, array_type):
+        # output0 will only pass when using check_error_stat=mean
+        res0 = IterationResult(
+            outputs={
+                "output0": array_type([0, 1, 2, 3], dtype=np.float32),
+                "output1": array_type([0, 1, 2, 3], dtype=np.float32),
+            }
+        )
+        res1 = IterationResult(
+            outputs={
+                "output0": array_type((0.15, 1.25, 2.5, 3.75), dtype=np.float32),
+                "output1": array_type((0, 1, 2, 3), dtype=np.float32),
+            }
+        )
+
+        atol = 0.4125
+        assert not CompareFunc.simple(atol=atol)(res0, res1)["output0"]
+
+        assert CompareFunc.simple(check_error_stat=check_error_stat, atol=atol)(
+            res0, res1
+        )["output0"]
+        assert CompareFunc.simple(check_error_stat=check_error_stat, atol=atol)(
+            res0, res1
+        )["output1"]
+
+    def test_invalid_error_stat(self, array_type):
+        res0 = IterationResult(
+            outputs={"output": array_type([0, 1, 2, 3], dtype=np.float32)}
+        )
+        res1 = IterationResult(
+            outputs={"output": array_type([0.15, 1.25, 2.5, 3.75], dtype=np.float32)}
+        )
+
+        with pytest.raises(PolygraphyException, match="Invalid choice"):
+            CompareFunc.simple(check_error_stat="invalid-stat")(res0, res1)
+
+    @pytest.mark.parametrize("check_error_stat", ["max", "median", "mean", "elemwise"])
+    @pytest.mark.parametrize("val0, val1", [(np.nan, 0.15), (0.15, np.nan)])
+    def test_nans_always_fail(self, check_error_stat, val0, val1, array_type):
+        res0 = IterationResult(outputs={"output": array_type([val0], dtype=np.float32)})
+        res1 = IterationResult(outputs={"output": array_type([val1], dtype=np.float32)})
+
+        assert not CompareFunc.simple(check_error_stat=check_error_stat)(res0, res1)[
+            "output"
+        ]
+
+    @pytest.mark.parametrize("infinities_compare_equal", (False, True))
+    @pytest.mark.parametrize("val", (np.inf, -np.inf))
+    def test_infinities_compare_equal(self, infinities_compare_equal, val, array_type):
+        res0 = IterationResult(outputs={"output": array_type([val], dtype=np.float32)})
+        res1 = IterationResult(outputs={"output": array_type([val], dtype=np.float32)})
+
+        cf = CompareFunc.simple(infinities_compare_equal=infinities_compare_equal)
+        assert bool(cf(res0, res1)["output"]) == infinities_compare_equal
+
+
+@pytest.mark.parametrize("array_type", [np.array, build_torch])
+class TestIndicesCompareFunc:
+    @pytest.mark.parametrize(
+        "out0,out1,index_tolerance,expected",
+        [
+            ([0, 1, 2, 3], [0, 1, 2, 3], 0, True),
+            # Check that dictionaries work for index tolerance
+            ([0, 1, 2, 3], [0, 1, 2, 3], {"": 0}, True),
+            ([0, 1, 2, 3], [0, 1, 2, 3], {"output": 0}, True),
+            ([[0, 1], [0, 1], [0, 1]], [[0, 1], [0, 1], [0, 1]], 0, True),
+            ([1, 0, 2, 3], [0, 1, 2, 3], 0, False),
+            ([1, 0, 2, 3], [0, 1, 2, 3], 1, True),
+            ([0, 1, 2, 3], [0, 1, 3, 2], 1, True),
+            # Last 'index_tolerance' indices should be ignored.
+            ([1, 0, 2, 7], [0, 1, 2, 3], 0, False),
+            ([1, 0, 2, 7], [0, 1, 2, 3], 1, True),
+            ([[2, 3, 4], [5, 6, 9]], [[3, 2, 4], [5, 9, 6]], 0, False),
+            ([[2, 3, 4], [5, 6, 9]], [[3, 2, 4], [5, 9, 6]], 1, True),
+            ([0, 1, 2, 3, 4, 5, 6], [0, 3, 2, 1, 4, 5, 6], 0, False),
+            ([0, 1, 2, 3, 4, 5, 6], [0, 3, 2, 1, 4, 5, 6], 2, True),
+        ],
+    )
+    def test_index_tolerance(self, out0, out1, index_tolerance, expected, array_type):
+        res0 = IterationResult(outputs={"output": array_type(out0, dtype=np.int32)})
+        res1 = IterationResult(outputs={"output": array_type(out1, dtype=np.int32)})
+
+        assert (
+            CompareFunc.indices(index_tolerance=index_tolerance)(res0, res1)["output"]
+            == expected
+        )

tools/Polygraphy/tests/comparator/test_data_loader.py

@@ -0,0 +1,302 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+from collections import OrderedDict
+
+import numpy as np
+import torch
+import pytest
+
+from polygraphy import constants, util
+from polygraphy.common import TensorMetadata
+from polygraphy.comparator import DataLoader
+from polygraphy.comparator.data_loader import DataLoaderCache
+from polygraphy.datatype import DataType
+from tests.models.meta import ONNX_MODELS
+from polygraphy.exception import PolygraphyException
+
+
+def meta(dtype):
+    return (
+        TensorMetadata()
+        .add("X", dtype=dtype, shape=(4, 4))
+        .add("Y", dtype=dtype, shape=(5, 5))
+    )
+
+
+class TestDataLoader:
+    @pytest.mark.parametrize("dtype", [np.int32, bool, np.float32, np.int64])
+    def test_default_ranges(self, dtype):
+        data_loader = DataLoader(input_metadata=meta(dtype))
+        x, y = data_loader[0].values()
+        assert np.all((x >= 0) & (x <= 1))
+        assert np.all((y >= 0) & (y <= 1))
+
+    def test_can_override_shape(self):
+        model = ONNX_MODELS["dynamic_identity"]
+
+        shape = (1, 1, 4, 5)
+        custom_input_metadata = TensorMetadata().add("X", dtype=None, shape=shape)
+        data_loader = DataLoader(input_metadata=custom_input_metadata)
+        # Simulate what the comparator does
+        data_loader.input_metadata = model.input_metadata
+
+        feed_dict = data_loader[0]
+        assert tuple(feed_dict["X"].shape) == shape
+
+    @pytest.mark.parametrize(
+        "min_shape, max_shape, expected",
+        [
+            # When both min/max are set, use min.
+            ((2, 3, 2, 2), (4, 3, 2, 2), (2, 3, 2, 2)),
+            # When only one of min/max are set, use whichever one is set.
+            ((2, 3, 2, 2), None, (2, 3, 2, 2)),
+            (None, (4, 3, 2, 2), (4, 3, 2, 2)),
+            # When min/max are not set, override with the default shape value.
+            (None, None, (constants.DEFAULT_SHAPE_VALUE, 3, 2, 2)),
+        ],
+    )
+    def test_can_use_min_max_shape(self, min_shape, max_shape, expected):
+        shape = (-1, 3, 2, 2)
+
+        data_loader = DataLoader()
+        data_loader.input_metadata = TensorMetadata().add(
+            "X", dtype=np.float32, shape=shape, min_shape=min_shape, max_shape=max_shape
+        )
+
+        feed_dict = data_loader[0]
+        assert tuple(feed_dict["X"].shape) == expected
+
+    @pytest.mark.parametrize("dtype", [np.int32, bool, np.float32, np.int64])
+    @pytest.mark.parametrize("range_val", [0, 1])
+    def test_range_min_max_equal(self, dtype, range_val):
+        data_loader = DataLoader(
+            input_metadata=meta(dtype), val_range=(range_val, range_val)
+        )
+        feed_dict = data_loader[0]
+        assert np.all(feed_dict["X"] == range_val)
+        assert np.all(feed_dict["Y"] == range_val)
+
+    @pytest.mark.parametrize(
+        "range",
+        [
+            (0, 1, np.int32),
+            (5.0, 5.5, np.float32),
+            (0, 1, bool),
+            (float("inf"), float("inf"), np.float32),
+            (float("-inf"), float("inf"), np.float32),
+            (0, float("inf"), np.float32),
+            (float("-inf"), 0, np.float32),
+        ],
+    )
+    def test_val_ranges(self, range):
+        min_val, max_val, dtype = range
+        data_loader = DataLoader(
+            input_metadata=meta(dtype), val_range=(min_val, max_val)
+        )
+        feed_dict = data_loader[0]
+        assert np.all((feed_dict["X"] >= min_val) & (feed_dict["X"] <= max_val))
+
+    @pytest.mark.parametrize("dtype", [np.int32, np.int64, np.float32])
+    def test_val_range_dict(self, dtype):
+        val_range = {"X": (2, 5), "Y": (-1, 2)}
+        data_loader = DataLoader(input_metadata=meta(dtype), val_range=val_range)
+        feed_dict = data_loader[0]
+        assert np.all((feed_dict["X"] >= 2) & (feed_dict["X"] <= 5))
+        assert np.all((feed_dict["Y"] >= -1) & (feed_dict["Y"] <= 2))
+
+    @pytest.mark.parametrize("dtype", [np.int32, np.int64, np.float32])
+    def test_val_range_dict_default(self, dtype):
+        val_range = {"": (6, 8), "Y": (-3, 4)}
+        data_loader = DataLoader(input_metadata=meta(dtype), val_range=val_range)
+        feed_dict = data_loader[0]
+        assert np.all((feed_dict["X"] >= 6) & (feed_dict["X"] <= 8))
+        assert np.all((feed_dict["Y"] >= -3) & (feed_dict["Y"] <= 4))
+
+    @pytest.mark.parametrize("dtype", [np.int32, np.int64, np.float32])
+    def test_val_range_dict_fallback(self, dtype):
+        val_range = {"Y": (-3, 4)}
+        data_loader = DataLoader(input_metadata=meta(dtype), val_range=val_range)
+        feed_dict = data_loader[0]
+        assert np.all((feed_dict["X"] >= 0) & (feed_dict["X"] <= 1))
+        assert np.all((feed_dict["Y"] >= -3) & (feed_dict["Y"] <= 4))
+
+    def test_shape_tensor_detected(self):
+        INPUT_DATA = (1, 2, 3)
+        input_meta = TensorMetadata().add("X", dtype=np.int32, shape=(3,))
+        # This contains the shape values
+        overriden_meta = TensorMetadata().add("X", dtype=np.int32, shape=INPUT_DATA)
+        data_loader = DataLoader(input_metadata=overriden_meta)
+        data_loader.input_metadata = input_meta
+
+        feed_dict = data_loader[0]
+        assert np.all(feed_dict["X"] == INPUT_DATA)  # values become INPUT_DATA
+
+    def test_no_shape_tensor_false_positive_negative_dims(self):
+        INPUT_DATA = (-100, 2, 4)
+        # This should NOT be detected as a shape tensor
+        input_meta = TensorMetadata().add("X", dtype=np.int32, shape=(3,))
+        overriden_meta = TensorMetadata().add("X", dtype=np.int32, shape=INPUT_DATA)
+        data_loader = DataLoader(input_metadata=overriden_meta)
+        data_loader.input_metadata = input_meta
+
+        feed_dict = data_loader[0]
+        assert feed_dict["X"].shape == (
+            3,
+        )  # Shape IS (3, ), because this is NOT a shape tensor
+        assert np.any(
+            feed_dict["X"] != INPUT_DATA
+        )  # Contents are not INPUT_DATA, since it's not treated as a shape value
+
+    def test_no_shape_tensor_false_positive_float(self):
+        INPUT_DATA = (-100, -50, 0)
+        # Float cannot be a shape tensor
+        input_meta = TensorMetadata().add("X", dtype=np.float32, shape=(3,))
+        overriden_meta = TensorMetadata().add("X", dtype=np.float32, shape=INPUT_DATA)
+        data_loader = DataLoader(input_metadata=overriden_meta)
+        data_loader.input_metadata = input_meta
+
+        feed_dict = data_loader[0]
+        assert feed_dict["X"].shape == (3,)  # Values are NOT (3, )
+        assert np.any(feed_dict["X"] != INPUT_DATA)  # Values are NOT (3, )
+
+    def test_non_user_provided_inputs_never_shape_tensors(self):
+        # If the user didn't provide metadata, then the value can never be a shape tensor.
+        input_meta = TensorMetadata().add("X", dtype=np.int32, shape=(3,))
+        data_loader = DataLoader()
+        data_loader.input_metadata = input_meta
+
+        feed_dict = data_loader[0]
+        assert feed_dict["X"].shape == (3,)  # Treat as a normal tensor
+
+    @pytest.mark.parametrize("dtype", [np.float32, np.int32])
+    @pytest.mark.parametrize("data_loader_backend_module", ["torch", "numpy"])
+    def test_generate_scalar(self, dtype, data_loader_backend_module):
+        data_loader = DataLoader(
+            input_metadata=TensorMetadata().add("input", dtype=dtype, shape=[]),
+            data_loader_backend_module=data_loader_backend_module,
+        )
+
+        scalar = data_loader[0]["input"]
+        assert isinstance(
+            scalar,
+            np.ndarray if data_loader_backend_module == "numpy" else torch.Tensor,
+        )
+        assert scalar.shape == tuple()
+
+    def test_error_on_unsupported_numpy_type(self):
+        input_meta = TensorMetadata().add("X", dtype=DataType.BFLOAT16, shape=(3,))
+        data_loader = DataLoader()
+        data_loader.input_metadata = input_meta
+
+        with pytest.raises(
+            PolygraphyException,
+            match="Please use a custom data loader to provide inputs.",
+        ):
+            data_loader[0]
+
+    def test_bf16_supported_torch(self):
+        input_meta = TensorMetadata().add("X", dtype=DataType.BFLOAT16, shape=(3,))
+        data_loader = DataLoader(data_loader_backend_module="torch")
+        data_loader.input_metadata = input_meta
+
+        assert util.array.is_torch(data_loader[0]["X"])
+    
+    @pytest.mark.parametrize("name, should_match", [
+        ("inp_*",      [True for _ in range(12)]),
+        ("inp_?",      [False, False, False, *[True for _ in range(9)]]),
+        ("inp_[abc]",  [*[False for _ in range(6)], True, True, True, False, False, False]),
+        ("inp_[!abc]", [False, False, False, True, True, True, False, False, False, True, True, True]),
+    ])
+    def test_input_name_with_wildcards(self, name, should_match):
+        match_case = [
+            "inp_foo", "inp_bar", "inp_123", "inp_1", "inp_s", "inp_k",
+            "inp_a", "inp_b", "inp_c", "inp_d", "inp_e", "inp_f",
+        ]
+        input_meta = TensorMetadata().add(name, dtype=np.float32, shape=(2, 2, 3))
+        data_loader = DataLoader(input_metadata=input_meta)
+        data_loader.input_metadata = TensorMetadata()
+        for case in match_case:
+            data_loader.input_metadata.add(case, dtype=np.float32, shape=(-1, 2, 3))
+
+        res = [data_loader[0][name].shape == (2, 2, 3) for name in data_loader[0]]
+        assert res == should_match 
+
+build_torch = lambda a, **kwargs: util.array.to_torch(np.array(a, **kwargs))
+
+
+@pytest.mark.parametrize("array_type", [np.array, build_torch])
+class TestDataLoaderCache:
+    def test_can_cast_dtype(self, array_type):
+        # Ensure that the data loader can only be used once
+        def load_data():
+            yield {"X": array_type(np.ones((1, 1), dtype=np.float32))}
+
+        cache = DataLoaderCache(load_data())
+
+        fp32_meta = TensorMetadata().add("X", dtype=DataType.FLOAT32, shape=(1, 1))
+        cache.set_input_metadata(fp32_meta)
+        feed_dict = cache[0]
+        assert util.array.dtype(feed_dict["X"]) == DataType.FLOAT32
+
+        fp64_meta = TensorMetadata().add("X", dtype=DataType.FLOAT64, shape=(1, 1))
+        cache.set_input_metadata(fp64_meta)
+        feed_dict = cache[0]
+        assert util.array.dtype(feed_dict["X"]) == DataType.FLOAT64
+
+    # If one input isn't in the cache, we shouldn't give up looking
+    # for other inputs
+    def test_will_not_give_up_on_first_cache_miss(self, array_type):
+        SHAPE = (32, 32)
+
+        DATA = [OrderedDict()]
+        DATA[0]["X"] = array_type(np.zeros(SHAPE, dtype=np.int64))
+        DATA[0]["Y"] = array_type(np.zeros(SHAPE, dtype=np.int64))
+
+        cache = DataLoaderCache(DATA)
+        cache.set_input_metadata(
+            TensorMetadata()
+            .add("X", DataType.INT64, shape=SHAPE)
+            .add("Y", DataType.INT64, SHAPE)
+        )
+
+        # Populate the cache with bad X but good Y.
+        # The data loader cache should fail to coerce X to the right shape and then reload it from the data loader.
+        cache.cache[0] = OrderedDict()
+        cache.cache[0]["X"] = array_type(np.ones((64, 64), dtype=np.int64))
+        cache.cache[0]["Y"] = array_type(np.ones(SHAPE, dtype=np.int64))
+
+        feed_dict = cache[0]
+        # Cache cannot reuse X, so it'll reload - we'll get all 0s from the data loader
+        assert util.array.all(feed_dict["X"] == 0)
+        # Cache can reuse Y, even though it's after X, so we'll get ones from the cache
+        assert util.array.all(feed_dict["Y"] == 1)
+
+    # The cache should ignore extra data generated by the data loader
+    def test_ignores_extra_data(self, array_type):
+        SHAPE = (32, 32)
+
+        DATA = [OrderedDict()]
+        DATA[0]["X"] = array_type(np.zeros(SHAPE, dtype=np.int64))
+        DATA[0]["Y"] = array_type(np.zeros(SHAPE, dtype=np.int64))
+
+        cache = DataLoaderCache(DATA)
+
+        cache.set_input_metadata(TensorMetadata().add("X", DataType.INT64, shape=SHAPE))
+
+        feed_dict = cache[0]
+        assert list(feed_dict.keys()) == ["X"]
+        assert util.array.all(feed_dict["X"] == 0)

tools/Polygraphy/tests/comparator/test_postprocess.py

@@ -0,0 +1,57 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import numpy as np
+import pytest
+from polygraphy import util
+from polygraphy.comparator import PostprocessFunc, IterationResult
+
+build_torch = lambda a, **kwargs: util.array.to_torch(np.array(a, **kwargs))
+
+
+@pytest.mark.parametrize("array_type", [np.array, build_torch])
+class TestTopK:
+    def test_basic(self, array_type):
+        arr = array_type([1, 2, 3, 4, 5], dtype=np.float32)
+        func = PostprocessFunc.top_k(k=3)
+        top_k = func(IterationResult({"x": arr}))
+        assert util.array.equal(top_k["x"], array_type([4, 3, 2]))
+
+    def test_k_can_exceed_array_len(self, array_type):
+        arr = array_type([1, 2, 3, 4, 5], dtype=np.float32)
+        func = PostprocessFunc.top_k(k=10)
+        top_k = func(IterationResult({"x": arr}))
+        assert util.array.equal(top_k["x"], array_type([4, 3, 2, 1, 0]))
+
+    def test_per_output_top_k(self, array_type):
+        arr = array_type([1, 2, 3, 4, 5], dtype=np.float32)
+        func = PostprocessFunc.top_k(k={"": 10, "y": 2})
+        top_k = func(IterationResult({"x": arr, "y": arr}))
+        assert util.array.equal(top_k["x"], array_type([4, 3, 2, 1, 0]))
+        assert util.array.equal(top_k["y"], array_type([4, 3]))
+
+    def test_per_output_top_k_axis(self, array_type):
+        arr = array_type([[5, 6, 5], [6, 5, 6]], dtype=np.float32)
+        func = PostprocessFunc.top_k(k={"": (1, 0), "y": (1, 1)})
+        top_k = func(IterationResult({"x": arr, "y": arr}))
+        assert util.array.equal(top_k["x"], array_type([[1, 0, 1]]))
+        assert util.array.equal(top_k["y"], array_type([[1], [0]]))
+
+    def test_top_k_half(self, array_type):
+        arr = array_type([1, 2, 3, 4, 5], dtype=np.float16)
+        func = PostprocessFunc.top_k(k=3)
+        top_k = func(IterationResult({"x": arr}))
+        assert util.array.equal(top_k["x"], array_type([4, 3, 2]))

tools/Polygraphy/tests/comparator/test_struct.py

@@ -0,0 +1,163 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import contextlib
+
+import numpy as np
+import pytest
+import torch
+
+from polygraphy import config, util
+from polygraphy.comparator import IterationResult, RunResults
+from polygraphy.comparator.struct import LazyArray
+from polygraphy.exception import PolygraphyException
+
+
+def make_outputs():
+    return {"dummy_out": np.zeros((4, 4))}
+
+
+def make_iter_results(runner_name):
+    return [IterationResult(outputs=make_outputs(), runner_name=runner_name)] * 2
+
+
+@pytest.fixture()
+def run_results():
+    results = RunResults()
+    results.append(("runner0", make_iter_results("runner0")))
+    results.append(("runner1", make_iter_results("runner1")))
+    return results
+
+
+class TestRunResults:
+    def test_items(self, run_results):
+        for name, iteration_results in run_results.items():
+            assert isinstance(name, str)
+            assert isinstance(iteration_results, list)
+            for iter_res in iteration_results:
+                assert isinstance(iter_res, IterationResult)
+
+    def test_keys(self, run_results):
+        assert list(run_results.keys()) == ["runner0", "runner1"]
+
+    def test_values(self, run_results):
+        for iteration_results in run_results.values():
+            for iter_res in iteration_results:
+                assert isinstance(iter_res, IterationResult)
+
+    def test_getitem(self, run_results):
+        assert isinstance(run_results["runner0"][0], IterationResult)
+        assert isinstance(run_results[0][1][0], IterationResult)
+        assert run_results[0][1] == run_results["runner0"]
+        assert run_results[1][1] == run_results["runner1"]
+
+    def test_getitem_out_of_bounds(self, run_results):
+        with pytest.raises(IndexError):
+            run_results[2]
+
+        with pytest.raises(PolygraphyException, match="does not exist in this"):
+            run_results["runner2"]
+
+    def test_setitem(self, run_results):
+        def check_results(results, is_none=False):
+            for iter_res in results["runner1"]:
+                if is_none:
+                    assert not iter_res
+                    assert iter_res.runner_name == "custom_runner"
+                else:
+                    assert iter_res
+                    assert iter_res.runner_name
+
+        check_results(run_results)
+
+        iter_results = [IterationResult(outputs=None, runner_name=None)]
+        run_results["runner1"] = iter_results
+
+        check_results(run_results, is_none=True)
+
+    def test_setitem_out_of_bounds(self, run_results):
+        iter_results = [IterationResult(outputs=None, runner_name="new")]
+        run_results["runner2"] = iter_results
+
+        assert len(run_results) == 3
+        assert run_results["runner2"][0].runner_name == "new"
+
+    def test_contains(self, run_results):
+        assert "runner0" in run_results
+        assert "runner1" in run_results
+        assert "runner3" not in run_results
+
+    def test_add_new(self):
+        results = RunResults()
+        results.add([make_outputs()], runner_name="custom")
+
+        iter_results = results["custom"]
+        assert len(iter_results) == 1
+        assert all(
+            isinstance(iter_result, IterationResult) for iter_result in iter_results
+        )
+
+    def test_add_new_default_name(self):
+        results = RunResults()
+        results.add([make_outputs()])
+
+        name = results[0][0]
+        iter_results = results[name]
+        assert len(iter_results) == 1
+        assert all(
+            isinstance(iter_result, IterationResult) for iter_result in iter_results
+        )
+
+
+@pytest.mark.parametrize("module", [torch, np])
+class TestLazyArray:
+    @pytest.mark.parametrize("set_threshold", [True, False])
+    def test_unswapped_array(self, set_threshold, module):
+        with contextlib.ExitStack() as stack:
+            if set_threshold:
+
+                def reset_array_swap():
+                    config.ARRAY_SWAP_THRESHOLD_MB = -1
+
+                stack.callback(reset_array_swap)
+
+                config.ARRAY_SWAP_THRESHOLD_MB = 8
+
+            small_shape = (7 * 1024 * 1024,)
+            small_array = module.ones(small_shape, dtype=module.uint8)
+            lazy = LazyArray(small_array)
+            assert util.array.equal(small_array, lazy.arr)
+            assert lazy.tmpfile is None
+
+            assert util.array.equal(small_array, lazy.load())
+
+    def test_swapped_array(self, module):
+        with contextlib.ExitStack() as stack:
+
+            def reset_array_swap():
+                config.ARRAY_SWAP_THRESHOLD_MB = -1
+
+            stack.callback(reset_array_swap)
+
+            config.ARRAY_SWAP_THRESHOLD_MB = 8
+
+            large_shape = (9 * 1024 * 1024,)
+            large_array = module.ones(large_shape, dtype=module.uint8)
+            lazy = LazyArray(large_array)
+            assert lazy.arr is None
+            assert lazy.tmpfile is not None
+
+            assert util.array.equal(large_array, lazy.load())