chore: import upstream snapshot with attribution

tools/Polygraphy/tests/README.md

@@ -0,0 +1,34 @@
+# Tests
+
+The tests directory closely mirrors the structure of the main `polygraphy` directory.
+
+## Adding Tests
+
+For a given submodule, add tests into the corresponding directory under `tests/`.
+
+## Parallel Test Execution
+
+By default, the Polygraphy build system runs tests in parallel. However, some tests
+may not be good candidates for parallel exection - for example, performance tests.
+You can selectively disable parallel execution for these tests using the `pytest.mark.serial`
+marker. The build system will exclude these tests from parallel execution and run them serially instead.
+
+For example:
+```python
+@pytest.mark.serial
+def my_not_parallel_test():
+    ...
+```
+
+## Slow Tests
+
+Some tests are long-running, so we'd prefer to avoid running them during local development.
+You can mark tests with the `pytest.mark.slow` marker so that those tests are only run when
+the `RUN_ALL_TESTS` make option is enabled.
+
+For example:
+```python
+@pytest.mark.slow
+def my_long_running_test():
+    ...
+```

tools/Polygraphy/tests/backend/base/test_loader.py

@@ -0,0 +1,32 @@
+#
+# 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 pytest
+from polygraphy.backend.base import BaseLoader
+from polygraphy.exception import PolygraphyInternalException
+
+
+def test_loader_checks_call_constant_method():
+    class BadLoader(BaseLoader):
+        def __init__(self):
+            self.x = 2
+
+        def call_impl(self):
+            self.x = 3
+
+    with pytest.raises(PolygraphyInternalException):
+        BadLoader()()

tools/Polygraphy/tests/backend/base/test_runner.py

@@ -0,0 +1,29 @@
+#
+# 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.backend.onnxrt import OnnxrtRunner, SessionFromOnnx
+from polygraphy.exception import PolygraphyException
+from tests.models.meta import ONNX_MODELS
+
+
+def test_infer_raises_if_runner_inactive():
+    runner = OnnxrtRunner(SessionFromOnnx(ONNX_MODELS["identity"].loader))
+    feed_dict = {"x": np.ones((1, 1, 2, 2), dtype=np.float32)}
+
+    with pytest.raises(PolygraphyException, match="Must be activated"):
+        runner.infer(feed_dict)

tools/Polygraphy/tests/backend/common/test_loader.py

@@ -0,0 +1,77 @@
+#
+# 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 os
+from textwrap import dedent
+
+import pytest
+import tensorrt as trt
+from polygraphy import util
+from polygraphy.backend.common import InvokeFromScript, invoke_from_script
+from polygraphy.exception import PolygraphyException
+
+
+class TestImporter:
+    @pytest.mark.parametrize("loader", [InvokeFromScript, invoke_from_script])
+    def test_import_from_script(self, loader):
+        script = dedent(
+            """
+            from polygraphy.backend.trt import CreateNetwork
+            from polygraphy import func
+            import tensorrt as trt
+
+            @func.extend(CreateNetwork())
+            def load_network(builder, network):
+                inp = network.add_input("input", dtype=trt.float32, shape=(1, 1))
+                out = network.add_identity(inp).get_output(0)
+                network.mark_output(out)
+            """
+        )
+
+        with util.NamedTemporaryFile("w+", suffix=".py") as f:
+            f.write(script)
+            f.flush()
+            os.fsync(f.fileno())
+
+            if loader == InvokeFromScript:
+                load_network = loader(f.name, "load_network")
+                builder, network = load_network()
+            else:
+                builder, network = loader(f.name, "load_network")
+            with builder, network:
+                assert isinstance(builder, trt.Builder)
+                assert isinstance(network, trt.INetworkDefinition)
+                assert network.num_layers == 1
+                assert network.get_layer(0).type == trt.LayerType.IDENTITY
+
+    def test_import_non_existent(self):
+        script = dedent(
+            """
+            def example():
+                pass
+            """
+        )
+
+        with util.NamedTemporaryFile("w+", suffix=".py") as f:
+            f.write(script)
+            f.flush()
+            os.fsync(f.fileno())
+
+            with pytest.raises(
+                PolygraphyException, match="Could not import symbol: non_existent from"
+            ):
+                invoke_from_script(f.name, "non_existent")

tools/Polygraphy/tests/backend/onnx/test_loader.py

@@ -0,0 +1,380 @@
+#
+# 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 os
+import tempfile
+
+import numpy as np
+import onnx
+import onnx_graphsurgeon as gs
+import pytest
+
+from polygraphy import constants
+from polygraphy.backend.onnx import (
+    ConvertToFp16,
+    FoldConstants,
+    ModifyOutputs,
+    OnnxFromBytes,
+    OnnxFromPath,
+    OnnxFromTfGraph,
+    SaveOnnx,
+    SetUpperBound,
+    extract_subgraph,
+    fold_constants,
+    gs_from_onnx,
+    infer_shapes,
+    onnx_from_path,
+)
+from polygraphy.common import TensorMetadata
+from polygraphy.logger import G_LOGGER
+from tests.helper import is_file_non_empty
+from tests.models.meta import ONNX_MODELS, TF_MODELS
+
+
+class TestLoggerCallbacks:
+    @pytest.mark.parametrize("sev", G_LOGGER.SEVERITY_LETTER_MAPPING.keys())
+    def test_set_severity(self, sev):
+        G_LOGGER.module_severity = sev
+
+
+class TestOnnxFromPath:
+    def test_basic(self):
+        loader = OnnxFromPath(ONNX_MODELS["identity"].path)
+        model = loader()
+        assert isinstance(model, onnx.ModelProto)
+        assert len(model.graph.node) == 1
+
+    @pytest.mark.serial
+    def test_warn_if_impl_methods_called(self, check_warnings_on_loader_impl_methods):
+        check_warnings_on_loader_impl_methods(
+            OnnxFromPath(ONNX_MODELS["identity"].path)
+        )
+
+    def test_external_data(self):
+        model = ONNX_MODELS["ext_weights"]
+        loader = OnnxFromPath(model.path, model.ext_data)
+        assert isinstance(loader(), onnx.ModelProto)
+
+    def test_ignore_external_data(self):
+        model = ONNX_MODELS["ext_weights"]
+        loader = OnnxFromPath(model.path, ignore_external_data=True)
+        onnx_model = loader()
+
+        assert isinstance(onnx_model, onnx.ModelProto)
+        assert all(init.data_location == 1 for init in onnx_model.graph.initializer)
+
+
+class TestOnnxFromBytes:
+    def test_basic(self):
+        loader = OnnxFromBytes(ONNX_MODELS["identity"].loader)
+        model = loader()
+        assert isinstance(model, onnx.ModelProto)
+        assert len(model.graph.node) == 1
+
+
+class TestGsFromOnnx:
+    def test_basic(self):
+        graph = gs_from_onnx(OnnxFromPath(ONNX_MODELS["identity"].path))
+        assert isinstance(graph, gs.Graph)
+
+
+class TestExportOnnxFromTf:
+    def test_no_optimize(self):
+        pytest.importorskip("tensorflow")
+        loader = OnnxFromTfGraph(TF_MODELS["identity"].loader, optimize=False)
+        model = loader()
+
+    def test_opset(self):
+        pytest.importorskip("tensorflow")
+        loader = OnnxFromTfGraph(TF_MODELS["identity"].loader, opset=9)
+        model = loader()
+        assert model.opset_import[0].version == 9
+
+
+class TestModifyOnnx:
+    @pytest.mark.parametrize("copy", [True, False])
+    def test_layerwise(self, copy):
+        original_model = onnx_from_path(ONNX_MODELS["identity_identity"].path)
+        loader = ModifyOutputs(original_model, outputs=constants.MARK_ALL, copy=copy)
+        model = loader()
+        assert len(original_model.graph.output) == 1 or not copy
+        assert len(model.graph.output) == 2
+
+    @pytest.mark.parametrize("output", ["identity_out_0", "identity_out_2"])
+    def test_custom_outputs(self, output):
+        loader = ModifyOutputs(
+            OnnxFromPath(ONNX_MODELS["identity_identity"].path), outputs=[output]
+        )
+        model = loader()
+        assert len(model.graph.output) == 1
+        assert model.graph.output[0].name == output
+
+    def test_exclude_outputs_with_layerwise(self):
+        loader = ModifyOutputs(
+            OnnxFromPath(ONNX_MODELS["identity_identity"].path),
+            outputs=constants.MARK_ALL,
+            exclude_outputs=["identity_out_2"],
+        )
+        model = loader()
+        assert len(model.graph.output) == 1
+        assert model.graph.output[0].name == "identity_out_0"
+
+
+@pytest.mark.parametrize("allow_onnxruntime", [True, False])
+class TestInferShapes:
+    def check_model(self, model):
+        # Find all intermediate tensors to check if they have shapes.
+        tensors = set()
+        for node in model.graph.node:
+            tensors.update(node.output)
+        tensors -= {out.name for out in model.graph.output}
+
+        assert len(model.graph.value_info) >= len(tensors)
+        for val in model.graph.value_info:
+            assert val.type.tensor_type.HasField("shape")
+
+    def test_model(self, allow_onnxruntime):
+        original_model = onnx_from_path(ONNX_MODELS["identity_identity"].path)
+        model = infer_shapes(original_model, allow_onnxruntime=allow_onnxruntime)
+        self.check_model(model)
+
+    def test_path(self, allow_onnxruntime):
+        model = infer_shapes(
+            ONNX_MODELS["identity_identity"].path, allow_onnxruntime=allow_onnxruntime
+        )
+        self.check_model(model)
+
+    @pytest.mark.parametrize("set_data_dir", [True, False])
+    def test_external_data(self, set_data_dir, allow_onnxruntime):
+        model = ONNX_MODELS["ext_weights_same_dir"]
+        model = infer_shapes(
+            model.path,
+            external_data_dir=model.ext_data if set_data_dir else None,
+            allow_onnxruntime=allow_onnxruntime,
+        )
+        self.check_model(model)
+
+    def test_save_to_disk_on_size_threshold(self, allow_onnxruntime):
+        model = onnx_from_path(ONNX_MODELS["const_foldable"].path)
+        model = infer_shapes(
+            model, save_to_disk_threshold_bytes=0, allow_onnxruntime=allow_onnxruntime
+        )
+        self.check_model(model)
+
+
+class TestConvertToFp16:
+    @pytest.mark.parametrize("copy", [True, False])
+    def test_basic(self, copy):
+        # Precondition.
+        original_model = onnx_from_path(ONNX_MODELS["identity_identity"].path)
+        assert original_model.graph.input[0].type.tensor_type.elem_type == onnx.TensorProto.FLOAT or not copy
+
+        # Under test.
+        loader = ConvertToFp16(original_model, copy=copy)
+        model = loader()
+
+        # Postcondition.
+        graph = gs_from_onnx(model)
+        graph.toposort()
+
+        assert graph.inputs[0].dtype == "float32"
+        assert graph.nodes[0].op == "Cast"
+        assert graph.nodes[1].op == "Identity"
+        assert graph.nodes[2].op == "Identity"
+        assert graph.nodes[3].op == "Cast"
+        assert graph.outputs[0].dtype == "float32"
+
+
+class TestFoldConstants:
+    @pytest.mark.parametrize("fold_shapes", [True, False])
+    @pytest.mark.parametrize("partitioning", [None, "basic", "recursive"])
+    @pytest.mark.parametrize("copy", [True, False])
+    @pytest.mark.parametrize("allow_onnxruntime_shape_inference", [True, False])
+    def test_basic(
+        self, partitioning, fold_shapes, copy, allow_onnxruntime_shape_inference
+    ):
+        original_model = onnx_from_path(ONNX_MODELS["const_foldable"].path)
+        loader = FoldConstants(
+            original_model,
+            partitioning=partitioning,
+            fold_shapes=fold_shapes,
+            copy=copy,
+            error_ok=False,
+            allow_onnxruntime_shape_inference=allow_onnxruntime_shape_inference,
+        )
+        model = loader()
+        assert len(original_model.graph.node) != 1 or not copy
+        assert len(model.graph.node) == 1
+
+    @pytest.mark.parametrize(
+        "size_threshold, expect_folding",
+        [
+            (None, True),
+            (0, False),
+            (10 << 20, True),
+            (10 << 20 - 1, False),
+        ],
+    )
+    def test_size_threshold(self, size_threshold, expect_folding):
+        model = onnx_from_path(ONNX_MODELS["constant_fold_bloater"].path)
+        model = fold_constants(model, size_threshold=size_threshold)
+
+        if expect_folding:
+            assert len(model.graph.node) == 0
+        else:
+            assert len(model.graph.node) == 1
+            assert model.graph.node[0].op_type == "Tile"
+
+
+class TestSetUpperBound:
+    @pytest.mark.parametrize("global_upper_bound", [False, True])
+    @pytest.mark.parametrize("specified_upper_bound", [False, True])
+    def test_set_upper_bound(
+        self,
+        global_upper_bound,
+        specified_upper_bound,
+    ):
+        original_model = onnx_from_path(ONNX_MODELS["unbounded_dds"].path)
+        upper_bound_dict = {}
+        if not global_upper_bound and not specified_upper_bound:
+            upper_bound_dict[""] = 1000
+            upper_bound = 1000
+        if global_upper_bound:
+            upper_bound_dict[""] = 2000
+            upper_bound = 2000
+        if specified_upper_bound:
+            upper_bound_dict["cast_out_6"] = 4000
+            upper_bound = 4000
+
+        loader = SetUpperBound(
+            original_model,
+            upper_bounds=upper_bound_dict,
+        )
+
+        model = loader()
+        graph = gs_from_onnx(model)
+
+        # Check if there is a Min operator in the modified model
+        find_min = False
+        for node in graph.nodes:
+            if node.op == "Min":
+                find_min = True
+                # Check if the Min operator's second input is a constant tensor
+                assert isinstance(node.inputs[1], gs.Constant)
+
+                val = node.inputs[1].values
+                # Check if the constant value equals the target upper bound
+                assert val == upper_bound
+        assert find_min
+
+
+class TestSaveOnnx:
+    def test_save_onnx(self):
+        with tempfile.TemporaryDirectory() as outdir:
+            outpath = os.path.join(outdir, "test", "nested")
+            loader = SaveOnnx(OnnxFromPath(ONNX_MODELS["identity"].path), path=outpath)
+            loader()
+            assert is_file_non_empty(outpath)
+
+    def test_external_data(self):
+        with tempfile.NamedTemporaryFile(dir=".") as path, tempfile.NamedTemporaryFile(dir=".") as data:
+            rpath_name = os.path.basename(data.name)
+            model = OnnxFromPath(ONNX_MODELS["const_foldable"].path)
+            loader = SaveOnnx(
+                model, path.name, external_data_path=rpath_name, size_threshold=0
+            )
+            loader()
+            assert is_file_non_empty(path.name)
+            assert is_file_non_empty(data.name)
+
+
+@pytest.fixture()
+def extract_model():
+    input_metadata = TensorMetadata().add("X", dtype=np.float32, shape=(64, 64))
+    output_metadata = TensorMetadata().add(
+        "identity_out_0", dtype=np.float32, shape=None
+    )
+    return (
+        onnx_from_path(ONNX_MODELS["identity_identity"].path),
+        input_metadata,
+        output_metadata,
+    )
+
+
+class TestExtractSubgraph:
+    def check_model(self, model):
+        graph = gs_from_onnx(model)
+        assert len(graph.nodes) == 1
+
+        assert len(graph.inputs) == 1
+        assert graph.inputs[0].name == "X"
+        assert graph.inputs[0].shape is not None
+        assert graph.inputs[0].dtype is not None
+
+        assert len(graph.outputs) == 1
+        assert graph.outputs[0].name == "identity_out_0"
+        assert graph.outputs[0].dtype is not None
+
+    def test_extract_onnx_model(self, extract_model):
+        original_model, input_meta, output_meta = extract_model
+        model = extract_subgraph(original_model, input_meta, output_meta)
+
+        assert original_model.graph.output[0].name == "identity_out_2"
+        self.check_model(model)
+
+    def test_extract_onnx_model_no_input_meta(self, extract_model):
+        model, _, output_meta = extract_model
+        model = extract_subgraph(model, output_metadata=output_meta)
+        self.check_model(model)
+
+    def test_extract_onnx_model_no_output_meta(self, extract_model):
+        model, input_meta, _ = extract_model
+        model = extract_subgraph(model, input_metadata=input_meta)
+        assert model.graph.output[0].name == "identity_out_2"
+
+    def test_extract_onnx_gs_graph(self, extract_model):
+        model, input_meta, output_meta = extract_model
+        graph = gs_from_onnx(model)
+        subgraph = extract_subgraph(graph, input_meta, output_meta)
+        # Make sure original graph isn't modified.
+        assert len(graph.nodes) == 2
+
+        assert isinstance(subgraph, gs.Graph)
+        assert len(subgraph.nodes) == 1
+
+        assert len(subgraph.inputs) == 1
+        assert subgraph.inputs[0].name == "X"
+
+        assert len(subgraph.outputs) == 1
+        assert subgraph.outputs[0].name == "identity_out_0"
+
+    def test_extract_passes_no_input_shape(self, extract_model):
+        model, input_meta, output_meta = extract_model
+        input_meta["X"].shape = None
+        model = extract_subgraph(model, input_meta, output_meta)
+        self.check_model(model)
+
+    def test_extract_passes_no_input_dtype(self, extract_model):
+        model, input_meta, output_meta = extract_model
+        input_meta["X"].dtype = None
+        model = extract_subgraph(model, input_meta, output_meta)
+        self.check_model(model)
+
+    def test_extract_passes_no_output_shape(self, extract_model):
+        model, input_meta, output_meta = extract_model
+        output_meta["identity_out_0"].shape = None
+        model = extract_subgraph(model, input_meta, output_meta)
+        self.check_model(model)

tools/Polygraphy/tests/backend/onnx/test_util.py

@@ -0,0 +1,33 @@
+#
+# 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 polygraphy.backend.onnx import onnx_from_path, gs_from_onnx
+from polygraphy.backend.onnx import util as onnx_util
+from tests.models.meta import ONNX_MODELS
+
+
+def test_get_num_nodes():
+    model = onnx_from_path(ONNX_MODELS["scan"].path)
+    assert onnx_util.get_num_nodes(model) == 3  # Should count subgraph nodes.
+
+
+def test_get_unbounded_dds_tensors():
+    model = onnx_from_path(ONNX_MODELS["unbounded_dds"].path)
+    graph = gs_from_onnx(model)
+    tensors = onnx_util.get_unbounded_dds_tensors(graph)
+    assert len(tensors) == 1
+    assert tensors[0].name == "cast_out_6"

tools/Polygraphy/tests/backend/onnxrt/test_loader.py

@@ -0,0 +1,88 @@
+#
+# 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 polygraphy.backend.onnxrt import SessionFromOnnx
+from polygraphy.exception import PolygraphyException
+from tests.models.meta import ONNX_MODELS
+
+import onnxruntime as onnxrt
+import pytest
+
+
+class TestSessionFromOnnx:
+    def test_defaults(self):
+        model = ONNX_MODELS["identity"]
+        loader = SessionFromOnnx(model.loader)
+        sess = loader()
+
+        assert sess
+        assert isinstance(sess, onnxrt.InferenceSession)
+        assert sess.get_providers() == ["CPUExecutionProvider"]
+
+    @pytest.mark.parametrize(
+        "providers,expected",
+        [
+            (["cpu"], ["CPUExecutionProvider"]),
+            (["CPU"], ["CPUExecutionProvider"]),
+        ],
+    )
+    def test_provider_matching(self, providers, expected):
+        model = ONNX_MODELS["identity"]
+        loader = SessionFromOnnx(model.loader, providers=providers)
+        sess = loader()
+
+        assert sess
+        assert isinstance(sess, onnxrt.InferenceSession)
+        assert sess.get_providers() == expected
+
+    @pytest.mark.skipif(
+        "TensorrtExecutionProvider" not in onnxrt.get_available_providers(),
+        reason="Skip test if TensorrtExecutionProvider is not available",
+    )
+    @pytest.mark.parametrize(
+        "providers,expected_dict",
+        [
+            # Searches for 'tensorrt' as the execution provider's name
+            (["tensorrt", "cpu"], {"TensorrtExecutionProvider": {}, "CPUExecutionProvider": {}}),
+            # Searches for the execution provider's name if the item is a tuple in the format (EP name, EP options)
+            (
+                    [("TensorrtExecutionProvider", {"trt_op_types_to_exclude": "Add"}), "CPUExecutionProvider"],
+                    {"TensorrtExecutionProvider": {"trt_op_types_to_exclude": "Add"}, "CPUExecutionProvider": {}}
+            ),
+        ],
+    )
+    def test_provider_with_options(self, providers, expected_dict):
+        model = ONNX_MODELS["identity"]
+        loader = SessionFromOnnx(model.loader, providers=providers)
+        sess = loader()
+
+        assert sess
+        assert isinstance(sess, onnxrt.InferenceSession)
+        assert sess.get_providers() == list(expected_dict.keys())
+
+        provider_options = sess.get_provider_options()
+        for k, v in provider_options.items():
+            if expected_dict.get(k, None):
+                assert set(expected_dict[k].items()).issubset(v.items())
+
+    def test_invalid_providers_raise_errors(self):
+        model = ONNX_MODELS["identity"]
+        loader = SessionFromOnnx(model.loader, providers=["cpu", "not_a_real_provider"])
+        with pytest.raises(
+            PolygraphyException,
+            match="Could not find specified ONNX-Runtime execution provider",
+        ):
+            loader()

tools/Polygraphy/tests/backend/onnxrt/test_runner.py

@@ -0,0 +1,105 @@
+#
+# 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
+import torch
+
+from polygraphy.backend.onnxrt import OnnxrtRunner, SessionFromOnnx
+from polygraphy.exception import PolygraphyException
+from polygraphy.logger import G_LOGGER
+from tests.models.meta import ONNX_MODELS
+
+
+class TestLoggerCallbacks:
+    @pytest.mark.parametrize("sev", G_LOGGER.SEVERITY_LETTER_MAPPING.keys())
+    def test_set_severity(self, sev):
+        G_LOGGER.module_severity = sev
+
+
+class TestOnnxrtRunner:
+    def test_can_name_runner(self):
+        NAME = "runner"
+        runner = OnnxrtRunner(None, name=NAME)
+        assert runner.name == NAME
+
+    def test_basic(self):
+        model = ONNX_MODELS["identity"]
+        with OnnxrtRunner(SessionFromOnnx(model.loader)) as runner:
+            assert runner.is_active
+            model.check_runner(runner)
+            assert runner.last_inference_time() is not None
+        assert not runner.is_active
+
+    def test_torch_tensors(self):
+        model = ONNX_MODELS["identity"]
+        with OnnxrtRunner(SessionFromOnnx(model.loader)) as runner:
+            arr = torch.ones((1, 1, 2, 2), dtype=torch.float32)
+            outputs = runner.infer({"x": arr})
+
+            assert isinstance(outputs["y"], torch.Tensor)
+            assert torch.equal(outputs["y"], arr)
+
+    @pytest.mark.serial
+    def test_warn_if_impl_methods_called(self, check_warnings_on_runner_impl_methods):
+        model = ONNX_MODELS["identity"]
+        runner = OnnxrtRunner(SessionFromOnnx(model.loader))
+        check_warnings_on_runner_impl_methods(runner)
+
+    def test_shape_output(self):
+        model = ONNX_MODELS["reshape"]
+        with OnnxrtRunner(SessionFromOnnx(model.loader)) as runner:
+            model.check_runner(runner)
+
+    def test_dim_param_preserved(self):
+        model = ONNX_MODELS["dim_param"]
+        with OnnxrtRunner(SessionFromOnnx(model.loader)) as runner:
+            input_meta = runner.get_input_metadata(use_numpy_dtypes=False)
+            # In Polygraphy, we only use None to indicate a dynamic input dimension - not strings.
+            assert len(input_meta) == 1
+            for _, (_, shape) in input_meta.items():
+                assert shape == ["dim0", 16, 128]
+
+    @pytest.mark.parametrize(
+        "names, err",
+        [
+            (["fake-input", "x"], "Extra inputs in"),
+            (["fake-input"], "The following inputs were not found"),
+            ([], "The following inputs were not found"),
+        ],
+    )
+    def test_error_on_wrong_name_feed_dict(self, names, err):
+        model = ONNX_MODELS["identity"]
+        with OnnxrtRunner(SessionFromOnnx(model.loader)) as runner:
+            with pytest.raises(PolygraphyException, match=err):
+                runner.infer(
+                    {
+                        name: np.ones(shape=(1, 1, 2, 2), dtype=np.float32)
+                        for name in names
+                    }
+                )
+
+    def test_error_on_wrong_dtype_feed_dict(self):
+        model = ONNX_MODELS["identity"]
+        with OnnxrtRunner(SessionFromOnnx(model.loader)) as runner:
+            with pytest.raises(PolygraphyException, match="unexpected dtype."):
+                runner.infer({"x": np.ones(shape=(1, 1, 2, 2), dtype=np.int32)})
+
+    def test_error_on_wrong_shape_feed_dict(self):
+        model = ONNX_MODELS["identity"]
+        with OnnxrtRunner(SessionFromOnnx(model.loader)) as runner:
+            with pytest.raises(PolygraphyException, match="incompatible shape."):
+                runner.infer({"x": np.ones(shape=(1, 1, 3, 2), dtype=np.float32)})

tools/Polygraphy/tests/backend/pluginref/test_runner.py

@@ -0,0 +1,100 @@
+#
+# 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.backend.onnx import GsFromOnnx, OnnxFromPath
+from polygraphy.backend.pluginref import PluginRefRunner
+from polygraphy.exception import PolygraphyException
+from polygraphy.logger import G_LOGGER
+from tests.models.meta import ONNX_MODELS
+
+
+class TestLoggerCallbacks:
+    @pytest.mark.parametrize("sev", G_LOGGER.SEVERITY_LETTER_MAPPING.keys())
+    def test_set_severity(self, sev):
+        G_LOGGER.module_severity = sev
+
+
+class TestPluginRefRunner:
+    def test_can_name_runner(self):
+        NAME = "runner"
+        runner = PluginRefRunner(None, name=NAME)
+        assert runner.name == NAME
+
+    def test_basic(self):
+        model = ONNX_MODELS["identity"]
+        with PluginRefRunner(GsFromOnnx(OnnxFromPath(model.path))) as runner:
+            assert runner.is_active
+            model.check_runner(runner)
+        assert not runner.is_active
+
+    @pytest.mark.serial
+    def test_warn_if_impl_methods_called(self, check_warnings_on_runner_impl_methods):
+        model = ONNX_MODELS["identity"]
+        runner = PluginRefRunner(GsFromOnnx(OnnxFromPath(model.path)))
+        check_warnings_on_runner_impl_methods(runner)
+
+    def test_works_on_multiple_nodes(self):
+        model = ONNX_MODELS["identity_identity"]
+        with PluginRefRunner(GsFromOnnx(OnnxFromPath(model.path))) as runner:
+            model.check_runner(runner)
+
+    def test_fail_on_unsupported_node(self):
+        model = ONNX_MODELS["and"]
+        with PluginRefRunner(GsFromOnnx(OnnxFromPath(model.path))) as runner:
+            with pytest.raises(
+                PolygraphyException,
+                match="does not have a reference implementation registered!",
+            ):
+                runner.infer(
+                    {
+                        "x": np.ones(shape=(3, 4), dtype=bool),
+                        "y": np.ones(shape=(3, 4), dtype=bool),
+                    }
+                )
+
+    @pytest.mark.parametrize(
+        "names, err",
+        [
+            (["fake-input", "x"], "Extra inputs in"),
+            (["fake-input"], "The following inputs were not found"),
+            ([], "The following inputs were not found"),
+        ],
+    )
+    def test_error_on_wrong_name_feed_dict(self, names, err):
+        model = ONNX_MODELS["identity"]
+        with PluginRefRunner(GsFromOnnx(OnnxFromPath(model.path))) as runner:
+            with pytest.raises(PolygraphyException, match=err):
+                runner.infer(
+                    {
+                        name: np.ones(shape=(1, 1, 2, 2), dtype=np.float32)
+                        for name in names
+                    }
+                )
+
+    def test_error_on_wrong_dtype_feed_dict(self):
+        model = ONNX_MODELS["identity"]
+        with PluginRefRunner(GsFromOnnx(OnnxFromPath(model.path))) as runner:
+            with pytest.raises(PolygraphyException, match="unexpected dtype."):
+                runner.infer({"x": np.ones(shape=(1, 1, 2, 2), dtype=np.int32)})
+
+    def test_error_on_wrong_shape_feed_dict(self):
+        model = ONNX_MODELS["identity"]
+        with PluginRefRunner(GsFromOnnx(OnnxFromPath(model.path))) as runner:
+            with pytest.raises(PolygraphyException, match="incompatible shape."):
+                runner.infer({"x": np.ones(shape=(1, 1, 3, 2), dtype=np.float32)})

tools/Polygraphy/tests/backend/tf/test_loader.py

@@ -0,0 +1,81 @@
+#
+# 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 os
+import tempfile
+
+import pytest
+from polygraphy import constants, util
+from polygraphy.backend.tf import (
+    GraphFromFrozen,
+    ModifyGraphOutputs,
+    SaveGraph,
+    graph_from_frozen,
+)
+from polygraphy.logger import G_LOGGER
+from tests.helper import is_file_non_empty
+from tests.models.meta import TF_MODELS
+
+tf = pytest.importorskip("tensorflow")
+
+
+class TestLoggerCallbacks:
+    @pytest.mark.parametrize("sev", G_LOGGER.SEVERITY_LETTER_MAPPING.keys())
+    def test_set_severity(self, sev):
+        G_LOGGER.module_severity = sev
+
+
+class TestFrozenGraphLoader:
+    def test_load_graph(self):
+        with tf.compat.v1.Graph().as_default() as graph:
+            inp = tf.placeholder(shape=(1, 1, 1, 1), dtype=tf.float32)
+            out = tf.identity(inp)
+
+        graph, outputs = graph_from_frozen(graph)
+        assert graph
+        assert outputs
+
+    def test_load_pb(self):
+        tf_loader = GraphFromFrozen(TF_MODELS["identity"].path)
+        tf_loader()
+
+
+class TestModifyGraph:
+    def test_layerwise(self):
+        load_frozen = GraphFromFrozen(TF_MODELS["identity"].path)
+        modify_tf = ModifyGraphOutputs(load_frozen, outputs=constants.MARK_ALL)
+
+        graph, outputs = modify_tf()
+        assert graph
+        assert outputs
+
+
+class TestSaveGraph:
+    def test_save_pb(self):
+        with util.NamedTemporaryFile() as outpath:
+            tf_loader = SaveGraph(
+                GraphFromFrozen(TF_MODELS["identity"].path), path=outpath.name
+            )
+            tf_loader()
+            assert is_file_non_empty(outpath.name)
+
+    def test_save_tensorboard(self):
+        with tempfile.TemporaryDirectory() as outdir:
+            tf_loader = SaveGraph(
+                GraphFromFrozen(TF_MODELS["identity"].path), tensorboard_dir=outdir
+            )
+            tf_loader()
+            assert os.path.exists(tf_loader.tensorboard_dir)

tools/Polygraphy/tests/backend/tf/test_runner.py

@@ -0,0 +1,93 @@
+#
+# 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.backend.tf import SessionFromGraph, TfRunner
+from polygraphy.exception import PolygraphyException
+from tests.helper import is_file_non_empty
+from tests.models.meta import TF_MODELS
+
+pytest.importorskip("tensorflow")
+
+
+class TestTfRunner:
+    def test_can_name_runner(self):
+        NAME = "runner"
+        runner = TfRunner(None, name=NAME)
+        assert runner.name == NAME
+
+    def test_basic(self):
+        model = TF_MODELS["identity"]
+        with TfRunner(SessionFromGraph(model.loader)) as runner:
+            assert runner.is_active
+            model.check_runner(runner)
+            assert runner.last_inference_time() is not None
+        assert not runner.is_active
+
+    @pytest.mark.serial
+    def test_warn_if_impl_methods_called(self, check_warnings_on_runner_impl_methods):
+        model = TF_MODELS["identity"]
+        runner = TfRunner(SessionFromGraph(model.loader))
+        check_warnings_on_runner_impl_methods(runner)
+
+    @pytest.mark.skip(reason="Non-trivial to set up - requires CUPTI")
+    def test_save_timeline(self):
+        model = TF_MODELS["identity"]
+        with util.NamedTemporaryFile() as outpath:
+            with TfRunner(
+                SessionFromGraph(model.loader),
+                allow_growth=True,
+                save_timeline=outpath.name,
+            ) as runner:
+                model.check_runner(runner)
+                assert is_file_non_empty(outpath.name)
+
+    @pytest.mark.parametrize(
+        "names, err",
+        [
+            (["fake-input", "Input:0"], "Extra inputs in"),
+            (["fake-input"], "The following inputs were not found"),
+            ([], "The following inputs were not found"),
+        ],
+    )
+    def test_error_on_wrong_name_feed_dict(self, names, err):
+        model = TF_MODELS["identity"]
+        with TfRunner(SessionFromGraph(model.loader)) as runner:
+            with pytest.raises(PolygraphyException, match=err):
+                runner.infer(
+                    {
+                        name: np.ones(shape=(1, 15, 25, 30), dtype=np.float32)
+                        for name in names
+                    }
+                )
+
+    def test_error_on_wrong_dtype_feed_dict(self):
+        model = TF_MODELS["identity"]
+        with TfRunner(SessionFromGraph(model.loader)) as runner:
+            with pytest.raises(PolygraphyException, match="unexpected dtype."):
+                runner.infer(
+                    {"Input:0": np.ones(shape=(1, 15, 25, 30), dtype=np.int32)}
+                )
+
+    def test_error_on_wrong_shape_feed_dict(self):
+        model = TF_MODELS["identity"]
+        with TfRunner(SessionFromGraph(model.loader)) as runner:
+            with pytest.raises(PolygraphyException, match="incompatible shape."):
+                runner.infer(
+                    {"Input:0": np.ones(shape=(1, 1, 25, 30), dtype=np.float32)}
+                )

tools/Polygraphy/tests/backend/trt/test_algorithm_selector.py

@@ -0,0 +1,340 @@
+#
+# 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 namedtuple
+
+import pytest
+import tensorrt as trt
+from polygraphy import config, util
+from polygraphy.backend.trt import (
+    Algorithm,
+    TacticRecorder,
+    TacticReplayData,
+    TacticReplayer,
+    TensorInfo,
+)
+from polygraphy.exception import PolygraphyException
+
+# Skip all tests in this file if TensorRT-RTX is enabled
+if config.USE_TENSORRT_RTX:
+    pytest.skip("Algorithm selector tests are not compatible with TensorRT-RTX", allow_module_level=True)
+
+
+FakeAlgorithmContext = namedtuple(
+    "FakeAlgorithmContext", ["name", "num_inputs", "num_outputs"]
+)
+FakeAlgorithm = namedtuple("FakeAlgorithm", ["algorithm_variant", "io_info"])
+FakeAlgorithm.get_algorithm_io_info = lambda this, index: this.io_info[index]
+
+FakeAlgorithmVariant = namedtuple("FakeAlgorithmVariant", ["implementation", "tactic"])
+
+
+def fake_context(name):
+    return FakeAlgorithmContext(name=name, num_inputs=1, num_outputs=1)
+
+
+def make_tensor_info(
+    dtype=trt.float32,
+    strides=(1, 2, 3),
+    vectorized_dim=-1,
+    components_per_element=1,
+):
+    return TensorInfo(dtype, strides, vectorized_dim, components_per_element)
+
+
+def fake_algo(implementation=6, tactic=0, io=None):
+    io_info = [make_tensor_info()] * 2
+    if io:
+        io_info = []
+        for dtype, strides in io:
+            io_info.append(
+                TensorInfo(
+                    dtype=dtype,
+                    strides=strides,
+                    vectorized_dim=-1,
+                    components_per_element=1,
+                )
+            )
+
+    trt_algo = FakeAlgorithm(
+        algorithm_variant=FakeAlgorithmVariant(implementation, tactic), io_info=io_info
+    )
+    return trt_algo
+
+
+class TestTensorInfo:
+    @pytest.mark.parametrize(
+        "left, right, expected",
+        [
+            (
+                TensorInfo(trt.float32, (1, 2, 3), -1, 1),
+                TensorInfo(trt.float32, (1, 2, 3), -1, 1),
+                True,
+            ),
+            # Different data type
+            (
+                TensorInfo(trt.float32, (1, 2, 3), -1, 1),
+                TensorInfo(trt.float16, (1, 2, 3), -1, 1),
+                False,
+            ),
+            # Different vectotrization
+            (
+                TensorInfo(trt.float32, (1, 2, 3), -1, 1),
+                TensorInfo(trt.float32, (1, 2, 3), 0, 2),
+                False,
+            ),
+        ],
+    )
+    def test_equality(self, left, right, expected):
+        assert (left == right) == expected
+
+
+class TestAlgorithm:
+    @pytest.mark.parametrize(
+        "left, right, expected",
+        [
+            (
+                Algorithm(
+                    6,
+                    1,
+                    inputs=[make_tensor_info(trt.float32)],
+                    outputs=[make_tensor_info(trt.float32)],
+                ),
+                Algorithm(
+                    6,
+                    1,
+                    inputs=[make_tensor_info(trt.float32)],
+                    outputs=[make_tensor_info(trt.float32)],
+                ),
+                True,
+            ),  # Same
+            (
+                Algorithm(
+                    6,
+                    1,
+                    inputs=[make_tensor_info(trt.float32)],
+                    outputs=[make_tensor_info(trt.float32)],
+                ),
+                Algorithm(
+                    7,
+                    1,
+                    inputs=[make_tensor_info(trt.float32)],
+                    outputs=[make_tensor_info(trt.float32)],
+                ),
+                False,
+            ),  # Different implementation
+            (
+                Algorithm(
+                    6,
+                    2,
+                    inputs=[make_tensor_info(trt.float32)],
+                    outputs=[make_tensor_info(trt.float32)],
+                ),
+                Algorithm(
+                    6,
+                    1,
+                    inputs=[make_tensor_info(trt.float32)],
+                    outputs=[make_tensor_info(trt.float32)],
+                ),
+                False,
+            ),  # Different tactic
+            (
+                Algorithm(
+                    6,
+                    1,
+                    inputs=[make_tensor_info(trt.float32)],
+                    outputs=[make_tensor_info(trt.float32)],
+                ),
+                Algorithm(
+                    6,
+                    1,
+                    inputs=[make_tensor_info(trt.int8)],
+                    outputs=[make_tensor_info(trt.float32)],
+                ),
+                False,
+            ),  # Different input data type
+            (
+                Algorithm(
+                    6,
+                    1,
+                    inputs=[make_tensor_info(trt.float32)],
+                    outputs=[make_tensor_info(trt.float32)],
+                ),
+                Algorithm(
+                    6,
+                    1,
+                    inputs=[make_tensor_info(trt.float32)],
+                    outputs=[make_tensor_info(trt.int8)],
+                ),
+                False,
+            ),  # Different output data type
+            (
+                Algorithm(
+                    6,
+                    1,
+                    inputs=[make_tensor_info(trt.float32)] * 2,
+                    outputs=[make_tensor_info(trt.float32)],
+                ),
+                Algorithm(
+                    6,
+                    1,
+                    inputs=[make_tensor_info(trt.float32)],
+                    outputs=[make_tensor_info(trt.float32)],
+                ),
+                False,
+            ),  # Different number of inputs
+            (
+                Algorithm(
+                    6,
+                    1,
+                    inputs=[make_tensor_info(trt.float32)],
+                    outputs=[make_tensor_info(trt.float32)] * 2,
+                ),
+                Algorithm(
+                    6,
+                    1,
+                    inputs=[make_tensor_info(trt.float32)],
+                    outputs=[make_tensor_info(trt.float32)],
+                ),
+                False,
+            ),  # Different number of outputs
+        ],
+    )
+    def test_equality(self, left, right, expected):
+        assert (left == right) == expected
+
+
+@pytest.fixture(params=[True, False], ids=["path", "object"])
+def replay(request):
+    """
+    Returns:
+        Tuple[FakeAlgorithmContext, Algorithm, FakeAlgorithm,
+              Union[str, TacticReplayData], Union[str, TacticReplayData]]:
+                This fixture returns 5 things:
+                1. A fake TensorRT algorithm context
+                2. A Polygraphy Algorithm instance
+                3. A fake TensorRT algorithm (with the same information as (2))
+                4. An input tactic replay data, populated with the Polygraphy Algorithm (2), either
+                    as a ``TacticReplayData`` instance, or a path.
+                5. An output tactic replay data, empty, either as a ``TacticReplayData`` instance, or
+                    a path.
+    """
+    jsonify = request.param
+
+    name = "node_of_y"
+    context = fake_context(name)
+
+    trt_algo = fake_algo()
+    poly_algo = Algorithm.from_trt(context, trt_algo)
+
+    in_replay_data = TacticReplayData().add(name, poly_algo)
+    out_replay_data = TacticReplayData()
+    if jsonify:
+        inpath = util.NamedTemporaryFile("w")
+        in_replay_data.save(inpath.name)
+        in_replay_data = inpath.name
+
+        outpath = util.NamedTemporaryFile("r")
+        out_replay_data = outpath.name
+
+    yield context, poly_algo, trt_algo, in_replay_data, out_replay_data
+
+
+class TestReplayer:
+    def test_basic(self, replay):
+        context, _, algo, replay_data, _ = replay
+        replayer = TacticReplayer(replay_data)
+        selected = replayer.select_algorithms(
+            context, [fake_algo(implementation=2), algo, fake_algo(tactic=1)]
+        )
+        assert selected == [1]
+
+    def test_new_layer_falls_back(self, replay):
+        _, _, _, replay_data, _ = replay
+        replayer = TacticReplayer(replay_data)
+        selected = replayer.select_algorithms(
+            fake_context(name="new_layer"),
+            [fake_algo(2, 1), fake_algo(3, 4), fake_algo(5, 6)],
+        )
+        assert selected == [0, 1, 2]
+
+    def test_missing_algo_fails(self, replay):
+        context, _, _, replay_data, _ = replay
+        replayer = TacticReplayer(replay_data)
+        with pytest.raises(
+            PolygraphyException, match="was not provided by TensorRT as a choice"
+        ):
+            assert replayer.select_algorithms(context, [fake_algo(2, 1)]) == [0]
+
+    @pytest.mark.parametrize(
+        "algo",
+        [
+            fake_algo(2),
+            fake_algo(tactic=2),
+            fake_algo(
+                io=[
+                    (trt.float32, (1, 2)),
+                    (trt.float32, (1, 2)),
+                ]
+            ),
+            fake_algo(
+                io=[
+                    (trt.int8, (1, 2)),
+                    (trt.float32, (1, 2)),
+                ]
+            ),
+            fake_algo(
+                io=[
+                    (trt.float32, (1, 2)),
+                    (trt.float32, (1, 2)),
+                ]
+            ),
+            fake_algo(
+                io=[
+                    (trt.float32, (1, 2)),
+                    (trt.int32, (1, 2)),
+                ]
+            ),
+        ],
+    )
+    def test_different_algo_fails(self, replay, algo):
+        context, _, _, replay_data, _ = replay
+        replayer = TacticReplayer(replay_data)
+        with pytest.raises(
+            PolygraphyException, match="was not provided by TensorRT as a choice"
+        ):
+            assert replayer.select_algorithms(context, [algo]) == [0]
+
+    def test_fails_if_wrong_selected(self, replay):
+        context, _, _, replay_data, _ = replay
+        replayer = TacticReplayer(replay_data)
+        # We should be able to check tactics even if we're not recording them.
+        with pytest.raises(
+            PolygraphyException, match="TensorRT selected a tactic different"
+        ):
+            replayer.report_algorithms([context], [fake_algo(implementation=9)])
+
+
+class TestRecorder:
+    def test_basic(self, replay):
+        context, poly_algo, algo, _, replay_data = replay
+        assert isinstance(replay_data, str) or not replay_data
+        replayer = TacticRecorder(replay_data)
+
+        replayer.report_algorithms([context], [algo])
+        assert len(replayer.data) == 1
+        assert replayer.data[context.name] == poly_algo

tools/Polygraphy/tests/backend/trt/test_calibrator.py

@@ -0,0 +1,330 @@
+#
+# 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
+import tensorrt as trt
+import torch
+
+from polygraphy import config, cuda, util
+from polygraphy.backend.trt import (
+    Calibrator,
+    CreateConfig,
+    Profile,
+    engine_from_network,
+    get_trt_logger,
+    network_from_onnx_bytes,
+)
+from polygraphy.common import TensorMetadata
+from polygraphy.comparator import DataLoader
+from polygraphy.datatype import DataType
+from polygraphy.exception import PolygraphyException
+from tests.helper import get_file_size, is_file_non_empty
+from tests.models.meta import ONNX_MODELS
+
+# Skip all tests in this file if TensorRT-RTX is enabled
+if config.USE_TENSORRT_RTX:
+    pytest.skip("Calibrator tests are not compatible with TensorRT-RTX", allow_module_level=True)
+
+
+@pytest.fixture(scope="session")
+def identity_builder_network():
+    builder, network, parser = network_from_onnx_bytes(ONNX_MODELS["identity"].loader)
+    with builder, network, parser:
+        yield builder, network
+
+
+@pytest.fixture(scope="session")
+def dynamic_identity_builder_network():
+    builder, network, parser = network_from_onnx_bytes(ONNX_MODELS["dynamic_identity"].loader)
+    with builder, network, parser:
+        yield builder, network
+
+
+@pytest.fixture(scope="session")
+def multi_input_builder_network():
+    builder, network, parser = network_from_onnx_bytes(ONNX_MODELS["reducable"].loader)
+    with builder, network, parser:
+        yield builder, network
+
+
+def generate_data(num_batches):
+    for item in [np.ones((1, 1, 2, 2), dtype=np.float32)] * num_batches:
+        yield {"x": item}
+
+
+class TestCalibrator:
+    def check_calibrator_cleanup(self, calibrator):
+        # Calibrator buffers should be freed after the build
+        assert all([buf.allocated_nbytes == 0 for buf in calibrator.device_buffers.values()])
+
+    @pytest.mark.parametrize(
+        "BaseClass",
+        [
+            trt.IInt8Calibrator,
+            trt.IInt8LegacyCalibrator,
+            trt.IInt8EntropyCalibrator,
+            trt.IInt8EntropyCalibrator2,
+            trt.IInt8MinMaxCalibrator,
+        ],
+    )
+    def test_calibrator_basic(self, identity_builder_network, BaseClass):
+        builder, network = identity_builder_network
+        NUM_BATCHES = 2
+
+        data = [{"x": np.ones((1, 1, 2, 2), dtype=np.float32)}] * NUM_BATCHES
+        calibrator = Calibrator(data, BaseClass=BaseClass)
+
+        create_config = CreateConfig(int8=True, calibrator=calibrator)
+        with engine_from_network((builder, network), create_config):
+            assert calibrator.num_batches == NUM_BATCHES
+        self.check_calibrator_cleanup(calibrator)
+
+    def test_host_data_copied_to_device(self):
+        with Calibrator(generate_data(1)) as calibrator:
+            [ptr] = calibrator.get_batch(names=["x"])
+            v = cuda.DeviceView(ptr, shape=(1, 1, 2, 2), dtype=np.float32)
+            arr = v.numpy()
+            assert arr.shape == (1, 1, 2, 2)
+            assert np.all(arr == 1)
+        self.check_calibrator_cleanup(calibrator)
+
+    def test_calibrator_data_and_ordering_correct(self):
+        def generate_multidata(num_batches):
+            for _ in range(num_batches):
+                shape = (4, 5)
+                yield {
+                    "x0": np.zeros(shape, dtype=np.float32),
+                    "x1": cuda.DeviceArray(shape=shape, dtype=np.float32).copy_from(np.ones(shape, dtype=np.float32)),
+                    "x2": cuda.DeviceArray(shape=shape, dtype=np.float32)
+                    .copy_from(np.ones(shape, dtype=np.float32) * 2)
+                    .ptr,
+                }
+
+        NUM_BATCHES = 2
+        with Calibrator(generate_multidata(NUM_BATCHES)) as calibrator:
+            for _ in range(NUM_BATCHES):
+                ptrs = calibrator.get_batch(names=["x0", "x1", "x2"])
+                for index, ptr in enumerate(ptrs):
+                    v = cuda.DeviceView(ptr, shape=(4, 5), dtype=np.float32)
+                    assert np.all(v.numpy() == index)
+        self.check_calibrator_cleanup(calibrator)
+
+    def test_calibrator_generator_data(self, identity_builder_network):
+        builder, network = identity_builder_network
+        NUM_BATCHES = 2
+
+        calibrator = Calibrator(generate_data(NUM_BATCHES))
+
+        create_config = CreateConfig(int8=True, calibrator=calibrator)
+        with engine_from_network((builder, network), create_config):
+            assert calibrator.num_batches == NUM_BATCHES
+        self.check_calibrator_cleanup(calibrator)
+
+    # We should be able to mix DeviceView with NumPy arrays and PyTorch tensors.
+    @pytest.mark.parametrize("mode", ["array", "view", "pointer", "torch"])
+    def test_calibrator_device_buffers_multiinput(self, multi_input_builder_network, mode):
+        def generate_dev_data(num_batches):
+            with cuda.DeviceArray(shape=(1,), dtype=np.float32) as x:
+                for _ in range(num_batches):
+                    x.copy_from(np.ones((1,), dtype=np.float32))
+                    xdata = {
+                        "array": x,
+                        "view": cuda.DeviceView(x.ptr, x.shape, x.dtype),
+                        "pointer": x.ptr,
+                        "torch": torch.ones((1,), dtype=torch.float32),
+                    }[mode]
+                    yield {"X0": xdata, "Y0": np.zeros((1,), dtype=np.float32)}
+
+        builder, network = multi_input_builder_network
+        NUM_BATCHES = 2
+
+        calibrator = Calibrator(generate_dev_data(NUM_BATCHES))
+
+        create_config = CreateConfig(int8=True, calibrator=calibrator)
+        with engine_from_network((builder, network), create_config):
+            assert calibrator.num_batches == NUM_BATCHES
+        self.check_calibrator_cleanup(calibrator)
+
+    # We want the calibrator to inter-op with TRT APIs seamlessly
+    def test_calibrator_outside_polygraphy(self, identity_builder_network):
+        builder, network = identity_builder_network
+        NUM_BATCHES = 2
+
+        config = builder.create_builder_config()
+        config.set_flag(trt.BuilderFlag.INT8)
+        calibrator = Calibrator(generate_data(NUM_BATCHES))
+        config.int8_calibrator = calibrator
+        runtime = trt.Runtime(get_trt_logger())
+        engine = runtime.deserialize_cuda_engine(builder.build_serialized_network(network, config))
+
+        assert engine
+        self.check_calibrator_cleanup(calibrator)
+
+    def test_calibrator_with_path_name_cache(self, identity_builder_network):
+        builder, network = identity_builder_network
+        data = [{"x": np.ones((1, 1, 2, 2), dtype=np.float32)}]
+
+        with util.NamedTemporaryFile() as cache:
+            calibrator = Calibrator(data, cache=cache.name)
+            create_config = CreateConfig(int8=True, calibrator=calibrator)
+            with engine_from_network((builder, network), create_config):
+                assert is_file_non_empty(cache.name)
+        self.check_calibrator_cleanup(calibrator)
+
+    @pytest.mark.parametrize("mode", ["wb+", "rb", "wb"])
+    def test_calibrator_with_file_object_cache(self, identity_builder_network, mode):
+        builder, network = identity_builder_network
+        data = [{"x": np.ones((1, 1, 2, 2), dtype=np.float32)}]
+
+        with util.NamedTemporaryFile(mode=mode) as cache:
+            calibrator = Calibrator(data, cache=cache)
+            create_config = CreateConfig(int8=True, calibrator=calibrator)
+            with engine_from_network((builder, network), create_config):
+                if mode != "rb":
+                    assert is_file_non_empty(cache.name)
+        self.check_calibrator_cleanup(calibrator)
+
+    # read_calibration_cache should work even if an explicit cache is not provided
+    # This way, it is possible to calibrate quickly when calibrating multiple times.
+    def test_calibrator_caches_without_explicit_cache(self, identity_builder_network):
+        builder, network = identity_builder_network
+        data = [{"x": np.ones((1, 1, 2, 2), dtype=np.float32)}]
+
+        calibrator = Calibrator(data)
+        # First, populate the cache
+        create_config = CreateConfig(int8=True, calibrator=calibrator)
+        with engine_from_network((builder, network), create_config):
+            pass
+
+        # Check that the internal cache is populated
+        assert calibrator.read_calibration_cache()
+        self.check_calibrator_cleanup(calibrator)
+
+    def test_calibrator_rechecks_cache_on_reset(self, identity_builder_network):
+        builder, network = identity_builder_network
+        data = [{"x": np.ones((1, 1, 2, 2), dtype=np.float32)}]
+
+        with util.NamedTemporaryFile(mode="wb+") as cache:
+            calibrator = Calibrator(data, cache=cache.name)
+            # First, populate the cache
+            create_config = CreateConfig(int8=True, calibrator=calibrator)
+            with engine_from_network((builder, network), create_config):
+                pass
+
+            # Ensure that now the calibrator will read from the cache when reset
+            calibrator.reset()
+            assert calibrator.cache_contents is None
+            assert len(calibrator.read_calibration_cache()) == get_file_size(cache.name)
+
+        self.check_calibrator_cleanup(calibrator)
+
+    @pytest.mark.parametrize(
+        "names",
+        [
+            (["fake-input", "x"]),
+            (["fake-input"]),
+        ],
+    )
+    def test_calibrator_invalid_input_fails(self, identity_builder_network, names):
+        builder, network = identity_builder_network
+
+        data = [{name: np.ones((1, 1, 2, 2), dtype=np.float32) for name in names}]
+        calibrator = Calibrator(data)
+
+        create_config = CreateConfig(int8=True, calibrator=calibrator)
+
+        with pytest.raises(PolygraphyException):
+            with engine_from_network((builder, network), create_config):
+                pass
+        self.check_calibrator_cleanup(calibrator)
+
+    @pytest.mark.parametrize(
+        "expected_meta,meta,should_pass",
+        [
+            (
+                TensorMetadata().add(name="input", dtype=np.float32, shape=(1, 3, 28, 28)),
+                TensorMetadata().add(name="input", dtype=np.float32, shape=(1, 3, 28, 28)),
+                True,
+            ),
+            (
+                TensorMetadata().add(name="input", dtype=np.float32, shape=(-1, None, 28, 28)),
+                TensorMetadata().add(name="input", dtype=np.float32, shape=(1, 3, 28, 28)),
+                True,
+            ),
+            # Wrong data type
+            (
+                TensorMetadata().add(name="input", dtype=np.float32, shape=(1, 3, 28, 28)),
+                TensorMetadata().add(name="input", dtype=np.float64, shape=(1, 3, 28, 28)),
+                False,
+            ),
+            # Wrong shape
+            (
+                TensorMetadata().add(name="input", dtype=np.float32, shape=(1, 3, 28, 28)),
+                TensorMetadata().add(name="input", dtype=np.float32, shape=(1, 2, 28, 28)),
+                False,
+            ),
+        ],
+    )
+    def test_calibrator_checks_input_metadata(self, expected_meta, meta, should_pass):
+        data = [
+            {
+                name: np.ones(shape=shape, dtype=DataType.to_dtype(dtype, "numpy"))
+                for name, (dtype, shape) in meta.items()
+            }
+        ]
+        calibrator = Calibrator(data)
+        calibrator.set_input_metadata(expected_meta)
+
+        with calibrator:
+            assert (calibrator.get_batch(list(expected_meta.keys())) is not None) == should_pass
+        self.check_calibrator_cleanup(calibrator)
+
+    def test_calibrator_forces_float32_data(self):
+        data_loader = DataLoader()
+
+        calibrator = Calibrator(data_loader)
+
+        meta = TensorMetadata().add("input", dtype=DataType.FLOAT16, shape=(1, 2, 3))
+        calibrator.set_input_metadata(meta)
+
+        data = data_loader[0]["input"]
+        # TRT requires all calibration inputs to be provided in FP32 regardless of the data type
+        # in the original model.
+        assert util.array.dtype(data) == DataType.FLOAT32
+
+    # TensorRT does not support changing input shapes during calibration
+    @pytest.mark.xfail
+    def test_calibrator_dynamic_shapes(self, dynamic_identity_builder_network):
+        builder, network = dynamic_identity_builder_network
+
+        SHAPES = [(1, 2, 1, 1), (1, 2, 3, 3)]
+
+        def generate_dynamic_shaped_data():
+            for shape in SHAPES:
+                yield {"X": np.ones(shape=shape, dtype=np.float32)}
+
+        calibrator = Calibrator(generate_dynamic_shaped_data())
+
+        create_config = CreateConfig(
+            int8=True,
+            calibrator=calibrator,
+            profiles=[Profile().add(name="X", min=(1, 2, 1, 1), opt=(1, 2, 2, 2), max=(1, 2, 4, 4))],
+        )
+        with engine_from_network((builder, network), create_config) as engine:
+            assert calibrator.num_batches == 2
+            assert engine
+        self.check_calibrator_cleanup(calibrator)

tools/Polygraphy/tests/backend/trt/test_config.py

@@ -0,0 +1,542 @@
+import contextlib
+import os
+import tempfile
+
+import pytest
+
+from polygraphy import mod, util
+from polygraphy.backend.trt import (
+    Calibrator,
+    Profile,
+    network_from_onnx_bytes,
+    postprocess_config,
+)
+from polygraphy.common.struct import BoundedShape
+from polygraphy.comparator import DataLoader
+from polygraphy.datatype import DataType
+from polygraphy import config as polygraphy_config
+from tests.helper import has_dla
+from tests.models.meta import ONNX_MODELS
+
+# Import CreateConfigRTX conditionally for TensorRT-RTX builds
+if polygraphy_config.USE_TENSORRT_RTX:
+    import tensorrt_rtx as trt
+    from polygraphy.backend.tensorrt_rtx import CreateConfigRTX as CreateConfig
+else:
+    import tensorrt as trt
+    from polygraphy.backend.trt import CreateConfig
+
+@pytest.fixture(scope="session")
+def identity_builder_network():
+    builder, network, parser = network_from_onnx_bytes(ONNX_MODELS["identity"].loader)
+    with builder, network, parser:
+        yield builder, network
+
+
+class TestCreateConfig:
+    def test_defaults(self, identity_builder_network):
+        builder, network = identity_builder_network
+        loader = CreateConfig()
+        assert loader.timing_cache_path is None
+
+        with loader(builder, network) as config:
+            assert not config.get_flag(trt.BuilderFlag.DISABLE_TIMING_CACHE)
+            with contextlib.suppress(AttributeError):
+                if polygraphy_config.USE_TENSORRT_RTX:
+                    assert config.get_flag(trt.BuilderFlag.TF32)
+                else:
+                    assert not config.get_flag(trt.BuilderFlag.TF32)
+            with contextlib.suppress(AttributeError):
+                assert not config.get_flag(trt.BuilderFlag.SPARSE_WEIGHTS)
+            assert not config.get_flag(trt.BuilderFlag.FP16)
+            assert not config.get_flag(trt.BuilderFlag.INT8)
+            if mod.version(trt.__version__) >= mod.version("8.7"):
+                assert not config.get_flag(trt.BuilderFlag.BF16)
+            if mod.version(trt.__version__) >= mod.version("8.6"):
+                assert not config.get_flag(trt.BuilderFlag.FP8)
+                assert not config.get_flag(trt.BuilderFlag.VERSION_COMPATIBLE)
+                assert not config.get_flag(trt.BuilderFlag.EXCLUDE_LEAN_RUNTIME)
+                if not polygraphy_config.USE_TENSORRT_RTX:
+                    assert (
+                        config.hardware_compatibility_level
+                        == trt.HardwareCompatibilityLevel.NONE
+                    )
+            if mod.version(trt.__version__) >= mod.version("10.2") and not polygraphy_config.USE_TENSORRT_RTX:
+                assert (
+                    config.runtime_platform
+                    == trt.RuntimePlatform.SAME_AS_BUILD
+                )
+            assert config.num_optimization_profiles == 1
+            if not polygraphy_config.USE_TENSORRT_RTX:
+                assert config.int8_calibrator is None
+            with contextlib.suppress(AttributeError):
+                if mod.version(trt.__version__) >= mod.version("10.0") or polygraphy_config.USE_TENSORRT_RTX:
+                    assert config.get_tactic_sources() == 24
+                elif mod.version(trt.__version__) >= mod.version("8.7"):
+                    assert config.get_tactic_sources() == 29
+                elif mod.version(trt.__version__) >= mod.version("8.5"):
+                    assert config.get_tactic_sources() == 31
+            if mod.version(trt.__version__) >= mod.version("8.7"):
+                assert not config.get_flag(trt.BuilderFlag.ERROR_ON_TIMING_CACHE_MISS)
+            if mod.version(trt.__version__) >= mod.version("8.7"):
+                assert not config.get_flag(trt.BuilderFlag.DISABLE_COMPILATION_CACHE)
+            with contextlib.suppress(AttributeError):
+                assert not config.get_flag(trt.BuilderFlag.OBEY_PRECISION_CONSTRAINTS)
+            with contextlib.suppress(AttributeError):
+                if polygraphy_config.USE_TENSORRT_RTX:
+                    assert config.engine_capability == trt.EngineCapability.STANDARD
+                else:
+                    assert config.engine_capability == trt.EngineCapability.DEFAULT
+            with contextlib.suppress(AttributeError):
+                assert not config.get_flag(trt.BuilderFlag.DIRECT_IO)
+
+    @pytest.mark.parametrize(
+        "engine_capability",
+        [
+            trt.EngineCapability.STANDARD,
+            trt.EngineCapability.SAFETY,
+            trt.EngineCapability.DLA_STANDALONE,
+        ],
+    )
+    def test_engine_capability(self, identity_builder_network, engine_capability):
+        builder, network = identity_builder_network
+        loader = CreateConfig(engine_capability=engine_capability)
+        with loader(builder, network) as config:
+            assert config.engine_capability == engine_capability
+
+    @pytest.mark.parametrize("flag", ["obey", "prefer", None])
+    def test_precision_constraints(self, identity_builder_network, flag):
+        builder, network = identity_builder_network
+        loader = CreateConfig(precision_constraints=flag)
+        with loader(builder, network) as config:
+            obey_set = config.get_flag(trt.BuilderFlag.OBEY_PRECISION_CONSTRAINTS)
+            prefer_set = config.get_flag(trt.BuilderFlag.PREFER_PRECISION_CONSTRAINTS)
+            if flag == "obey":
+                assert obey_set and not prefer_set
+            elif flag == "prefer":
+                assert not obey_set and prefer_set
+            else:
+                assert not obey_set and not prefer_set
+
+    @pytest.mark.skipif(
+        mod.version(trt.__version__) < mod.version("8.6") and not polygraphy_config.USE_TENSORRT_RTX,
+        reason="Unsupported before TRT 8.6",
+    )
+    @pytest.mark.parametrize(
+        "kwargs, expected_flag",
+        [
+            ({"version_compatible": True}, "VERSION_COMPATIBLE"),
+            (
+                {"version_compatible": True, "exclude_lean_runtime": True},
+                "EXCLUDE_LEAN_RUNTIME",
+            ),
+        ],
+    )
+    def test_version_compatibility_flags(
+        self, identity_builder_network, kwargs, expected_flag
+    ):
+        builder, network = identity_builder_network
+        loader = CreateConfig(**kwargs)
+        with loader(builder, network) as config:
+            assert config.get_flag(getattr(trt.BuilderFlag, expected_flag))
+
+    def test_direct_io(self, identity_builder_network):
+        builder, network = identity_builder_network
+        loader = CreateConfig(direct_io=True)
+        with loader(builder, network) as config:
+            assert config.get_flag(trt.BuilderFlag.DIRECT_IO)
+
+    @pytest.mark.parametrize("flag", [True, False])
+    def test_restricted(self, identity_builder_network, flag):
+        builder, network = identity_builder_network
+        loader = CreateConfig(restricted=flag)
+        with loader(builder, network) as config:
+            assert config.get_flag(trt.BuilderFlag.SAFETY_SCOPE) == flag
+
+    @pytest.mark.parametrize(
+        "arg_name, flag_type",
+        [
+            ("refittable", trt.BuilderFlag.REFIT),
+        ]
+        + (
+            [
+                (
+                    "disable_compilation_cache",
+                    trt.BuilderFlag.DISABLE_COMPILATION_CACHE,
+                ),
+            ]
+            if mod.version(trt.__version__) >= mod.version("9.0")
+            else []
+        )
+        + (
+            [
+                ("strip_plan", trt.BuilderFlag.STRIP_PLAN),
+            ]
+            if mod.version(trt.__version__) >= mod.version("10.0")
+            else []
+        )
+        + (
+            [
+                ("fp16", trt.BuilderFlag.FP16),
+                ("int8", trt.BuilderFlag.INT8),
+                ("allow_gpu_fallback", trt.BuilderFlag.GPU_FALLBACK),
+                ("tf32", trt.BuilderFlag.TF32),
+            ]
+            + (
+                [
+                    ("bf16", trt.BuilderFlag.BF16),
+                ]
+                if mod.version(trt.__version__) >= mod.version("8.7")
+                else []
+            )
+            + (
+                [
+                    ("fp8", trt.BuilderFlag.FP8),
+                ]
+                if mod.version(trt.__version__) >= mod.version("8.6")
+                else []
+            )
+            if not polygraphy_config.USE_TENSORRT_RTX
+            else []
+        ),
+    )
+    @pytest.mark.parametrize("value", [True, False])
+    def test_flags(self, identity_builder_network, arg_name, flag_type, value):
+        builder, network = identity_builder_network
+        loader = CreateConfig(**{arg_name: value})
+        with loader(builder, network) as config:
+            assert config.get_flag(flag_type) == value
+
+    @pytest.mark.parametrize("flag", [True, False])
+    def test_sparse_weights(self, identity_builder_network, flag):
+        builder, network = identity_builder_network
+        loader = CreateConfig(sparse_weights=flag)
+        with loader(builder, network) as config:
+            assert config.get_flag(trt.BuilderFlag.SPARSE_WEIGHTS) == flag
+
+    @pytest.mark.skipif(
+        polygraphy_config.USE_TENSORRT_RTX,
+        reason="TensorRT-RTX does not support DLA"
+    )
+    def test_use_dla(self, identity_builder_network):
+        builder, network = identity_builder_network
+        loader = CreateConfig(use_dla=True)
+        with loader(builder, network) as config:
+            assert config.default_device_type == trt.DeviceType.DLA
+            if has_dla():
+                assert config.DLA_core == 0
+
+    with contextlib.suppress(AttributeError):
+        TACTIC_SOURCES_CASES = [
+            (None, 31),  # By default, all sources are enabled.
+            ([], 0),
+            ([trt.TacticSource.CUBLAS], 1),
+            ([trt.TacticSource.CUBLAS_LT], 2),
+            ([trt.TacticSource.CUDNN], 4),
+            ([trt.TacticSource.CUBLAS, trt.TacticSource.CUBLAS_LT], 3),
+            ([trt.TacticSource.CUBLAS, trt.TacticSource.CUDNN], 5),
+            ([trt.TacticSource.CUBLAS_LT, trt.TacticSource.CUDNN], 6),
+            (
+                [
+                    trt.TacticSource.CUDNN,
+                    trt.TacticSource.CUBLAS,
+                    trt.TacticSource.CUBLAS_LT,
+                ],
+                7,
+            ),
+            (
+                [
+                    trt.TacticSource.CUDNN,
+                    trt.TacticSource.CUBLAS,
+                    trt.TacticSource.CUBLAS_LT,
+                    trt.TacticSource.EDGE_MASK_CONVOLUTIONS,
+                ],
+                15,
+            ),
+            (
+                [
+                    trt.TacticSource.CUDNN,
+                    trt.TacticSource.CUBLAS,
+                    trt.TacticSource.CUBLAS_LT,
+                    trt.TacticSource.EDGE_MASK_CONVOLUTIONS,
+                    trt.TacticSource.JIT_CONVOLUTIONS,
+                ],
+                31,
+            ),
+        ]
+
+        if mod.version(trt.__version__) >= mod.version("10.0") or polygraphy_config.USE_TENSORRT_RTX:
+            TACTIC_SOURCES_CASES[0] = (None, 24)
+        elif mod.version(trt.__version__) >= mod.version("8.7"):
+            TACTIC_SOURCES_CASES[0] = (None, 29)
+
+    @pytest.mark.parametrize("sources, expected", TACTIC_SOURCES_CASES)
+    def test_tactic_sources(self, identity_builder_network, sources, expected):
+        builder, network = identity_builder_network
+        loader = CreateConfig(tactic_sources=sources)
+        with loader(builder, network) as config:
+            assert config.get_tactic_sources() == expected
+
+    @pytest.mark.skipif(
+        mod.version(trt.__version__) < mod.version("8.7") and not polygraphy_config.USE_TENSORRT_RTX,
+        reason="API was added in TRT 8.7",
+    )
+    @pytest.mark.parametrize("flag", [True, False])
+    def test_error_on_timing_cache_miss(self, identity_builder_network, flag):
+        builder, network = identity_builder_network
+        loader = CreateConfig(error_on_timing_cache_miss=flag)
+        with loader(builder, network) as config:
+            assert config.get_flag(trt.BuilderFlag.ERROR_ON_TIMING_CACHE_MISS) == flag
+
+    @pytest.mark.skipif(
+        polygraphy_config.USE_TENSORRT_RTX,
+        reason="TensorRT-RTX does not support calibrators"
+    )
+    def test_calibrator_metadata_set(self, identity_builder_network):
+        builder, network = identity_builder_network
+        calibrator = Calibrator(DataLoader())
+        loader = CreateConfig(int8=True, calibrator=calibrator)
+        with loader(builder, network) as config:
+            assert config.int8_calibrator
+            assert "x" in calibrator.data_loader.input_metadata
+            meta = calibrator.data_loader.input_metadata["x"]
+            assert meta.shape == BoundedShape((1, 1, 2, 2))
+            assert meta.dtype == DataType.FLOAT32
+
+    def test_multiple_profiles(self, identity_builder_network):
+        builder, network = identity_builder_network
+        profiles = [
+            Profile().add("x", (1, 2, 1, 1), (1, 2, 2, 2), (1, 2, 4, 4)),
+            Profile().add("x", (1, 2, 4, 4), (1, 2, 8, 8), (1, 2, 16, 16)),
+        ]
+        loader = CreateConfig(profiles=profiles)
+        with loader(builder, network) as config:
+            assert config.num_optimization_profiles == 2
+
+    @pytest.mark.parametrize("path_mode", [True, False], ids=["path", "file-like"])
+    def test_timing_cache(self, identity_builder_network, path_mode):
+        builder, network = identity_builder_network
+        with util.NamedTemporaryFile() as cache:
+            loader = CreateConfig(load_timing_cache=cache.name if path_mode else cache)
+            with loader(builder, network) as config:
+                assert config.get_timing_cache()
+
+    def test_fall_back_to_empty_timing_cache(self, identity_builder_network):
+        """Tests that passing in a nonexistent timing cache path is non-fatal"""
+        builder, network = identity_builder_network
+        with tempfile.TemporaryDirectory() as tmpdir:
+            cache_name = os.path.join(tmpdir, "casper")
+            loader = CreateConfig(load_timing_cache=cache_name)
+            with loader(builder, network) as config:
+                assert config.get_timing_cache()
+
+    def test_empty_timing_cache_when_default(self, identity_builder_network):
+        builder, network = identity_builder_network
+        loader = CreateConfig()
+        with loader(builder, network) as config:
+            cache = config.get_timing_cache()
+            with cache.serialize() as buffer:
+                cache_size = len(bytes(buffer))
+
+            cache.reset()
+            with cache.serialize() as buffer:
+                new_cache_size = len(bytes(buffer))
+            assert cache_size == new_cache_size
+
+    def test_profiling_verbosity(self, identity_builder_network):
+        builder, network = identity_builder_network
+        expected = trt.ProfilingVerbosity.NONE
+        loader = CreateConfig(profiling_verbosity=expected)
+        with loader(builder, network) as config:
+            assert config.profiling_verbosity == expected
+
+    with contextlib.suppress(AttributeError):
+        POOL_LIMITS = [
+            {trt.MemoryPoolType.WORKSPACE: 25},
+            {trt.MemoryPoolType.DLA_MANAGED_SRAM: 25},
+            {trt.MemoryPoolType.DLA_LOCAL_DRAM: 25},
+            {trt.MemoryPoolType.DLA_GLOBAL_DRAM: 25},
+            # Multiple limits
+            {
+                trt.MemoryPoolType.DLA_LOCAL_DRAM: 20,
+                trt.MemoryPoolType.DLA_GLOBAL_DRAM: 25,
+                trt.MemoryPoolType.WORKSPACE: 39,
+            },
+        ]
+
+        # @pytest.mark.skipif(
+        #     config.USE_TENSORRT_RTX,
+        #     reason="TensorRT-RTX does not support DLA memory pools"
+        # )
+        @pytest.mark.parametrize("pool_limits", POOL_LIMITS)
+        def test_memory_pool_limits(self, pool_limits, identity_builder_network):
+            if any("dla" in key.name.lower() for key in pool_limits) and not has_dla():
+                pytest.skip("DLA is not available on this system")
+
+            builder, network = identity_builder_network
+            loader = CreateConfig(memory_pool_limits=pool_limits)
+            with loader(builder, network) as config:
+                for pool_type, pool_size in pool_limits.items():
+                    assert config.get_memory_pool_limit(pool_type) == pool_size
+
+    @pytest.mark.parametrize(
+        "preview_features",
+        [
+            [trt.PreviewFeature.PROFILE_SHARING_0806]
+            if mod.version(trt.__version__) >= mod.version("10.0")
+            else (
+                [trt.PreviewFeature.ALIASED_PLUGIN_IO_10_03]
+                if polygraphy_config.USE_TENSORRT_RTX
+                else [trt.PreviewFeature.FASTER_DYNAMIC_SHAPES_0805]
+            ),
+        ],
+    )
+    def test_preview_features(self, identity_builder_network, preview_features):
+        builder, network = identity_builder_network
+        loader = CreateConfig(preview_features=preview_features)
+        with loader(builder, network) as config:
+            # Check that only the enabled preview features are on.
+            for pf in trt.PreviewFeature.__members__.values():
+                expected = pf in preview_features
+                # TensorRT-RTX enables PROFILE_SHARING_0806 by default and can't be disabled
+                if polygraphy_config.USE_TENSORRT_RTX and pf == trt.PreviewFeature.PROFILE_SHARING_0806:
+                    expected = True
+                assert config.get_preview_feature(pf) == expected
+
+    @pytest.mark.skipif(
+        polygraphy_config.USE_TENSORRT_RTX,
+        reason="TensorRT-RTX does not support quantization_flag API"
+    )
+    @pytest.mark.parametrize(
+        "quantization_flags",
+        [
+            [],
+            [trt.QuantizationFlag.CALIBRATE_BEFORE_FUSION],
+        ],
+    )
+    def test_quantization_flags(self, identity_builder_network, quantization_flags):
+        builder, network = identity_builder_network
+        loader = CreateConfig(quantization_flags=quantization_flags)
+        with loader(builder, network) as config:
+            # Check that only the enabled quantization flags are on.
+            for qf in trt.QuantizationFlag.__members__.values():
+                assert config.get_quantization_flag(qf) == (qf in quantization_flags)
+
+    @pytest.mark.skipif(
+        mod.version(trt.__version__) < mod.version("8.6") and not polygraphy_config.USE_TENSORRT_RTX,
+        reason="Unsupported for TRT versions prior to 8.6",
+    )
+    @pytest.mark.parametrize("level", range(6))
+    def test_builder_optimization_level(self, identity_builder_network, level):
+        builder, network = identity_builder_network
+        loader = CreateConfig(builder_optimization_level=level)
+        with loader(builder, network) as config:
+            assert config.builder_optimization_level == level
+
+    if mod.version(trt.__version__) >= mod.version("8.6"):
+
+        @pytest.mark.parametrize(
+            "level",
+            [
+                trt.HardwareCompatibilityLevel.NONE,
+                trt.HardwareCompatibilityLevel.AMPERE_PLUS,
+            ],
+        )
+        def test_hardware_compatibility_level(self, identity_builder_network, level):
+            builder, network = identity_builder_network
+            loader = CreateConfig(hardware_compatibility_level=level)
+            with loader(builder, network) as config:
+                assert config.hardware_compatibility_level == level
+
+    if mod.version(trt.__version__) >= mod.version("10.2"):
+
+        @pytest.mark.parametrize(
+            "platform",
+            [
+                trt.RuntimePlatform.SAME_AS_BUILD,
+                trt.RuntimePlatform.WINDOWS_AMD64,
+            ],
+        )
+        def test_runtime_platform(self, identity_builder_network, platform):
+            builder, network = identity_builder_network
+            loader = CreateConfig(runtime_platform=platform)
+            with loader(builder, network) as config:
+                assert config.runtime_platform == platform
+
+    @pytest.mark.skipif(
+        mod.version(trt.__version__) < mod.version("8.6") and not polygraphy_config.USE_TENSORRT_RTX,
+        reason="Unsupported for TRT versions prior to 8.6",
+    )
+    @pytest.mark.parametrize("num_streams", range(3))
+    def test_max_aux_streams(self, identity_builder_network, num_streams):
+        builder, network = identity_builder_network
+        loader = CreateConfig(max_aux_streams=num_streams)
+        with loader(builder, network) as config:
+            assert config.max_aux_streams == num_streams
+
+    @pytest.mark.skipif(
+        mod.version(trt.__version__) < mod.version("9.0") and not polygraphy_config.USE_TENSORRT_RTX,
+        reason="API was added in TRT 9.0",
+    )
+    def test_progress_monitor(self, identity_builder_network):
+        class DummyProgressMonitor(trt.IProgressMonitor):
+            def __init__(self):
+                trt.IProgressMonitor.__init__(self)
+
+            def phase_start(self, phase_name, parent_phase, num_steps):
+                pass
+
+            def phase_finish(self, phase_name):
+                pass
+
+            def step_complete(self, phase_name, step):
+                return True
+
+        builder, network = identity_builder_network
+        progress_monitor = DummyProgressMonitor()
+        loader = CreateConfig(progress_monitor=progress_monitor)
+        with loader(builder, network) as config:
+            assert config.progress_monitor == progress_monitor
+
+    if mod.version(trt.__version__) >= mod.version("10.8") and not polygraphy_config.USE_TENSORRT_RTX:
+        @pytest.mark.parametrize(
+            "level",
+            [
+                trt.TilingOptimizationLevel.NONE,
+                trt.TilingOptimizationLevel.FAST,
+                trt.TilingOptimizationLevel.MODERATE,
+                trt.TilingOptimizationLevel.FULL,
+            ],
+        )
+        def test_tiling_optimization_level(self, identity_builder_network, level):
+            builder, network = identity_builder_network
+            loader = CreateConfig(tiling_optimization_level=level)
+            with loader(builder, network) as config:
+                assert config.tiling_optimization_level == level
+
+
+class TestPostprocessConfig:
+    def test_with_config(self, identity_builder_network):
+        builder, network = identity_builder_network
+        config = CreateConfig()(builder, network)
+        assert not config.get_flag(trt.BuilderFlag.INT8)
+
+        config = postprocess_config(
+            config,
+            func=lambda builder, network, config: config.set_flag(trt.BuilderFlag.INT8),
+            builder=builder,
+            network=network,
+        )
+        assert config.get_flag(trt.BuilderFlag.INT8)
+
+    def test_with_config_callable(self, identity_builder_network):
+        builder, network = identity_builder_network
+        config = CreateConfig()
+
+        config = postprocess_config(
+            config,
+            func=lambda builder, network, config: config.set_flag(trt.BuilderFlag.INT8),
+            builder=builder,
+            network=network,
+        )
+        assert config.get_flag(trt.BuilderFlag.INT8)

tools/Polygraphy/tests/backend/trt/test_loader.py

@@ -0,0 +1,664 @@
+#
+# 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 __future__ import annotations
+
+import sys
+import pytest
+
+from polygraphy import config, constants, mod, util
+from polygraphy.backend.trt import (
+    Calibrator,
+    EngineBytesFromNetwork,
+    EngineFromBytes,
+    EngineFromNetwork,
+    EngineFromPath,
+    LoadPlugins,
+    LoadRuntime,
+    ModifyNetworkOutputs,
+    NetworkFromOnnxBytes,
+    Profile,
+    SaveEngine,
+    buffer_from_engine,
+    bytes_from_engine,
+    create_config,
+    create_network,
+    engine_from_network,
+    get_trt_logger,
+    modify_network_outputs,
+    network_from_onnx_bytes,
+    network_from_onnx_path,
+    onnx_like_from_network,
+    postprocess_network,
+    set_layer_precisions,
+    set_tensor_datatypes,
+    set_tensor_formats,
+)
+from polygraphy.common.struct import BoundedShape
+from polygraphy.comparator import DataLoader
+from polygraphy.datatype import DataType
+from polygraphy.exception import PolygraphyException
+from tests.helper import get_file_size, is_file_non_empty
+from tests.models.meta import ONNX_MODELS
+
+# Import CreateConfigRTX conditionally for TensorRT-RTX builds
+if config.USE_TENSORRT_RTX:
+    import tensorrt_rtx as trt
+    from polygraphy.backend.tensorrt_rtx import CreateConfigRTX as CreateConfig
+else:
+    import tensorrt as trt
+    from polygraphy.backend.trt import CreateConfig
+
+
+##
+## Fixtures
+##
+
+
+@pytest.fixture(scope="session")
+def identity_engine():
+    network_loader = NetworkFromOnnxBytes(ONNX_MODELS["identity"].loader)
+    engine_loader = EngineFromNetwork(network_loader)
+    with engine_loader() as engine:
+        yield engine
+
+
+@pytest.fixture(scope="session")
+def identity_vc_engine_bytes():
+    flags = [trt.OnnxParserFlag.NATIVE_INSTANCENORM]
+    config = CreateConfig(version_compatible=True)
+    network_loader = NetworkFromOnnxBytes(ONNX_MODELS["identity"].loader, flags=flags)
+    engine_loader = EngineBytesFromNetwork(network_loader, config=config)
+    with engine_loader() as engine_bytes:
+        yield engine_bytes
+
+
+@pytest.fixture(scope="session")
+def identity_builder_network():
+    builder, network, parser = network_from_onnx_bytes(ONNX_MODELS["identity"].loader)
+    yield builder, network
+
+
+@pytest.fixture(scope="session")
+def identity_network():
+    builder, network, parser = network_from_onnx_bytes(ONNX_MODELS["identity"].loader)
+    yield builder, network, parser
+
+
+@pytest.fixture(scope="session")
+def identity_identity_network():
+    builder, network, parser = network_from_onnx_bytes(
+        ONNX_MODELS["identity_identity"].loader
+    )
+    yield builder, network, parser
+
+
+@pytest.fixture(scope="session")
+def reshape_network():
+    builder, network, parser = network_from_onnx_bytes(ONNX_MODELS["reshape"].loader)
+    yield builder, network, parser
+
+
+@pytest.fixture(scope="session")
+def modifiable_network():
+    # Must return a loader since the network will be modified each time it's loaded.
+    return NetworkFromOnnxBytes(ONNX_MODELS["identity_identity"].loader)
+
+
+@pytest.fixture(scope="session")
+def modifiable_reshape_network():
+    # Must return a loader since the network will be modified each time it's loaded.
+    return NetworkFromOnnxBytes(ONNX_MODELS["reshape"].loader)
+
+
+##
+## Tests
+##
+
+
+class TestLoadPlugins:
+    @pytest.mark.skipif(
+        config.USE_TENSORRT_RTX,
+        reason="Plugin tests are not compatible with TensorRT-RTX"
+    )
+    def test_can_load_libnvinfer_plugins(self):
+        def get_plugin_names():
+            return [pc.name for pc in trt.get_plugin_registry().plugin_creator_list]
+
+        loader = LoadPlugins(
+            plugins=[
+                (
+                    "nvinfer_plugin.dll"
+                    if sys.platform.startswith("win")
+                    else "libnvinfer_plugin.so"
+                )
+            ]
+        )
+        loader()
+        assert get_plugin_names()
+
+
+class TestSerializedEngineLoader:
+    def test_serialized_engine_loader_from_lambda(self, identity_engine):
+        with util.NamedTemporaryFile() as outpath:
+            with open(outpath.name, "wb") as f, identity_engine.serialize() as buffer:
+                f.write(buffer)
+
+            loader = EngineFromBytes(lambda: open(outpath.name, "rb").read())
+            with loader() as engine:
+                assert isinstance(engine, trt.ICudaEngine)
+        
+    def test_serialized_engine_loader_from_buffer(self, identity_engine):
+        with identity_engine.serialize() as buffer:
+            loader = EngineFromBytes(buffer)
+            with loader() as engine:
+                assert isinstance(engine, trt.ICudaEngine)
+
+    def test_serialized_engine_loader_custom_runtime(self, identity_engine):
+        with identity_engine.serialize() as buffer:
+            loader = EngineFromBytes(buffer, runtime=trt.Runtime(get_trt_logger()))
+            with loader() as engine:
+                assert isinstance(engine, trt.ICudaEngine)
+
+
+@pytest.mark.skipif(
+    mod.version(trt.__version__) < mod.version("10.0") and not config.USE_TENSORRT_RTX, reason="API was added in TRT 10.0"
+)
+class TestSerializedEngineLoaderFromDisk:
+    def test_serialized_engine_loader_from_lambda(self, identity_engine):
+        with util.NamedTemporaryFile() as outpath:
+            with open(outpath.name, "wb") as f, identity_engine.serialize() as buffer:
+                f.write(buffer)
+
+            loader = EngineFromPath(lambda: outpath.name)
+            with loader() as engine:
+                assert isinstance(engine, trt.ICudaEngine)
+
+    def test_serialized_engine_loader_custom_runtime(self, identity_engine):
+        with util.NamedTemporaryFile() as outpath:
+            with open(outpath.name, "wb") as f, identity_engine.serialize() as buffer:
+                f.write(buffer)
+            
+            loader = EngineFromPath(lambda: outpath.name, runtime=trt.Runtime(get_trt_logger()))
+            with loader() as engine:
+                assert isinstance(engine, trt.ICudaEngine)
+
+
+@pytest.mark.skipif(
+    mod.version(trt.__version__) < mod.version("8.6"), reason="API was added in TRT 8.6"
+)
+@pytest.mark.skipif(
+    config.USE_TENSORRT_RTX,
+    reason="TensorRT-RTX does not have lean runtime shared objects"
+)
+class TestLoadRuntime:
+    def test_load_lean_runtime(self, nvinfer_lean_path):
+        loader = LoadRuntime(nvinfer_lean_path)
+        with loader() as runtime:
+            assert isinstance(runtime, trt.Runtime)
+
+
+@pytest.mark.skipif(
+    mod.version(trt.__version__) < mod.version("8.6") and not config.USE_TENSORRT_RTX, reason="API was added in TRT 8.6"
+)
+@pytest.mark.skipif(
+    config.USE_TENSORRT_RTX,
+    reason="TensorRT-RTX does not have libnvinfer_lean.so.1"
+)
+class TestSerializedVCEngineLoader:
+    def test_serialized_vc_engine_loader_from_lambda(self, identity_vc_engine_bytes):
+        with util.NamedTemporaryFile() as outpath:
+            with open(outpath.name, "wb") as f:
+                f.write(identity_vc_engine_bytes)
+
+            loader = EngineFromBytes(lambda: open(outpath.name, "rb").read())
+            with loader() as engine:
+                assert isinstance(engine, trt.ICudaEngine)
+
+    def test_serialized_engine_loader_from_buffer(self, identity_vc_engine_bytes):
+        loader = EngineFromBytes(identity_vc_engine_bytes)
+        with loader() as engine:
+            assert isinstance(engine, trt.ICudaEngine)
+
+
+class TestNetworkFromOnnxBytes:
+    def test_loader(self):
+        builder, network, parser = network_from_onnx_bytes(
+            ONNX_MODELS["identity"].loader
+        )
+        if not config.USE_TENSORRT_RTX:
+            assert not network.has_implicit_batch_dimension
+
+    @pytest.mark.parametrize(
+        "kwargs, flag",
+        (
+            [
+                (
+                    {"strongly_typed": True},
+                    trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED,
+                )
+            ]
+            if mod.version(trt.__version__) >= mod.version("8.7") and not config.USE_TENSORRT_RTX
+            else []
+        ),
+    )
+    def test_network_flags(self, kwargs, flag):
+        builder, network, parser = network_from_onnx_bytes(
+            ONNX_MODELS["identity"].loader, **kwargs
+        )
+        assert network.get_flag(flag)
+
+
+class TestNetworkFromOnnxPath:
+    def test_loader(self):
+        builder, network, parser = network_from_onnx_path(ONNX_MODELS["identity"].path)
+        if not config.USE_TENSORRT_RTX:
+            assert not network.has_implicit_batch_dimension
+
+    @pytest.mark.parametrize(
+        "kwargs, flag",
+        (
+            [
+                (
+                    {"strongly_typed": True},
+                    trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED,
+                )
+            ]
+            if mod.version(trt.__version__) >= mod.version("8.7") and not config.USE_TENSORRT_RTX
+            else []
+        ),
+    )
+    def test_network_flags(self, kwargs, flag):
+        builder, network, parser = network_from_onnx_path(
+            ONNX_MODELS["identity"].path, **kwargs
+        )
+        assert network.get_flag(flag)
+
+
+class TestModifyNetwork:
+    def test_mark_layerwise(self, modifiable_network):
+        load_network = ModifyNetworkOutputs(
+            modifiable_network, outputs=constants.MARK_ALL
+        )
+        builder, network, parser = load_network()
+
+        for layer in network:
+            for index in range(layer.num_outputs):
+                assert layer.get_output(index).is_network_output
+
+    def test_mark_custom_outputs(self, modifiable_network):
+        builder, network, parser = modify_network_outputs(
+            modifiable_network, outputs=["identity_out_0"]
+        )
+
+        assert network.num_outputs == 1
+        assert network.get_output(0).name == "identity_out_0"
+
+    def test_exclude_outputs_with_mark_layerwise(self, modifiable_network):
+        builder, network, parser = modify_network_outputs(
+            modifiable_network,
+            outputs=constants.MARK_ALL,
+            exclude_outputs=["identity_out_2"],
+        )
+
+        assert network.num_outputs == 1
+        assert network.get_output(0).name == "identity_out_0"
+
+    def test_mark_shape_outputs(self, modifiable_reshape_network):
+        builder, network, parser = modify_network_outputs(
+            modifiable_reshape_network, outputs=["output", "reduce_prod_out_gs_2"]
+        )
+
+        assert network.num_outputs == 2
+        assert network.get_output(1).name == "reduce_prod_out_gs_2"
+
+    def test_unmark_shape_outputs(self, modifiable_reshape_network):
+        builder, network, parser = modify_network_outputs(
+            modifiable_reshape_network,
+            outputs=constants.MARK_ALL,
+            exclude_outputs=["shape_out_gs_0", "reduce_prod_out_gs_2"],
+        )
+
+        assert network.num_outputs == 1
+
+    def test_mark_outputs_layer_with_optional_inputs(self):
+        builder, network = create_network()
+        inp = network.add_input("input", shape=(1, 3, 224, 224), dtype=trt.float32)
+        slice_layer = network.add_slice(
+            inp, (0, 0, 0, 0), (1, 3, 224, 224), (1, 1, 1, 1)
+        )
+
+        # Set a tensor for `stride` to increment `num_inputs` so we have some inputs
+        # which are `None` in between.
+        slice_layer.set_input(3, inp)
+        assert slice_layer.num_inputs == 4
+
+        slice = slice_layer.get_output(0)
+        slice.name = "Slice"
+
+        builder, network = modify_network_outputs((builder, network), outputs=["Slice"])
+        assert network.num_outputs == 1
+        assert network.get_output(0).name == "Slice"
+        assert network.get_output(0) == slice
+
+
+class TestPostprocessNetwork:
+    def test_basic(self, modifiable_network):
+        """Tests that the callback is actually invoked by Polygraphy."""
+        func_called = False
+
+        def func(network):
+            nonlocal func_called
+            func_called = True
+            assert isinstance(network, trt.INetworkDefinition)
+
+        builder, network, parser = postprocess_network(modifiable_network, func)
+        assert func_called
+
+    def test_kwargs(self, modifiable_network):
+        """Tests that callbacks that use **kwargs work as expected."""
+        func_called = False
+
+        def func(**kwargs):
+            nonlocal func_called
+            func_called = True
+            assert isinstance(kwargs["network"], trt.INetworkDefinition)
+
+        builder, network, parser = postprocess_network(modifiable_network, func)
+        assert func_called
+
+    @pytest.mark.skipif(
+        config.USE_TENSORRT_RTX,
+        reason="TensorRT-RTX uses strongly typed networks where layer precision cannot be set"
+    )
+    def test_modify_network(self, modifiable_network):
+        """Tests that the network passed in is properly modified by the callback."""
+
+        # Performs the equivalent of set_layer_precisions
+        def func(network):
+            for layer in network:
+                if layer.name == "onnx_graphsurgeon_node_1":
+                    layer.precision = trt.float16
+                if layer.name == "onnx_graphsurgeon_node_3":
+                    layer.precision = trt.int8
+
+        builder, network, parser = postprocess_network(modifiable_network, func)
+
+        assert network[0].precision == trt.float16
+        assert network[1].precision == trt.int8
+
+    def test_negative_non_callable(self, modifiable_network):
+        """Tests that PostprocessNetwork properly rejects `func` objects that
+        are not callable."""
+
+        with pytest.raises(PolygraphyException, match=r"Object .* is not a callable"):
+            builder, network, parser = postprocess_network(modifiable_network, None)
+
+
+class TestSetLayerPrecisions:
+    @pytest.mark.skipif(
+        config.USE_TENSORRT_RTX,
+        reason="TensorRT-RTX uses strongly typed networks where layer precision cannot be set"
+    )
+    def test_basic(self, modifiable_network):
+        builder, network, parser = set_layer_precisions(
+            modifiable_network,
+            layer_precisions={
+                "onnx_graphsurgeon_node_1": trt.float16,
+                "onnx_graphsurgeon_node_3": trt.int8,
+            },
+        )
+
+        assert network[0].precision == trt.float16
+        assert network[1].precision == trt.int8
+
+
+class TestSetTensorDatatypes:
+    @pytest.mark.skipif(
+        config.USE_TENSORRT_RTX,
+        reason="TensorRT-RTX uses strongly typed networks where tensor datatypes cannot be set"
+    )
+    def test_basic(self, modifiable_network):
+        builder, network, parser = set_tensor_datatypes(
+            modifiable_network,
+            tensor_datatypes={
+                "X": trt.float16,
+                "identity_out_2": trt.float16,
+            },
+        )
+
+        assert network.get_input(0).dtype == trt.float16
+        assert network.get_output(0).dtype == trt.float16
+
+
+class TestSetTensorFormats:
+    def test_basic(self, modifiable_network):
+        builder, network, parser = set_tensor_formats(
+            modifiable_network,
+            tensor_formats={
+                "X": [trt.TensorFormat.LINEAR, trt.TensorFormat.CHW4],
+                "identity_out_2": [trt.TensorFormat.HWC8],
+            },
+        )
+
+        assert network.get_input(0).allowed_formats == (
+            1 << int(trt.TensorFormat.LINEAR) | 1 << int(trt.TensorFormat.CHW4)
+        )
+        assert network.get_output(0).allowed_formats == 1 << int(trt.TensorFormat.HWC8)
+
+
+class TestEngineBytesFromNetwork:
+    def test_can_build(self, identity_network):
+        loader = EngineBytesFromNetwork(identity_network)
+        with loader() as serialized_engine:
+            assert isinstance(serialized_engine, trt.IHostMemory)
+
+
+class TestEngineFromNetwork:
+    def test_defaults(self, identity_network):
+        loader = EngineFromNetwork(identity_network)
+        assert loader.timing_cache_path is None
+
+    def test_can_build_with_parser_owning(self, identity_network):
+        loader = EngineFromNetwork(identity_network)
+        with loader() as engine:
+            assert isinstance(engine, trt.ICudaEngine)
+
+    def test_can_build_without_parser_non_owning(self, identity_builder_network):
+        builder, network = identity_builder_network
+        loader = EngineFromNetwork((builder, network))
+        with loader() as engine:
+            assert isinstance(engine, trt.ICudaEngine)
+
+    def test_custom_runtime(self, identity_builder_network):
+        builder, network = identity_builder_network
+        loader = EngineFromNetwork(
+            (builder, network), runtime=trt.Runtime(get_trt_logger())
+        )
+        with loader() as engine:
+            assert isinstance(engine, trt.ICudaEngine)
+
+    @pytest.mark.skipif(
+        config.USE_TENSORRT_RTX,
+        reason="TensorRT-RTX does not support calibrators"
+    )
+    @pytest.mark.parametrize(
+        "use_config_loader, set_calib_profile",
+        [(True, None), (False, False), (False, True)],
+    )
+    def test_can_build_with_calibrator(
+        self, identity_builder_network, use_config_loader, set_calib_profile
+    ):
+        builder, network = identity_builder_network
+        calibrator = Calibrator(DataLoader())
+
+        def check_calibrator():
+            # CreateConfig and EngineFromNetwork should set the input metadata for the calibrator,
+            # which in turn should be passed to the data loader.
+            assert calibrator.input_metadata is not None
+            assert "x" in calibrator.data_loader.input_metadata
+            meta = calibrator.data_loader.input_metadata["x"]
+            assert meta.shape == BoundedShape((1, 1, 2, 2))
+            assert meta.dtype == DataType.FLOAT32
+
+        if use_config_loader:
+            config = create_config(builder, network, int8=True, calibrator=calibrator)
+            check_calibrator()
+        else:
+            config = builder.create_builder_config()
+            config.set_flag(trt.BuilderFlag.INT8)
+            config.int8_calibrator = calibrator
+            # Since this network has static shapes, we shouldn't need to set a calibration profile.
+            if set_calib_profile:
+                calib_profile = (
+                    Profile().fill_defaults(network).to_trt(builder, network)
+                )
+                config.add_optimization_profile(calib_profile)
+                config.set_calibration_profile(calib_profile)
+
+        loader = EngineFromNetwork((builder, network), config)
+        with loader():
+            pass
+        check_calibrator()
+
+        # Calibrator buffers should be freed after the build
+        assert all(
+            [buf.allocated_nbytes == 0 for buf in calibrator.device_buffers.values()]
+        )
+
+    @pytest.mark.parametrize("path_mode", [True, False], ids=["path", "file-like"])
+    def test_timing_cache_generate_and_append(self, path_mode):
+        with util.NamedTemporaryFile() as total_cache, util.NamedTemporaryFile() as identity_cache:
+
+            def build_engine(model, cache):
+                if not path_mode:
+                    cache.seek(0)
+                network_loader = NetworkFromOnnxBytes(ONNX_MODELS[model].loader)
+                # In non-path_mode, use the file-like object directly.
+                # Must load the cache with CreateConfig so that new data is appended
+                # instead of overwriting the previous cache.
+                loader = EngineFromNetwork(
+                    network_loader,
+                    CreateConfig(load_timing_cache=cache.name),
+                    save_timing_cache=cache.name if path_mode else cache,
+                )
+                with loader():
+                    pass
+                if not path_mode:
+                    cache.seek(0)
+
+            assert not total_cache.read()
+
+            build_engine("const_foldable", total_cache)
+            const_foldable_cache_size = get_file_size(total_cache.name)
+
+            # Build this network twice. Once with a fresh cache so we can determine its size.
+            assert get_file_size(identity_cache.name) == 0
+            build_engine("identity", identity_cache)
+            identity_cache_size = get_file_size(identity_cache.name)
+
+            build_engine("identity", total_cache)
+            total_cache_size = get_file_size(total_cache.name)
+
+            # The total cache should be larger than either of the individual caches.
+            assert (
+                total_cache_size >= const_foldable_cache_size
+                and total_cache_size >= identity_cache_size
+            )
+            # The total cache should also be smaller than or equal to the sum of the individual caches since
+            # header information should not be duplicated.
+            assert total_cache_size <= (const_foldable_cache_size + identity_cache_size)
+
+
+class TestBytesFromEngine:
+    def test_serialize_engine(self, identity_network):
+        with engine_from_network(identity_network) as engine:
+            serialized_engine = bytes_from_engine(engine)
+            assert isinstance(serialized_engine, bytes)
+
+
+class TestBufferFromEngine:
+
+    def test_should_return_IHostMemory(self, identity_engine: trt.ICudaEngine) -> None:
+        # Precondition.
+        engine = identity_engine
+
+        # Under test.
+        buffer = buffer_from_engine(engine)
+
+        # Postcondition.
+        assert isinstance(buffer, trt.IHostMemory)
+
+    def test_should_content_match_engine(self, identity_engine: trt.ICudaEngine) -> None:
+        """Test that `BufferFromEngine` returns a buffer with the same content as the engine."""
+        # Precondition.
+        engine = identity_engine
+
+        # Under test.
+        buffer = buffer_from_engine(engine)
+
+        # Postcondition.
+        assert bytes(buffer) == bytes(engine.serialize())
+
+
+class TestSaveEngine:
+
+    def test_should_write_serialized_engine_to_file(self, identity_network: trt.ICudaEngine) -> None:
+        # Precondition.
+        with util.NamedTemporaryFile(mode="wb+") as out_file:
+            name = out_file.name
+            engine = engine_from_network(identity_network)
+
+            # Under test.
+            save_engine = SaveEngine(engine, path=out_file)
+            save_engine()
+            out_file.flush()
+
+            # Postcondition.
+            assert is_file_non_empty(out_file.name)
+            out_file.seek(0)
+            assert bytes(engine.serialize()) == bytes(out_file.read())
+
+
+class TestOnnxLikeFromNetwork:
+    @pytest.mark.parametrize(
+        "model_name",
+        [
+            "identity",
+            "empty_tensor_expand",
+            "const_foldable",
+            "and",
+            "scan",
+            "dim_param",
+            "tensor_attr",
+        ],
+    )
+    def test_onnx_like_from_network(self, model_name):
+        assert onnx_like_from_network(
+            NetworkFromOnnxBytes(ONNX_MODELS[model_name].loader)
+        )
+
+
+class TestDefaultPlugins:
+    @pytest.mark.skipif(
+        config.USE_TENSORRT_RTX,
+        reason="Plugin tests are not compatible with TensorRT-RTX"
+    )
+    def test_default_plugins(self):
+        network_loader = NetworkFromOnnxBytes(ONNX_MODELS["roialign"].loader)
+        engine_loader = EngineFromNetwork(network_loader)
+        engine = engine_loader()
+        assert engine is not None

tools/Polygraphy/tests/backend/trt/test_profile.py

@@ -0,0 +1,119 @@
+#
+# 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 config
+if config.USE_TENSORRT_RTX:
+    import tensorrt_rtx as trt
+else:
+    import tensorrt as trt
+from polygraphy.backend.trt import Profile, create_network, network_from_onnx_bytes
+from tests.models.meta import ONNX_MODELS
+
+
+@pytest.fixture(scope="session")
+def dynamic_identity_network():
+    builder, network, parser = network_from_onnx_bytes(
+        ONNX_MODELS["dynamic_identity"].loader
+    )
+    with builder, network, parser:
+        yield builder, network, parser
+
+
+class TestProfile:
+    def test_can_add(self):
+        profile = Profile()
+        min, opt, max = (1, 1), (2, 2), (4, 4)
+        assert profile.add("input", min=min, opt=opt, max=max) is profile
+        shape_tuple = profile["input"]
+        assert shape_tuple.min == min
+        assert shape_tuple.opt == opt
+        assert shape_tuple.max == max
+
+    def test_fill_defaults_does_not_overwrite(self, dynamic_identity_network):
+        _, network, _ = dynamic_identity_network
+        profile = Profile().add("X", (1, 1, 1, 1), (1, 1, 2, 2), (1, 1, 3, 3))
+
+        assert profile.fill_defaults(network) is profile
+        assert profile["X"].min == (1, 1, 1, 1)
+        assert profile["X"].opt == (1, 1, 2, 2)
+        assert profile["X"].max == (1, 1, 3, 3)
+
+    def test_fill_defaults_scalar_shape_tensor(self):
+        _, network = create_network()
+        fill_shape = network.add_input("fill_shape", shape=tuple(), dtype=trt.int32)
+
+        # Need to add some other operations so TensorRT treats `fill_shape` as a shape tensor.
+        if config.USE_TENSORRT_RTX:
+            fill = network.add_fill(tuple(), trt.FillOperation.LINSPACE, trt.int32)
+        else:
+            fill = network.add_fill(tuple(), trt.FillOperation.LINSPACE)
+        fill.set_input(0, fill_shape)
+        fill.set_input(
+            1,
+            network.add_constant(
+                shape=tuple(), weights=np.array(0).astype(np.int32)
+            ).get_output(0),
+        )
+        fill.set_input(
+            2,
+            network.add_constant(
+                shape=tuple(), weights=np.array(1).astype(np.int32)
+            ).get_output(0),
+        )
+
+        network.mark_output(fill.get_output(0))
+
+        assert fill_shape.is_shape_tensor
+
+        profile = Profile()
+        profile.fill_defaults(network)
+
+        assert profile[fill_shape.name].min == (1,)
+        assert profile[fill_shape.name].opt == (1,)
+        assert profile[fill_shape.name].max == (1,)
+
+    def test_to_trt(self, dynamic_identity_network):
+        builder, network, _ = dynamic_identity_network
+        profile = Profile().add("X", (1, 2, 1, 1), (1, 2, 2, 2), (1, 2, 4, 4))
+
+        trt_profile = profile.to_trt(builder, network)
+        trt_profile.get_shape("X") == ((1, 2, 1, 1), (1, 2, 2, 2), (1, 2, 4, 4))
+
+    @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"
+        ]
+        builder, network = create_network()
+        for input in match_case:
+            network.add_input(input, shape=(-1, 2, 3), dtype=trt.float32)
+
+        profile = Profile()
+        profile.add(name, min=(2, 2, 3), opt=(2, 2, 3), max=(2, 2, 3))
+        profile.fill_defaults(network)
+        trt_prof = profile.to_trt(builder, network)
+        res = [trt_prof.get_shape(case) == [(2, 2, 3), (2, 2, 3), (2, 2, 3)] for case in match_case]
+        assert res == should_match 

tools/Polygraphy/tests/backend/trt/test_runner.py

@@ -0,0 +1,408 @@
+#
+# 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 threading
+
+import numpy as np
+import pytest
+import torch
+
+from polygraphy import config, cuda, mod
+from polygraphy.backend.trt import (
+    EngineFromNetwork,
+    NetworkFromOnnxBytes,
+    Profile,
+    TrtRunner,
+    engine_from_network,
+    network_from_onnx_bytes,
+)
+from polygraphy.backend.trt.runner import _get_array_on_cpu
+from polygraphy.exception import PolygraphyException
+from polygraphy.logger import G_LOGGER
+from tests.models.meta import ONNX_MODELS
+
+# Import CreateConfigRTX conditionally for TensorRT-RTX builds
+if config.USE_TENSORRT_RTX:
+    import tensorrt_rtx as trt
+    from polygraphy.backend.tensorrt_rtx import CreateConfigRTX as CreateConfig
+else:
+    import tensorrt as trt
+    from polygraphy.backend.trt import CreateConfig
+
+
+class TestLoggerCallbacks:
+    @pytest.mark.parametrize("sev", G_LOGGER.SEVERITY_LETTER_MAPPING.keys())
+    def test_set_severity(self, sev):
+        G_LOGGER.module_severity = sev
+
+
+@pytest.fixture(scope="class")
+def nonzero_engine():
+    model = ONNX_MODELS["nonzero"]
+    network_loader = NetworkFromOnnxBytes(model.loader)
+    return engine_from_network(network_loader)
+
+
+@pytest.fixture()
+def identity_engine():
+    model = ONNX_MODELS["identity"]
+    network_loader = NetworkFromOnnxBytes(model.loader)
+    return engine_from_network(network_loader)
+
+
+@pytest.fixture()
+def reducable_engine():
+    model = ONNX_MODELS["reducable"]
+    network_loader = NetworkFromOnnxBytes(model.loader)
+    return engine_from_network(network_loader)
+
+
+class TestTrtRunner:
+    def test_can_name_runner(self):
+        NAME = "runner"
+        runner = TrtRunner(None, name=NAME)
+        assert runner.name == NAME
+
+    def test_basic(self, identity_engine):
+        with TrtRunner(identity_engine) as runner:
+            assert runner.optimization_profile is None
+            assert runner.is_active
+            ONNX_MODELS["identity"].check_runner(runner)
+            assert runner.last_inference_time() is not None
+        assert not runner.is_active
+
+    @pytest.mark.serial
+    @pytest.mark.skipif(config.USE_TENSORRT_RTX, reason="TensorRT-RTX has different warning output behavior")
+    def test_warn_if_impl_methods_called(self, check_warnings_on_runner_impl_methods, identity_engine):
+        runner = TrtRunner(identity_engine)
+        check_warnings_on_runner_impl_methods(runner)
+
+    @pytest.mark.parametrize(
+        "inp, expected",
+        [
+            ([1, 0, 1, 1], [[0, 2, 3]]),
+            ([1, 0, 0, 1], [[0, 3]]),
+            ([0, 0, 0, 1], [[3]]),
+        ],
+    )
+    @pytest.mark.skipif(config.USE_TENSORRT_RTX, reason="TensorRT-RTX does not support data dependent shapes")
+    def test_data_dependent_shapes(self, nonzero_engine, inp, expected):
+        with TrtRunner(nonzero_engine) as runner:
+            outputs = runner.infer(
+                {
+                    "input": np.array(
+                        inp,
+                        dtype=(np.int32 if mod.version(trt.__version__) < mod.version("9.0") else np.int64),
+                    )
+                }
+            )
+            assert np.array_equal(outputs["nonzero_out_0"], np.array(expected, dtype=np.int32))
+
+    @pytest.mark.parametrize("copy_outputs_to_host", [True, False])
+    @pytest.mark.parametrize("device", ["cpu", "cuda"])
+    def test_torch_tensors(self, copy_outputs_to_host, identity_engine, device):
+        with TrtRunner(identity_engine) as runner:
+            arr = torch.ones([1, 1, 2, 2], dtype=torch.float32, device=device)
+            outputs = runner.infer({"x": arr}, copy_outputs_to_host=copy_outputs_to_host)
+            assert all(isinstance(t, torch.Tensor) for t in outputs.values())
+
+            assert torch.equal(outputs["y"].to("cpu"), arr.to("cpu"))
+
+            assert outputs["y"].device.type == ("cpu" if copy_outputs_to_host else "cuda")
+
+    def test_context(self, identity_engine):
+        with TrtRunner(identity_engine.create_execution_context) as runner:
+            ONNX_MODELS["identity"].check_runner(runner)
+
+    def test_device_buffer_order_matches_bindings(self, reducable_engine):
+        with TrtRunner(reducable_engine) as runner:
+            dev_buf_order = list(runner.device_input_buffers.keys())
+            for binding, dev_buf_name in zip(reducable_engine, dev_buf_order):
+                assert binding == dev_buf_name
+
+    def test_shape_output(self):
+        model = ONNX_MODELS["reshape"]
+        engine = engine_from_network(NetworkFromOnnxBytes(model.loader))
+        with engine, TrtRunner(engine.create_execution_context) as runner:
+            model.check_runner(runner)
+
+    def test_multithreaded_runners_from_engine(self, identity_engine):
+        with TrtRunner(identity_engine) as runner0, TrtRunner(identity_engine) as runner1:
+            t1 = threading.Thread(target=ONNX_MODELS["identity"].check_runner, args=(runner0,))
+            t2 = threading.Thread(target=ONNX_MODELS["identity"].check_runner, args=(runner1,))
+            t1.start()
+            t2.start()
+            t1.join()
+            t2.join()
+
+    @pytest.mark.parametrize("use_optimization_profile", [True, False])
+    @pytest.mark.skipif(
+        not config.USE_TENSORRT_RTX
+        and mod.version(trt.__version__) >= mod.version("8.6")
+        and mod.version(trt.__version__) < mod.version("8.7"),
+        reason="Bug in TRT 8.6",
+    )
+    def test_multiple_profiles(self, use_optimization_profile):
+        model = ONNX_MODELS["dynamic_identity"]
+        profile0_shapes = [
+            (1, 2, 1, 1),
+            (1, 2, 1, 1),
+            (1, 2, 1, 1),
+        ]  # Use min==opt==max to fix shapes in the engine.
+        profile1_shapes = [(1, 2, 1, 1), (1, 2, 2, 2), (1, 2, 4, 4)]
+        profile2_shapes = [(1, 2, 4, 4), (1, 2, 8, 8), (1, 2, 16, 16)]
+        network_loader = NetworkFromOnnxBytes(model.loader)
+        profiles = [
+            Profile().add("X", *profile0_shapes),
+            Profile().add("X", *profile1_shapes),
+            Profile().add("X", *profile2_shapes),
+        ]
+        config_loader = CreateConfig(profiles=profiles)
+        engine = engine_from_network(network_loader, config_loader)
+        context = engine.create_execution_context()
+
+        for index, shapes in enumerate([profile0_shapes, profile1_shapes, profile2_shapes]):
+            with TrtRunner(
+                context,
+                optimization_profile=index if use_optimization_profile else None,
+            ) as runner:
+                if not use_optimization_profile:
+                    runner.set_profile(index)
+
+                assert runner.context.active_optimization_profile == index
+                for shape in shapes:
+                    model.check_runner(runner, {"X": shape})
+
+    @pytest.mark.skipif(
+        not config.USE_TENSORRT_RTX and mod.version(trt.__version__) < mod.version("10.0"),
+        reason="Feature not present before 10.0",
+    )
+    @pytest.mark.parametrize("allocation_strategy", [None, "static", "profile", "runtime"])
+    def test_allocation_strategies(self, allocation_strategy):
+        if config.USE_TENSORRT_RTX and allocation_strategy == "runtime":
+            pytest.skip("TensorRT-RTX issues with runtime allocation strategy")
+
+        model = ONNX_MODELS["residual_block"]
+        profile0_shapes = [(1, 3, 224, 224), (1, 3, 224, 224), (1, 3, 224, 224)]
+        profile1_shapes = [(1, 3, 224, 224), (1, 3, 224, 224), (2, 3, 224, 224)]
+        profile2_shapes = [(1, 3, 224, 224), (1, 3, 224, 224), (4, 3, 224, 224)]
+        network_loader = NetworkFromOnnxBytes(model.loader)
+        profiles = [
+            Profile().add("gpu_0/data_0", *profile0_shapes),
+            Profile().add("gpu_0/data_0", *profile1_shapes),
+            Profile().add("gpu_0/data_0", *profile2_shapes),
+        ]
+        config_loader = CreateConfig(profiles=profiles)
+        engine = engine_from_network(network_loader, config_loader)
+
+        for index, shapes in enumerate([profile0_shapes, profile1_shapes, profile2_shapes]):
+            with TrtRunner(
+                engine,
+                optimization_profile=index,
+                allocation_strategy=allocation_strategy,
+            ) as runner:
+                for shape in shapes:
+                    model.check_runner(runner, {"gpu_0/data_0": shape})
+
+    def test_empty_tensor_with_dynamic_input_shape_tensor(self):
+        model = ONNX_MODELS["empty_tensor_expand"]
+        shapes = [(1, 2, 0, 3, 0), (2, 2, 0, 3, 0), (4, 2, 0, 3, 0)]
+        network_loader = NetworkFromOnnxBytes(model.loader)
+        profiles = [Profile().add("new_shape", *shapes)]
+        config_loader = CreateConfig(profiles=profiles)
+
+        with TrtRunner(EngineFromNetwork(network_loader, config_loader)) as runner:
+            for shape in shapes:
+                model.check_runner(runner, {"new_shape": shape})
+
+    @pytest.mark.parametrize(
+        "names, err",
+        [
+            (["fake-input", "x"], "Extra inputs in"),
+            (["fake-input"], "The following inputs were not found"),
+            ([], "The following inputs were not found"),
+        ],
+    )
+    @pytest.mark.parametrize("module", [torch, np])
+    def test_error_on_wrong_name_feed_dict(self, names, err, identity_engine, module):
+        with TrtRunner(identity_engine) as runner:
+            with pytest.raises(PolygraphyException, match=err):
+                runner.infer({name: module.ones((1, 1, 2, 2), dtype=module.float32) for name in names})
+
+    @pytest.mark.parametrize("module", [torch, np])
+    def test_error_on_wrong_dtype_feed_dict(self, identity_engine, module):
+        with TrtRunner(identity_engine) as runner:
+            with pytest.raises(PolygraphyException, match="unexpected dtype."):
+                runner.infer({"x": module.ones((1, 1, 2, 2), dtype=module.int32)})
+
+    @pytest.mark.parametrize("module", [torch, np])
+    def test_error_on_wrong_shape_feed_dict(self, identity_engine, module):
+        with TrtRunner(identity_engine) as runner:
+            with pytest.raises(PolygraphyException, match="incompatible shape."):
+                runner.infer({"x": module.ones((1, 1, 3, 2), dtype=module.float32)})
+
+    @pytest.mark.parametrize("use_view", [True, False])  # We should be able to use DeviceArray in place of DeviceView
+    def test_device_views(self, use_view, reducable_engine):
+        with TrtRunner(reducable_engine) as runner, cuda.DeviceArray((1,), dtype=np.float32) as x:
+            x.copy_from(np.ones((1,), dtype=np.float32))
+            outputs = runner.infer(
+                {
+                    "X0": x.view() if use_view else x,
+                    "Y0": np.ones((1,), dtype=np.float32),
+                }
+            )
+            assert outputs["identity_out_6"][0] == 2
+            assert outputs["identity_out_8"][0] == 2
+
+    def test_no_output_copy(self, identity_engine):
+        with TrtRunner(identity_engine) as runner:
+            inp = np.ones(shape=(1, 1, 2, 2), dtype=np.float32)
+            outputs = runner.infer({"x": inp}, copy_outputs_to_host=False)
+            assert isinstance(outputs["y"], cuda.DeviceView)
+            assert np.array_equal(outputs["y"].numpy(), inp)
+
+    def test_subsequent_infers_with_different_input_types(self, identity_engine):
+        with TrtRunner(identity_engine) as runner:
+            inp = np.ones(shape=(1, 1, 2, 2), dtype=np.float32)
+
+            def check(outputs):
+                assert np.all(outputs["y"] == inp)
+
+            check(runner.infer({"x": inp}))
+            check(runner.infer({"x": cuda.DeviceArray(shape=inp.shape, dtype=inp.dtype).copy_from(inp)}))
+
+            torch_outputs = runner.infer({"x": torch.from_numpy(inp)})
+            check({name: out.numpy() for name, out in torch_outputs.items()})
+            check(runner.infer({"x": inp}))
+
+    @pytest.mark.parametrize("use_view", [True, False])  # We should be able to use DeviceArray in place of DeviceView
+    def test_device_view_dynamic_shapes(self, use_view):
+        model = ONNX_MODELS["dynamic_identity"]
+        profiles = [
+            Profile().add("X", (1, 2, 1, 1), (1, 2, 2, 2), (1, 2, 4, 4)),
+        ]
+        runner = TrtRunner(EngineFromNetwork(NetworkFromOnnxBytes(model.loader), CreateConfig(profiles=profiles)))
+        with runner, cuda.DeviceArray(shape=(1, 2, 3, 3), dtype=np.float32) as arr:
+            inp = np.random.random_sample(size=(1, 2, 3, 3)).astype(np.float32)
+            arr.copy_from(inp)
+            outputs = runner.infer({"X": (cuda.DeviceView(arr.ptr, arr.shape, arr.dtype) if use_view else arr)})
+            assert np.all(outputs["Y"] == inp)
+            assert outputs["Y"].shape == (1, 2, 3, 3)
+
+    def test_cannot_use_device_view_shape_tensor(self):
+        model = ONNX_MODELS["empty_tensor_expand"]
+        with TrtRunner(EngineFromNetwork(NetworkFromOnnxBytes(model.loader))) as runner, cuda.DeviceArray(
+            shape=(5,),
+            dtype=(
+                np.int32
+                if mod.version(trt.__version__) < mod.version("9.0") and not config.USE_TENSORRT_RTX
+                else np.int64
+            ),
+        ) as arr:
+            with pytest.raises(PolygraphyException, match="it must reside in host memory"):
+                runner.infer({"data": np.ones((2, 0, 3, 0), dtype=np.float32), "new_shape": arr})
+
+    @pytest.mark.parametrize("hwc_input", [True, False], ids=["hwc_input", "chw_input"])
+    @pytest.mark.parametrize("hwc_output", [True, False], ids=["hwc_output", "chw_output"])
+    @pytest.mark.skipif(config.USE_TENSORRT_RTX, reason="TensorRT-RTX does not support custom I/O format networks")
+    def test_infer_chw_format(self, hwc_input, hwc_output):
+        model = ONNX_MODELS["identity_multi_ch"]
+        inp_shape = model.input_metadata["x"].shape
+        builder, network, parser = network_from_onnx_bytes(model.loader)
+
+        formats = 1 << int(trt.TensorFormat.HWC)
+        if hwc_input:
+            network.get_input(0).allowed_formats = formats
+        if hwc_output:
+            network.get_output(0).allowed_formats = formats
+
+        engine = engine_from_network((builder, network))
+
+        with TrtRunner(engine) as runner:
+            inp = np.random.normal(size=(inp_shape)).astype(np.float32)
+            if hwc_input:
+                inp = inp.transpose(0, 2, 3, 1)
+
+            outputs = runner.infer({"x": inp})
+            if hwc_input == hwc_output:  # output in CHW/HWC format and similarly shaped
+                assert np.allclose(outputs["y"], inp)
+            elif not hwc_input and hwc_output:  # output in HWC format and shaped (N, H, W, C)
+                assert np.allclose(outputs["y"].transpose(0, 3, 1, 2), inp)
+            else:  # hwc_input and not hwc_output: output in CHW format and shaped (N, C, H, W)
+                assert np.allclose(outputs["y"].transpose(0, 2, 3, 1), inp)
+
+    @pytest.mark.parametrize("use_torch", [True, False])
+    def test_get_array_on_cpu(self, use_torch):
+        shape = (4,)
+        with cuda.DeviceArray.raw(shape) as arr:
+            host_buffers = {}
+            stream = cuda.Stream()
+            host_arr = _get_array_on_cpu(arr, "test", host_buffers, stream, arr.nbytes, use_torch)
+
+            if use_torch:
+                assert isinstance(host_arr, torch.Tensor)
+            else:
+                assert isinstance(host_arr, np.ndarray)
+
+    @pytest.mark.skipif(
+        mod.version(trt.__version__) < mod.version("10.0") and not config.USE_TENSORRT_RTX,
+        reason="Feature not present before 10.0",
+    )
+    @pytest.mark.parametrize("budget", [None, -2, -1, 0, 0.5, 0.99, 1.0, 1000, np.inf])
+    def test_weight_streaming(self, budget):
+        model = ONNX_MODELS["matmul_2layer"]
+        network_loader = NetworkFromOnnxBytes(model.loader, strongly_typed=True)
+        config_loader = CreateConfig(weight_streaming=True)
+        engine = engine_from_network(network_loader, config_loader)
+
+        if budget == np.inf:
+            # set to max size - 1
+            budget = engine.streamable_weights_size - 1
+
+        kwargs = {"weight_streaming_budget": None, "weight_streaming_percent": None}
+        if budget is not None:
+            if 0 < budget <= 1:
+                kwargs["weight_streaming_percent"] = budget * 100
+            else:
+                kwargs["weight_streaming_budget"] = int(budget)
+
+        with TrtRunner(engine, optimization_profile=0, **kwargs) as runner:
+            model.check_runner(runner)
+
+    @pytest.mark.skipif(not config.USE_TENSORRT_RTX, reason="TensorRT-RTX not enabled")
+    def test_compute_capabilities_engine_building(self):
+        """Test compute capabilities integration with engine building"""
+        model = ONNX_MODELS["identity"]
+        network_loader = NetworkFromOnnxBytes(model.loader)
+
+        # Test --use-gpu flag
+        config_loader = CreateConfig(use_gpu=True)
+        engine = engine_from_network(network_loader, config_loader)
+        with TrtRunner(engine) as runner:
+            model.check_runner(runner)
+
+        # Test --compute-capabilities flag
+        config_loader = CreateConfig(compute_capabilities=[(7, 5), (8, 0), (8, 6)])
+        engine = engine_from_network(network_loader, config_loader)
+        with TrtRunner(engine) as runner:
+            model.check_runner(runner)
+
+    @pytest.mark.skipif(not config.USE_TENSORRT_RTX, reason="TensorRT-RTX not enabled")
+    def test_compute_capabilities_mutual_exclusion(self):
+        """Test that use_gpu and compute_capabilities are mutually exclusive"""
+        # Test mutual exclusion - should raise an exception
+        with pytest.raises(PolygraphyException, match="use_gpu and compute_capabilities are mutually exclusive"):
+            CreateConfig(use_gpu=True, compute_capabilities=[(7, 5)])

tools/Polygraphy/tests/backend/trt/test_util.py

@@ -0,0 +1,322 @@
+#
+# 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
+from textwrap import dedent
+
+import pytest
+
+from polygraphy import config, mod
+from polygraphy.backend.trt import Profile, create_network
+from polygraphy.backend.trt import util as trt_util
+
+# Import CreateConfigRTX conditionally for TensorRT-RTX builds
+if config.USE_TENSORRT_RTX:
+    import tensorrt_rtx as trt
+    from polygraphy.backend.tensorrt_rtx import CreateConfigRTX as CreateConfig
+else:
+    import tensorrt as trt
+    from polygraphy.backend.trt import CreateConfig
+
+
+@pytest.fixture(scope="session")
+def dummy_network():
+    builder, network = create_network()
+    network.add_input("X", dtype=trt.float32, shape=[-1])
+    with builder, network:
+        yield builder, network
+
+
+@pytest.fixture(scope="session")
+def layer_class_mapping():
+    return trt_util.get_layer_class_mapping()
+
+
+@pytest.mark.parametrize("layer_type", trt.LayerType.__members__.values())
+def test_all_layer_types_mapped(layer_class_mapping, layer_type):
+    waived_layers = [trt.LayerType.PLUGIN]
+    with contextlib.suppress(AttributeError):
+        waived_layers.append(trt.LayerType.PLUGIN_V3)
+    if layer_type in waived_layers:
+        pytest.skip("PLUGIN has no corresponding ILayer")
+    assert layer_type in layer_class_mapping
+
+
+# Can't use pytest.skip because we can't construct the test unless trt.MemoryPoolType exists.
+
+
+def adjust_memory_pool_limits_after_8_6(limits):
+    # Adjust tactic DRAM so we can match the output text reliably in update_expected_output.
+    if mod.version(trt.__version__) >= mod.version("8.6") or config.USE_TENSORRT_RTX:
+        limits[trt.MemoryPoolType.TACTIC_DRAM] = 1 << 30
+    return limits
+
+
+def update_expected_output(expected):
+    is_trt_10_plus = (
+        mod.version(trt.__version__) >= mod.version("10.0") or 
+        config.USE_TENSORRT_RTX
+    )
+    is_trt_8_6_plus = (
+        mod.version(trt.__version__) >= mod.version("8.6") or 
+        config.USE_TENSORRT_RTX
+    )
+    is_trt_8_7_plus = (
+        mod.version(trt.__version__) >= mod.version("8.7") or 
+        config.USE_TENSORRT_RTX
+    )
+    
+    if is_trt_8_6_plus:
+        if is_trt_10_plus:
+            expected = expected.replace(
+                "MiB]",
+                "MiB, TACTIC_DRAM: 1024.00 MiB, TACTIC_SHARED_MEMORY: 1024.00 MiB]",
+            )
+        else:
+            expected = expected.replace("MiB]", "MiB, TACTIC_DRAM: 1024.00 MiB]")
+
+        if "Preview Features" not in expected:
+            if not is_trt_10_plus:
+                expected = (
+                    dedent(expected).strip()
+                    + "\nPreview Features       | [FASTER_DYNAMIC_SHAPES_0805, DISABLE_EXTERNAL_TACTIC_SOURCES_FOR_CORE_0805]"
+                )
+            else:
+                preview_features = "[PROFILE_SHARING_0806"
+                if config.USE_TENSORRT_RTX:
+                    preview_features += ", RUNTIME_ACTIVATION_RESIZE_10_10"
+                preview_features += "]"
+                expected = (
+                    dedent(expected).strip()
+                    + f"\nPreview Features       | {preview_features}"
+                )
+
+    if is_trt_8_7_plus:
+        # CUBLAS_LT is not longer enabled by default
+        expected = expected.replace("CUBLAS_LT, ", "")
+
+    if is_trt_10_plus:
+        expected = expected.replace(
+            "EngineCapability.DEFAULT", "EngineCapability.STANDARD"
+        )
+        expected = expected.replace("CUBLAS, ", "")
+        expected = expected.replace("CUDNN, ", "")
+
+    return expected
+
+
+@pytest.mark.parametrize(
+    "create_config, expected",
+    # NOTE: We set workspace sizes here so we can have predictable output
+    [
+        (
+            CreateConfig(
+                memory_pool_limits=adjust_memory_pool_limits_after_8_6(
+                    {trt.MemoryPoolType.WORKSPACE: 16 << 20}
+                )
+            ),
+            update_expected_output(
+                """
+                Flags                  | [{}]
+                Engine Capability      | EngineCapability.DEFAULT
+                Memory Pools           | [WORKSPACE: 16.00 MiB]
+                Tactic Sources         | [CUBLAS, CUBLAS_LT, CUDNN, EDGE_MASK_CONVOLUTIONS, JIT_CONVOLUTIONS]
+                Profiling Verbosity    | ProfilingVerbosity.DETAILED
+                """.format("TF32" if config.USE_TENSORRT_RTX else "")
+            ),
+        ),
+        (
+            CreateConfig(
+                memory_pool_limits=adjust_memory_pool_limits_after_8_6(
+                    {trt.MemoryPoolType.WORKSPACE: 16 << 20}
+                ),
+                tactic_sources=[],
+            ),
+            update_expected_output(
+                """
+                Flags                  | [{}]
+                Engine Capability      | EngineCapability.DEFAULT
+                Memory Pools           | [WORKSPACE: 16.00 MiB]
+                Tactic Sources         | []
+                Profiling Verbosity    | ProfilingVerbosity.DETAILED
+                """.format("TF32" if config.USE_TENSORRT_RTX else "")
+            ),
+        ),
+        (
+            CreateConfig(
+                memory_pool_limits=adjust_memory_pool_limits_after_8_6(
+                    {trt.MemoryPoolType.WORKSPACE: 4 << 20}
+                )
+            ),
+            update_expected_output(
+                """
+                Flags                  | [{}]
+                Engine Capability      | EngineCapability.DEFAULT
+                Memory Pools           | [WORKSPACE: 4.00 MiB]
+                Tactic Sources         | [CUBLAS, CUBLAS_LT, CUDNN, EDGE_MASK_CONVOLUTIONS, JIT_CONVOLUTIONS]
+                Profiling Verbosity    | ProfilingVerbosity.DETAILED
+                """.format("TF32" if config.USE_TENSORRT_RTX else "")
+            ),
+        ),
+        (
+            CreateConfig(
+                memory_pool_limits=adjust_memory_pool_limits_after_8_6(
+                    {trt.MemoryPoolType.WORKSPACE: 16 << 20}
+                ),
+                **({} if config.USE_TENSORRT_RTX else {
+                    "fp16": True,
+                    "int8": True,
+                    "tf32": True,
+                }),
+                refittable=True,
+                precision_constraints="obey",
+            ),
+            update_expected_output(
+                """
+                Flags                  | [{}REFIT, TF32, OBEY_PRECISION_CONSTRAINTS]
+                Engine Capability      | EngineCapability.DEFAULT
+                Memory Pools           | [WORKSPACE: 16.00 MiB]
+                Tactic Sources         | [CUBLAS, CUBLAS_LT, CUDNN, EDGE_MASK_CONVOLUTIONS, JIT_CONVOLUTIONS]
+                Profiling Verbosity    | ProfilingVerbosity.DETAILED
+                """.format(
+                    "" if config.USE_TENSORRT_RTX else "FP16, INT8, ",
+                )
+            ),
+        ),
+        (
+            CreateConfig(
+                memory_pool_limits=adjust_memory_pool_limits_after_8_6(
+                    {trt.MemoryPoolType.WORKSPACE: 16 << 20}
+                ),
+                profiles=[
+                    Profile().add("X", [1], [1], [1]),
+                    Profile().add("X", [2], [2], [2]),
+                ],
+            ),
+            update_expected_output(
+                """
+                Flags                  | [{}]
+                Engine Capability      | EngineCapability.DEFAULT
+                Memory Pools           | [WORKSPACE: 16.00 MiB]
+                Tactic Sources         | [CUBLAS, CUBLAS_LT, CUDNN, EDGE_MASK_CONVOLUTIONS, JIT_CONVOLUTIONS]
+                Profiling Verbosity    | ProfilingVerbosity.DETAILED
+                Optimization Profiles  | 2 profile(s)
+                """.format("TF32" if config.USE_TENSORRT_RTX else "")
+            ),
+        ),
+    ] + ([] if config.USE_TENSORRT_RTX else [
+        (
+            CreateConfig(
+                memory_pool_limits=adjust_memory_pool_limits_after_8_6(
+                    {trt.MemoryPoolType.WORKSPACE: 16 << 20}
+                ),
+                use_dla=True,
+            ),
+            update_expected_output(
+                """
+                Flags                  | []
+                Engine Capability      | EngineCapability.DEFAULT
+                Memory Pools           | [WORKSPACE: 16.00 MiB, DLA_MANAGED_SRAM: 0.00 MiB, DLA_LOCAL_DRAM: 1024.00 MiB, DLA_GLOBAL_DRAM: 512.00 MiB]
+                Tactic Sources         | [CUBLAS, CUBLAS_LT, CUDNN, EDGE_MASK_CONVOLUTIONS, JIT_CONVOLUTIONS]
+                DLA                    | Default Device Type: DeviceType.DLA, Core: -1
+                Profiling Verbosity    | ProfilingVerbosity.DETAILED
+                """
+            ),
+        ),
+    ]) + [
+        (
+            (
+                CreateConfig(
+                    memory_pool_limits=adjust_memory_pool_limits_after_8_6(
+                        {trt.MemoryPoolType.WORKSPACE: 16 << 20}
+                    ),
+                    preview_features=[trt.PreviewFeature.PROFILE_SHARING_0806],
+                ),
+                update_expected_output(
+                    """
+                Flags                  | [{}]
+                Engine Capability      | EngineCapability.DEFAULT
+                Memory Pools           | [WORKSPACE: 16.00 MiB]
+                Tactic Sources         | [CUBLAS, CUBLAS_LT, CUDNN, EDGE_MASK_CONVOLUTIONS, JIT_CONVOLUTIONS]
+                Profiling Verbosity    | ProfilingVerbosity.DETAILED
+                Preview Features       | [PROFILE_SHARING_0806]
+                """.format("TF32" if config.USE_TENSORRT_RTX else "")
+                ),
+            )
+            if mod.version(trt.__version__) >= mod.version("10.0") or config.USE_TENSORRT_RTX
+            else (
+                CreateConfig(
+                    memory_pool_limits=adjust_memory_pool_limits_after_8_6(
+                        {trt.MemoryPoolType.WORKSPACE: 16 << 20}
+                    ),
+                    preview_features=(
+                        [trt.PreviewFeature.ALIASED_PLUGIN_IO_10_03] 
+                        if config.USE_TENSORRT_RTX 
+                        else [trt.PreviewFeature.FASTER_DYNAMIC_SHAPES_0805]
+                    ),
+                ),
+                update_expected_output(
+                    """
+                Flags                  | [{}]
+                Engine Capability      | EngineCapability.DEFAULT
+                Memory Pools           | [WORKSPACE: 16.00 MiB]
+                Tactic Sources         | [CUBLAS, CUBLAS_LT, CUDNN, EDGE_MASK_CONVOLUTIONS, JIT_CONVOLUTIONS]
+                Profiling Verbosity    | ProfilingVerbosity.DETAILED
+                Preview Features       | [{}]
+                """.format(
+                    "TF32" if config.USE_TENSORRT_RTX else "",
+                    "ALIASED_PLUGIN_IO_10_03" if config.USE_TENSORRT_RTX else "FASTER_DYNAMIC_SHAPES_0805"
+                )
+                ),
+            )
+        ),
+    ],
+    ids=[
+        "default",
+        "tactic-sources",
+        "memory-pool-limits",
+        "builder-flags" + ("-rtx" if config.USE_TENSORRT_RTX else ""),
+        "profiles",
+    ] + ([] if config.USE_TENSORRT_RTX else ["dla"]) + [
+        "preview-features",
+    ],
+)
+def test_str_from_config(create_config, expected, dummy_network):
+    config = create_config(*dummy_network)
+    actual = trt_util.str_from_config(config, dummy_network)
+    expected = dedent(expected).strip()
+    assert actual == expected
+
+
+def test_get_all_tensors_layer_with_null_inputs():
+    builder, network = create_network()
+    with builder, network:
+        inp = network.add_input("input", shape=(1, 3, 224, 224), dtype=trt.float32)
+        slice_layer = network.add_slice(
+            inp, (0, 0, 0, 0), (1, 3, 224, 224), (1, 1, 1, 1)
+        )
+
+        # Set a tensor for `stride` to increment `num_inputs` so we have some inputs
+        # which are `None` in between.
+        slice_layer.set_input(3, inp)
+        assert slice_layer.num_inputs == 4
+
+        slice = slice_layer.get_output(0)
+        slice.name = "Slice"
+        network.mark_output(slice)
+
+        assert trt_util.get_all_tensors(network) == {"input": inp, "Slice": slice}

tools/Polygraphy/tests/common/test_datatype.py

@@ -0,0 +1,183 @@
+#
+# 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 onnx
+import pytest
+import tensorrt as trt
+import torch
+
+from polygraphy import mod, util
+from polygraphy.datatype import DataType, DataTypeEntry
+
+DATATYPES = DataType.__members__.values()
+
+
+class TestDataType:
+    def compare_names(self, name, expected_name, replace_map):
+        # Names may not match up exactly, so use the replace_map to make adjustments to the
+        # foreign type before comparing against the Polygraphy type
+        for old, new in replace_map.items():
+            if name == old:
+                name = new
+        assert name == expected_name
+
+    @pytest.mark.parametrize("dtype", DATATYPES, ids=str)
+    def test_numpy(self, dtype):
+        if dtype in [
+            DataType.BFLOAT16,
+            DataType.FLOAT8E4M3FN,
+            DataType.FLOAT8E4M3FNUZ,
+            DataType.FLOAT8E5M2,
+            DataType.FLOAT8E5M2FNUZ,
+            DataType.INT4,
+            DataType.FLOAT4,
+        ]:
+            pytest.xfail("Type not supported by NumPy")
+
+        np_type = dtype.numpy()
+        assert DataType.to_dtype(dtype, "numpy") == np_type
+        assert np_type.itemsize == dtype.itemsize
+        self.compare_names(np_type.name, dtype.name, {"str": "string"})
+
+        assert isinstance(np_type, np.dtype)
+
+        assert DataType.from_dtype(np_type) == dtype
+
+    @pytest.mark.parametrize("dtype", DATATYPES, ids=str)
+    def test_onnxrt(self, dtype):
+        if dtype in [
+            DataType.INT4,
+            DataType.FLOAT4,
+        ]:
+            pytest.skip("Type not supported by ONNX-RT")
+
+        onnxrt_type = DataType.to_dtype(dtype, "onnxruntime")
+        assert dtype.onnxruntime() == onnxrt_type
+
+        assert isinstance(onnxrt_type, str)
+
+        self.compare_names(
+            onnxrt_type.replace("tensor(", "").replace(")", ""),
+            dtype.name,
+            {
+                "double": "float64",
+                "float": "float32",
+            },
+        )
+
+        assert DataType.from_dtype(onnxrt_type) == dtype
+
+    @pytest.mark.parametrize("dtype", DATATYPES, ids=str)
+    def test_onnx(self, dtype):
+        if dtype in [
+            DataType.INT4,
+            DataType.FLOAT4,
+        ]:
+            pytest.skip("Type not supported by ONNX")
+
+        onnx_type = dtype.onnx()
+        assert DataType.to_dtype(dtype, "onnx") == onnx_type
+
+        assert isinstance(onnx_type, int)
+
+        onnx_type_map = util.invert_dict(dict(onnx.TensorProto.DataType.items()))
+        self.compare_names(
+            onnx_type_map[onnx_type].lower(),
+            dtype.name,
+            {
+                "double": "float64",
+                "float": "float32",
+            },
+        )
+
+        assert DataType.from_dtype(onnx_type) == dtype
+
+    @pytest.mark.skipif(
+        mod.version(trt.__version__) < mod.version("8.7"),
+        reason="Unsupported before TRT 8.7",
+    )
+    @pytest.mark.parametrize("dtype", DATATYPES, ids=str)
+    def test_tensorrt(self, dtype):
+        unsupported_types = [
+            DataType.FLOAT64,
+            DataType.INT16,
+            DataType.UINT16,
+            DataType.UINT32,
+            DataType.UINT64,
+            DataType.STRING,
+            DataType.INT64,
+            DataType.FLOAT8E4M3FNUZ,
+            DataType.FLOAT8E5M2,
+            DataType.FLOAT8E5M2FNUZ,
+        ]
+        if  mod.version(trt.__version__) < mod.version("10.8"):
+            unsupported_types.append(DataType.FLOAT4)
+        if dtype in unsupported_types:
+            pytest.xfail("Type not supported by TensorRT")
+
+        tensorrt_dtype = dtype.tensorrt()
+        assert DataType.to_dtype(dtype, "tensorrt") == tensorrt_dtype
+
+        assert isinstance(tensorrt_dtype, trt.DataType)
+
+        self.compare_names(
+            tensorrt_dtype.name.lower(),
+            dtype.name,
+            {
+                "double": "float64",
+                "float": "float32",
+                "half": "float16",
+                "fp8": "float8e4m3fn",
+                "fp4": "float4",
+                "bf16": "bfloat16",
+            },
+        )
+
+        assert DataType.from_dtype(tensorrt_dtype) == dtype
+
+    @pytest.mark.parametrize("trt_dtype", trt.DataType.__members__.values())
+    def test_all_tensorrt_types_supported(self, trt_dtype):
+        dtype = DataType.from_dtype(trt_dtype, "tensorrt")
+        assert isinstance(dtype, DataTypeEntry)
+
+        assert dtype.tensorrt() == trt_dtype
+
+    @pytest.mark.parametrize("dtype", DATATYPES, ids=str)
+    def test_torch(self, dtype):
+        if dtype in [
+            DataType.FLOAT8E4M3FN,
+            DataType.FLOAT8E4M3FNUZ,
+            DataType.FLOAT8E5M2,
+            DataType.FLOAT8E5M2FNUZ,
+            DataType.UINT16,
+            DataType.UINT32,
+            DataType.UINT64,
+            DataType.STRING,
+            DataType.INT4,
+            DataType.FLOAT4,
+        ]:
+            pytest.xfail("Type not supported by Torch")
+
+        torch_type = dtype.torch()
+        assert DataType.to_dtype(dtype, "torch") == torch_type
+
+        assert isinstance(torch_type, torch.dtype)
+
+        self.compare_names(str(torch_type).replace("torch.", ""), dtype.name, {})
+
+        assert DataType.from_dtype(torch_type) == dtype

tools/Polygraphy/tests/common/test_interface.py

@@ -0,0 +1,70 @@
+#
+# 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 pytest
+
+from polygraphy.common.interface import TypedDict, TypedList
+from polygraphy.exception import PolygraphyException
+
+
+@pytest.fixture()
+def int_to_float():
+    class IntToFloat(TypedDict(lambda: int, lambda: float)):
+        pass
+
+    return IntToFloat()
+
+
+class TestTypedDict:
+    def test_wrong_type_set_item_value(self, int_to_float):
+        with pytest.raises(PolygraphyException, match="Unsupported value type"):
+            int_to_float[0] = "hi"
+
+    def test_wrong_type_set_item_key(self, int_to_float):
+        with pytest.raises(PolygraphyException, match="Unsupported key type"):
+            int_to_float["hi"] = 1.0
+
+    def test_wrong_type_update(self, int_to_float):
+        with pytest.raises(PolygraphyException, match="Unsupported key type"):
+            int_to_float.update({"hi": 1.0})
+
+
+@pytest.fixture()
+def ints():
+    class Ints(TypedList(lambda: int)):
+        pass
+
+    return Ints()
+
+
+class TestTypedList:
+    def test_wrong_type_append(self, ints):
+        with pytest.raises(PolygraphyException, match="Unsupported element type"):
+            ints.append(1.0)
+
+    def test_wrong_type_extend(self, ints):
+        with pytest.raises(PolygraphyException, match="Unsupported element type"):
+            ints.extend([0, 1, 2, 3, "surprise"])
+
+    def test_wrong_type_iadd(self, ints):
+        with pytest.raises(PolygraphyException, match="Unsupported element type"):
+            ints += [0, 1.0]
+
+    def test_wrong_type_setitem(self, ints):
+        ints.append(0)
+        with pytest.raises(PolygraphyException, match="Unsupported element type"):
+            ints[0] = 1.0

tools/Polygraphy/tests/common/test_struct.py

@@ -0,0 +1,36 @@
+#
+# 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 polygraphy.common import TensorMetadata
+from polygraphy.datatype import DataType
+
+
+class TestTensorMetadata:
+    def test_str(self):
+        meta = TensorMetadata().add("X", dtype=DataType.FLOAT32, shape=(64, 64))
+        assert str(meta) == "{X [dtype=float32, shape=(64, 64)]}"
+
+    def test_str_no_dtype(self):
+        meta = TensorMetadata().add("X", dtype=None, shape=(64, 64))
+        assert str(meta) == "{X [shape=(64, 64)]}"
+
+    def test_str_no_shape(self):
+        meta = TensorMetadata().add("X", dtype=DataType.FLOAT32, shape=None)
+        assert str(meta) == "{X [dtype=float32]}"
+
+    def test_str_no_meta(self):
+        meta = TensorMetadata().add("X", dtype=None, shape=None)
+        assert str(meta) == "{X}"

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())

tools/Polygraphy/tests/conftest.py

@@ -0,0 +1,215 @@
+#
+# 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 copy
+import ctypes.util
+import glob
+import logging
+import os
+import subprocess as sp
+import sys
+
+import pytest
+
+from tests.helper import ROOT_DIR
+
+
+@pytest.fixture()
+def sandboxed_install_run(virtualenv, script_runner):
+    """
+    A special fixture that runs commands, but sandboxes any `pip install`s in a virtual environment.
+    Packages from the test environment are still usable, but those in the virtual environment take precedence
+    """
+
+    VENV_PYTHONPATH = glob.glob(
+        os.path.join(virtualenv.virtualenv, "lib", "python*", "site-packages")
+    )[0]
+
+    class StatusWrapper:
+        def __init__(self, stdout=None, stderr=None, success=None) -> None:
+            self.stdout = stdout
+            self.stderr = stderr
+            self.success = success
+
+    def run_impl(command, cwd=None):
+        env = copy.copy(os.environ)
+        # Always prioritize our own copy of Polygraphy over anything in the venv.
+        env["PYTHONPATH"] = ROOT_DIR + os.pathsep + VENV_PYTHONPATH
+
+        print(f"Running command: {' '.join(command)}")
+
+        status = StatusWrapper()
+        if "pip" in command:
+            virtualenv.run(command, cwd=cwd)
+            status.success = True
+        elif command[0] == "polygraphy":
+            sr_status = script_runner.run(*command, cwd=cwd, env=env)
+            status.stdout = sr_status.stdout
+            status.success = sr_status.success
+        else:
+            sp_status = sp.run(
+                command, cwd=cwd, env=env, stdout=sp.PIPE, stderr=sp.PIPE
+            )
+
+            def try_decode(inp):
+                try:
+                    return inp.decode()
+                except UnicodeDecodeError:
+                    return inp
+
+            status.stdout = try_decode(sp_status.stdout)
+            status.stderr = try_decode(sp_status.stderr)
+            status.success = sp_status.returncode == 0
+
+        return status
+
+    return run_impl
+
+
+@pytest.fixture()
+def check_warnings_on_runner_impl_methods():
+    """
+    Fixture that ensures warnings are emitted when `_impl` methods of runners are called.
+    """
+
+    def check(runner):
+        import contextlib
+        import io
+
+        import numpy as np
+
+        from polygraphy.datatype import DataType
+
+        outfile = io.StringIO()
+        with contextlib.redirect_stdout(outfile), contextlib.redirect_stderr(outfile):
+            runner.activate()
+            # Check that NumPy dtypes are still returned by default
+            metadata = runner.get_input_metadata()
+            for dtype, _ in metadata.values():
+                assert isinstance(dtype, np.dtype)
+
+            metadata = runner.get_input_metadata(use_numpy_dtypes=False)
+            runner.infer(
+                {
+                    name: np.ones(shape, dtype=DataType.to_dtype(dtype, "numpy"))
+                    for name, (dtype, shape) in metadata.items()
+                }
+            )
+            runner.deactivate()
+
+            outfile.seek(0)
+            out = outfile.read()
+
+            def check_warning(method, warning_expected):
+                assert (
+                    f"Calling '{type(runner).__name__}.{method}_impl()' directly is not recommended. Please use '{method}()' instead."
+                    in out
+                ) == warning_expected
+
+            check_warning("get_input_metadata", warning_expected=False)
+            check_warning("activate", warning_expected=False)
+            check_warning("infer", warning_expected=False)
+            check_warning("deactivate", warning_expected=False)
+
+            runner.activate_impl()
+            metadata = runner.get_input_metadata_impl()
+            runner.infer_impl(
+                {
+                    name: np.ones(
+                        shape,
+                        dtype=DataType.to_dtype(DataType.from_dtype(dtype), "numpy"),
+                    )
+                    for name, (dtype, shape) in metadata.items()
+                }
+            )
+            runner.deactivate_impl()
+
+            outfile.seek(0)
+            out = outfile.read()
+            print(out)
+
+            check_warning("get_input_metadata", warning_expected=True)
+            check_warning("activate", warning_expected=True)
+            check_warning("infer", warning_expected=True)
+            check_warning("deactivate", warning_expected=True)
+
+    return check
+
+
+@pytest.fixture()
+def check_warnings_on_loader_impl_methods():
+    """
+    Fixture that ensures warnings are emitted when loader `_impl` methods are called.
+    """
+
+    def check(loader):
+        import contextlib
+        import io
+
+        outfile = io.StringIO()
+        with contextlib.redirect_stdout(outfile), contextlib.redirect_stderr(outfile):
+            warning_msg = f"Calling '{type(loader).__name__}.call_impl()' directly is not recommended. Please use '__call__()' instead."
+            loader.__call__()
+
+            outfile.seek(0)
+            out = outfile.read()
+
+            assert warning_msg not in out
+
+            loader.call_impl()
+
+            outfile.seek(0)
+            out = outfile.read()
+            print(out)
+
+            assert warning_msg in out
+
+    return check
+
+
+@pytest.fixture()
+@pytest.mark.skipif(
+    sys.platform.startswith("win"),
+    reason="Fixture has not been updated to work on Windows",
+)
+def nvinfer_lean_path():
+    lean_library_name = ctypes.util.find_library("nvinfer_lean")
+    for dirname in os.environ.get("LD_LIBRARY_PATH", "").split(os.path.pathsep) + [
+        "/usr/lib/x86_64-linux-gnu"
+    ]:
+        path = os.path.join(dirname, lean_library_name)
+        if os.path.exists(path):
+            return path
+
+    assert False, "Could not find nvinfer_lean!"
+
+
+@pytest.fixture()
+def tmp_python_log_file(tmp_path):
+    # backup original logging configuration
+    orig_handlers = logging.root.handlers[:]
+    orig_level = logging.root.level
+    logging.root.handlers = []
+    tmp_log_file = tmp_path / "test.log"
+    # setup logging to file
+    logging.basicConfig(filename=tmp_log_file, level=0)
+    try:
+        yield tmp_log_file
+    finally:
+        # revert back original configuration
+        logging.root.handlers = orig_handlers
+        logging.root.level = orig_level

tools/Polygraphy/tests/cuda/test_cuda.py

@@ -0,0 +1,186 @@
+#
+# 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
+import torch
+
+from polygraphy import util
+from polygraphy.cuda import DeviceArray, DeviceView, MemcpyKind, Stream, wrapper
+from tests.helper import time_func
+
+
+class TestDeviceView:
+    def test_basic(self):
+        with DeviceArray(shape=(1, 4, 2), dtype=np.float32) as arr:
+            v = DeviceView(arr.ptr, arr.shape, arr.dtype)
+            assert v.ptr == arr.ptr
+            assert v.shape == arr.shape
+            assert v.dtype == arr.dtype
+            assert v.nbytes == arr.nbytes
+            # For backwards compatibility
+            assert isinstance(arr.dtype, np.dtype)
+            assert isinstance(v.dtype, np.dtype)
+
+    def test_with_int_ptr(self):
+        ptr = 74892
+        v = DeviceView(ptr=ptr, shape=(1,), dtype=np.float32)
+        assert v.ptr == ptr
+
+    @pytest.mark.parametrize("module", [np, torch])
+    def test_copy_to(self, module):
+        with DeviceArray((2, 2), dtype=np.float32) as arr:
+            arr.copy_from(module.ones((2, 2), dtype=module.float32) * 4)
+
+            v = DeviceView(arr.ptr, arr.shape, arr.dtype)
+            host_buf = module.zeros((2, 2), dtype=module.float32)
+            v.copy_to(host_buf)
+
+            assert module.all(host_buf == 4)
+
+    def test_numpy(self):
+        with DeviceArray((2, 2), dtype=np.float32) as arr:
+            arr.copy_from(np.ones((2, 2), dtype=np.float32) * 4)
+
+            v = DeviceView(arr.ptr, arr.shape, arr.dtype)
+            assert np.all(v.numpy() == 4)
+
+
+class ResizeTestCase:
+    # *_bytes is the size of the allocated buffer, old/new are the apparent shapes of the buffer.
+    def __init__(self, old, old_size, new, new_size):
+        self.old = old
+        self.old_bytes = old_size * np.float32().itemsize
+        self.new = new
+        self.new_bytes = new_size * np.float32().itemsize
+
+
+RESIZES = [
+    ResizeTestCase(tuple(), 1, (1, 1, 1), 1),  # Reshape (no-op)
+    ResizeTestCase((2, 2, 2), 8, (1, 1), 8),  # Resize to smaller buffer
+    ResizeTestCase((2, 2, 2), 8, (9, 9), 81),  # Resize to larger buffer
+]
+
+
+class TestDeviceBuffer:
+    @pytest.mark.parametrize("shapes", RESIZES)
+    def test_device_buffer_resize(self, shapes):
+        with DeviceArray(shapes.old) as buf:
+            assert buf.allocated_nbytes == shapes.old_bytes
+            assert buf.shape == shapes.old
+            buf.resize(shapes.new)
+            assert buf.allocated_nbytes == shapes.new_bytes
+            assert buf.shape == shapes.new
+
+    @pytest.mark.serial  # Sometimes the GPU may run out of memory if too many other tests are also running.
+    def test_large_allocation(self):
+        dtype = np.byte
+        # See if we can alloc 3GB (bigger than value of signed int)
+        shape = (3 * 1024 * 1024 * 1024,)
+        with DeviceArray(shape=shape, dtype=dtype) as buf:
+            assert buf.allocated_nbytes == util.volume(shape) * np.dtype(dtype).itemsize
+
+    def test_device_buffer_memcpy_async(self):
+        shape = (1, 384)
+        arr = np.ones(shape, dtype=np.int32)
+
+        with DeviceArray(shape) as buf, Stream() as stream:
+            buf.copy_from(arr)
+
+            new_arr = np.empty(shape=shape, dtype=np.int32)
+            buf.copy_to(new_arr, stream)
+
+            stream.synchronize()
+
+            assert np.all(new_arr == arr)
+
+    def test_device_buffer_memcpy_sync(self):
+        shape = (1, 384)
+        arr = np.ones(shape, dtype=np.int32)
+
+        with DeviceArray(shape) as buf:
+            buf.copy_from(arr)
+
+            new_arr = np.empty(shape=shape, dtype=np.int32)
+            buf.copy_to(new_arr)
+
+            assert np.all(new_arr == arr)
+
+    def test_device_buffer_free(self):
+        buf = DeviceArray(shape=(64, 64), dtype=np.float32)
+        assert buf.allocated_nbytes == 64 * 64 * np.float32().itemsize
+
+        buf.free()
+        assert buf.allocated_nbytes == 0
+        assert buf.shape == tuple()
+
+    def test_empty_tensor_to_host(self):
+        with DeviceArray(shape=(5, 2, 0, 3, 0), dtype=np.float32) as buf:
+            assert util.volume(buf.shape) == 0
+
+            host_buf = np.empty(shape=(5, 2, 0, 3, 0), dtype=np.float32)
+            assert util.volume(host_buf.shape) == 0
+
+            buf.copy_to(host_buf)
+            assert host_buf.shape == buf.shape
+            assert host_buf.nbytes == 0
+            assert util.volume(host_buf.shape) == 0
+
+    @pytest.mark.flaky
+    @pytest.mark.serial
+    def test_copy_from_overhead(self):
+        host_buf = np.ones(shape=(4, 8, 512, 512), dtype=np.float32)
+        with DeviceArray(shape=host_buf.shape, dtype=host_buf.dtype) as dev_buf:
+            memcpy_time = time_func(
+                lambda: wrapper().memcpy(
+                    dst=dev_buf.ptr,
+                    src=host_buf.ctypes.data,
+                    nbytes=host_buf.nbytes,
+                    kind=MemcpyKind.HostToDevice,
+                )
+            )
+
+            copy_from_time = time_func(lambda: dev_buf.copy_from(host_buf))
+
+        print(f"memcpy time: {memcpy_time}, copy_from time: {copy_from_time}")
+        assert copy_from_time <= (memcpy_time * 1.12)
+
+    @pytest.mark.flaky
+    @pytest.mark.serial
+    def test_copy_to_overhead(self):
+        host_buf = np.ones(shape=(4, 8, 512, 512), dtype=np.float32)
+        with DeviceArray(shape=host_buf.shape, dtype=host_buf.dtype) as dev_buf:
+            memcpy_time = time_func(
+                lambda: wrapper().memcpy(
+                    dst=host_buf.ctypes.data,
+                    src=dev_buf.ptr,
+                    nbytes=host_buf.nbytes,
+                    kind=MemcpyKind.DeviceToHost,
+                )
+            )
+
+            copy_to_time = time_func(lambda: dev_buf.copy_to(host_buf))
+
+        print(f"memcpy time: {memcpy_time}, copy_to time: {copy_to_time}")
+        assert copy_to_time <= (memcpy_time * 1.12)
+
+    def test_raw(self):
+        with DeviceArray.raw((25,)) as buf:
+            assert buf.shape == (25,)
+            assert buf.nbytes == 25
+            buf.resize((30,))
+            assert buf.shape == (30,)
+            assert buf.nbytes == 30

tools/Polygraphy/tests/func/test_func.py

@@ -0,0 +1,187 @@
+#
+# 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 pytest
+from polygraphy import func
+from polygraphy.exception import PolygraphyException, PolygraphyInternalException
+
+
+class TestExtend:
+    def test_override_rv(self):
+        def x():
+            return 1
+
+        # Since y explicitly returns something, the return value of x is discarded.
+        @func.extend(x)
+        def y(elem):
+            assert elem == 1
+            return 2
+
+        assert y() == 2
+
+    def test_extend_named_parameters(self):
+        def x(arg0, arg1):
+            return arg0, arg1
+
+        @func.extend(x)
+        def y(elem0, elem1):
+            pass
+
+        arg0, arg1 = y(arg1=1, arg0=0)
+        assert arg0 == 0
+        assert arg1 == 1
+
+    def test_extend_0_args_1_rv(self):
+        def x():
+            return 1
+
+        @func.extend(x)
+        def y(elem):
+            assert elem == 1
+
+        assert y() == 1
+
+    def test_extend_0_args_2_rv(self):
+        def x():
+            return 1, 2
+
+        @func.extend(x)
+        def y(elem0, elem1):
+            assert elem0 == 1
+            assert elem1 == 2
+
+        assert y() == (1, 2)
+
+    def test_extend_1_args_0_rv(self):
+        def x(arg0):
+            pass
+
+        @func.extend(x)
+        def y():
+            pass
+
+        y(1)
+
+    def test_extend_1_args_1_rv(self):
+        def x(arg0):
+            assert arg0 == 1
+            return 3
+
+        @func.extend(x)
+        def y(elem):
+            assert elem == 3
+
+        assert y(1) == 3
+
+    def test_extend_2_args_2_rv(self):
+        def x(arg0, arg1):
+            assert arg0 == -1
+            assert arg1 == -1
+            return 1, 2
+
+        @func.extend(x)
+        def y(elem0, elem1):
+            assert elem0 == 1
+            assert elem1 == 2
+
+        assert y(-1, -1) == (1, 2)
+
+    def test_extend_can_modify_rv(self):
+        def x():
+            return []
+
+        @func.extend(x)
+        def y(lst):
+            lst.extend([1, 2, 3])
+
+        assert x() == []
+        assert y() == [1, 2, 3]
+
+    def test_extend_can_modify_rv_objects(self):
+        class ModifiableObj:
+            def __init__(self):
+                self.value = 0
+
+        def x():
+            return ModifiableObj()
+
+        @func.extend(x)
+        def y(mo):
+            mo.value = 1
+
+        assert x().value == 0
+        assert y().value == 1
+
+    def test_extend_incorrect_num_args(self):
+        def x():
+            return 1, 2
+
+        with pytest.raises(
+            PolygraphyException,
+            match=r"Function: y accepts 1 parameter\(s\), but needs to accept 2 parameter\(s\)",
+        ):
+
+            @func.extend(x)
+            def y(elem0):
+                assert elem0 == 1
+
+            y()
+
+    @pytest.mark.parametrize("args_mode", ["kwargs", "args", "mixed"])
+    def test_extend_forward_parameters(self, args_mode):
+        def x(x_arg0, x_arg1):
+            return x_arg0 + x_arg1
+
+        @func.extend(x)
+        def y(x_arg0, x_arg1, x_ret):
+            assert x_ret == x_arg0 + x_arg1
+
+        if args_mode == "kwargs":
+            assert y(x_arg0=2, x_arg1=1) == 3
+        elif args_mode == "args":
+            assert y(2, 1) == 3
+        else:
+            assert args_mode == "mixed"
+            assert y(2, x_arg1=1) == 3
+
+
+class TestConstantMethod:
+    def test_cannot_modify_attrs(self):
+        class Dummy:
+            def __init__(self):
+                self.x = 1
+
+            @func.constantmethod
+            def modify_x(self):
+                self.x = 2
+
+        d = Dummy()
+        with pytest.raises(
+            PolygraphyInternalException, match="was mutated in a constant method"
+        ):
+            d.modify_x()
+
+    def test_cannot_add_attrs(self):
+        class Dummy:
+            @func.constantmethod
+            def modify_x(self):
+                self.x = 2
+
+        d = Dummy()
+        with pytest.raises(
+            PolygraphyInternalException, match="was mutated in a constant method"
+        ):
+            d.modify_x()

tools/Polygraphy/tests/helper.py

@@ -0,0 +1,85 @@
+#
+# 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 os
+import time
+
+ROOT_DIR = os.path.realpath(os.path.join(os.path.dirname(__file__), os.path.pardir))
+
+# Use bin/polygraphy for any invocations of Polygraphy that don't use the script_runner fixture.
+POLYGRAPHY_CMD = [os.path.join(ROOT_DIR, "bin", "polygraphy")]
+
+# CLI tools and all their subtools
+ALL_TOOLS = {
+    "run": [],
+    "convert": [],
+    "inspect": ["data", "model", "tactics", "capability", "diff-tactics"],
+    "check": ["lint"],
+    "surgeon": [
+        "extract",
+        "insert",
+        "sanitize",
+        "prune",
+        "weight-strip",
+        "weight-reconstruct",
+    ],
+    "template": ["trt-network", "trt-config", "onnx-gs"],
+    "debug": ["build", "precision", "reduce", "repeat"],
+    "data": ["to-input"],
+}
+
+
+def get_file_size(path):
+    return os.stat(path).st_size
+
+
+def is_file_empty(path):
+    return get_file_size(path) == 0
+
+
+def is_file_non_empty(path):
+    return not is_file_empty(path)
+
+
+def time_func(func, warm_up=25, iters=50):
+    for _ in range(warm_up):
+        func()
+
+    start = time.time()
+    for _ in range(iters):
+        func()
+    end = time.time()
+    return (end - start) / float(iters)
+
+
+HAS_DLA = None
+
+
+def has_dla():
+    global HAS_DLA
+    if HAS_DLA is None:
+        import tensorrt as trt
+
+        from polygraphy.backend.trt import get_trt_logger
+
+        builder = trt.Builder(get_trt_logger())
+
+        try:
+            HAS_DLA = builder.num_DLA_cores > 0
+        except AttributeError:
+            HAS_DLA = False
+
+    return HAS_DLA

tools/Polygraphy/tests/logger/test_logger.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 inspect
+import logging
+
+import pytest
+
+from polygraphy import util
+from polygraphy.exception.exception import PolygraphyException
+from polygraphy.logger.logger import Logger, SeverityTrie
+
+
+# We don't use the global logger here because we would have to reset the state each time.
+class TestLogger:
+    def test_log_file(self):
+        logger = Logger()
+        with util.NamedTemporaryFile("w+") as log_file:
+            logger.log_file = log_file.name
+            assert logger.log_file == log_file.name
+            logger.info("Hello")
+
+            log_file.seek(0)
+            assert log_file.read() == "[I] Hello\n"
+
+    def test_line_info(self):
+        logger = Logger(line_info=True)
+        with util.NamedTemporaryFile("w+") as log_file:
+            logger.log_file = log_file.name
+
+            logger.info("Hello")
+            log_file.seek(0)
+            assert (
+                f"[I] [tests/logger/test_logger.py:{inspect.currentframe().f_lineno - 3}] Hello\n"
+                == log_file.read()
+            )
+
+    def test_severity_trie_with_no_default(self):
+        logger = Logger(severity={"backend/trt": 10})
+        assert logger.module_severity.get() == Logger.INFO
+        assert logger.module_severity.get("backend/trt") == 10
+
+    def test_callbacks_triggered(self):
+        logger = Logger()
+
+        num_times_called = 0
+
+        def callback(module_severity):
+            nonlocal num_times_called
+            num_times_called += 1
+
+        logger.register_callback(callback)
+        assert num_times_called == 1
+
+        # Callbacks should be triggered whenever severity is set
+        logger.module_severity = {"", Logger.INFO}
+        assert num_times_called == 2
+
+        # Callbacks should be triggered both when we enter and exit the context manager.
+        with logger.verbosity():
+            assert num_times_called == 3
+        assert num_times_called == 4
+
+    @pytest.mark.serial
+    def test_use_python_logging_system(self, tmp_python_log_file):
+        # Clear log file
+        with tmp_python_log_file.open("w") as fp:
+            fp.write("")
+
+        logger = Logger(severity=Logger.ULTRA_VERBOSE)
+        logger.use_python_logging_system = True
+        # add custom Polygraphy levels
+        logging.addLevelName(2, "ULTRA_VERBOSE")
+        logging.addLevelName(4, "SUPER_VERBOSE")
+        logging.addLevelName(6, "EXTRA_VERBOSE")
+        logging.addLevelName(22, "START")
+        logging.addLevelName(28, "FINISH")
+
+        # emit logs
+        logger.ultra_verbose("ultra verbose")
+        logger.super_verbose("super verbose")
+        logger.extra_verbose("extra verbose")
+        logger.verbose("verbose")
+        logger.info("info")
+        logger.start("start")
+        logger.finish("finish")
+        logger.warning("warning")
+        logger.error("error")
+        with pytest.raises(PolygraphyException):
+            logger.critical("critical")
+
+        # verify logs written in the log file
+        with tmp_python_log_file.open() as fp:
+            log_messages = fp.read()
+
+        # Remove lines containing "pytest_shutil.workspace"
+        log_messages = "\n".join(
+            line
+            for line in log_messages.splitlines()
+            if "pytest_shutil.workspace" not in line
+        ) + ("\n" if log_messages.endswith("\n") else "")
+
+        assert (
+            log_messages
+            == """\
+ULTRA_VERBOSE:Polygraphy:[U] ultra verbose
+SUPER_VERBOSE:Polygraphy:[S] super verbose
+EXTRA_VERBOSE:Polygraphy:[X] extra verbose
+DEBUG:Polygraphy:[V] verbose
+INFO:Polygraphy:[I] info
+START:Polygraphy:[I] start
+FINISH:Polygraphy:[I] finish
+WARNING:Polygraphy:[W] warning
+ERROR:Polygraphy:[E] error
+CRITICAL:Polygraphy:[!] critical
+"""
+        )
+
+
+class TestSeverityTrie:
+    @pytest.mark.parametrize(
+        "path,sev",
+        [
+            # Not in trie
+            ("backend", 30),
+            ("mod/importer.py", 30),
+            # Exact paths in trie
+            ("backend/onnx", 28),
+            ("backend/trt", 20),
+            # Submodule of path in trie
+            ("backend/trt/loader.py", 50),
+            ("backend/trt/runner.py", 20),
+            ("backend/onnx/loader.py", 28),
+            # Ensure paths with leading or duplicate slashes work
+            ("/backend", 30),
+            ("backend/////trt", 20),
+        ],
+    )
+    def test_get(self, path, sev):
+        # Duplicate slashes should be handled
+        trie = SeverityTrie(
+            {
+                "": 30,
+                "backend/trt": 20,
+                "backend/trt/loader.py": 50,
+                "backend///////onnx": 28,
+            }
+        )
+        assert trie.get(path) == sev

tools/Polygraphy/tests/mod/conftest.py

@@ -0,0 +1,29 @@
+#
+# 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 pytest
+from tests.helper import ROOT_DIR
+
+
+@pytest.fixture()
+def poly_venv(virtualenv):
+    virtualenv.env["PYTHONPATH"] = ROOT_DIR
+    virtualenv.env["LD_LIBRARY_PATH"] = ""
+
+    # Newer versions of setuptools break pytest-virtualenv
+    virtualenv.run([virtualenv.python, "-m", "pip", "install", "setuptools==59.6.0"])
+
+    return virtualenv

tools/Polygraphy/tests/mod/test_dependencies.py

@@ -0,0 +1,273 @@
+#
+# 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 glob
+import os
+import subprocess as sp
+
+import sys
+import pytest
+from polygraphy import mod, util
+from polygraphy.mod.importer import _version_ok
+
+from tests.helper import ALL_TOOLS, POLYGRAPHY_CMD, ROOT_DIR
+from tests.models.meta import ONNX_MODELS
+
+"""
+The tests here ensure that no additional dependencies are introduced into
+the various modules under Polygraphy.
+"""
+
+
+def is_submodule(path):
+    file_mod = (
+        os.path.isfile(path)
+        and path.endswith(".py")
+        and os.path.basename(path) != "__init__.py"
+    )
+    dir_mod = os.path.isdir(path) and os.path.isfile(os.path.join(path, "__init__.py"))
+    return file_mod or dir_mod
+
+
+MODULE_PATH = os.path.join(ROOT_DIR, "polygraphy")
+SUBMODULE_PATHS = [
+    os.path.relpath(os.path.splitext(path)[0], ROOT_DIR)
+    for path in glob.iglob(os.path.join(MODULE_PATH, "**"), recursive=True)
+    if is_submodule(path)
+]
+
+
+class TestPublicImports:
+    def test_no_extra_submodule_dependencies_required(self, poly_venv):
+        # Submodules should not require any extra dependencies to import.
+        for submodule_path in SUBMODULE_PATHS:
+            submodule_name = ".".join(submodule_path.split(os.path.sep))
+            cmd = [poly_venv.python, "-c", f"from {submodule_name} import *"]
+            print(" ".join(map(str, cmd)))
+            output = poly_venv.run(cmd, capture=True)
+            print(output)
+
+    def test_can_json_without_numpy(self, poly_venv):
+        cmd = [
+            poly_venv.python,
+            "-c",
+            "from polygraphy.json import to_json, from_json; x = to_json(1); x = from_json(x)",
+        ]
+        print(" ".join(map(str, cmd)))
+        output = poly_venv.run(cmd, capture=True)
+        print(output)
+
+
+class TestToolImports:
+    # We should be able to at least launch tools with no dependencies installed.
+    @pytest.mark.parametrize("tool, subtools", ALL_TOOLS.items())
+    def test_can_run_tool_without_deps(self, poly_venv, tool, subtools):
+        BASE_TOOL_CMD = [poly_venv.python, *POLYGRAPHY_CMD, tool, "-h"]
+
+        def check_tool(tool):
+            output = poly_venv.run(tool, capture=True)
+            assert "This tool could not be loaded due to an error:" not in output
+            assert "error:" not in output
+            assert "could not be loaded" not in output
+
+        check_tool(BASE_TOOL_CMD)
+
+        for subtool in subtools:
+            check_tool(BASE_TOOL_CMD + [subtool])
+
+
+class TestAutoinstallDeps:
+    @pytest.mark.parametrize(
+        "cmd",
+        [
+            ["run", ONNX_MODELS["identity"].path, "--onnxrt"],
+            ["run", ONNX_MODELS["identity"].path, "--trt"],
+            [
+                "surgeon",
+                "sanitize",
+                "--fold-constants",
+                ONNX_MODELS["const_foldable"].path,
+                "-o",
+                util.NamedTemporaryFile().name,
+            ],
+        ],
+    )
+    @pytest.mark.slow
+    def test_can_automatically_install_deps(self, poly_venv, cmd):
+        poly_venv.env["POLYGRAPHY_AUTOINSTALL_DEPS"] = "1"
+        cmd = [poly_venv.python, *POLYGRAPHY_CMD] + cmd
+        print(f"Running: {' '.join(map(str, cmd))}")
+        output = poly_venv.run(cmd, capture=True)
+        print(output)
+        assert "is required, but not installed. Attempting to install now" in output
+
+        # Make sure that no PyTorch dependency is accidentally introduced
+        assert "torch" not in poly_venv.installed_packages()
+
+    @pytest.mark.parametrize(
+        "new_ver, expected",
+        [
+            ("==1.4.2", "==1.4.2"),
+            (mod.LATEST_VERSION, ">=1.4.2"),
+        ],
+    )
+    def test_can_automatically_upgrade_deps(self, poly_venv, new_ver, expected):
+        poly_venv.env["POLYGRAPHY_AUTOINSTALL_DEPS"] = "1"
+
+        def get_colored_version():
+            return poly_venv.installed_packages()["colored"].version
+
+        poly_venv.run([poly_venv.python, "-m", "pip", "install", "colored==1.4.0"])
+        assert get_colored_version() == "1.4.0"
+
+        # Insert our own preferred version to make sure it upgrades.
+        poly_venv.run(
+            [
+                poly_venv.python,
+                "-c",
+                f"from polygraphy import mod; colored = mod.lazy_import('colored{new_ver}'); print(colored.__version__)",
+            ]
+        )
+        assert _version_ok(get_colored_version(), expected)
+
+    # Make sure the `requires` parameter of `lazy_import` functions as we expect.
+    @pytest.mark.parametrize("preinstall", [True, False])
+    @pytest.mark.parametrize(
+        "new_ver, expected",
+        [
+            ("==1.4.2", "==1.4.2"),
+        ],
+    )
+    def test_can_automatically_install_requirements(
+        self, poly_venv, new_ver, expected, preinstall
+    ):
+        poly_venv.env["POLYGRAPHY_AUTOINSTALL_DEPS"] = "1"
+
+        def get_colored_version():
+            return poly_venv.installed_packages()["colored"].version
+
+        if preinstall:
+            poly_venv.run([poly_venv.python, "-m", "pip", "install", "colored==1.4.0"])
+            assert get_colored_version() == "1.4.0"
+
+        # Insert our own preferred version to make sure it upgrades.
+        poly_venv.run(
+            [
+                poly_venv.python,
+                "-c",
+                f"from polygraphy import mod; "
+                f"requests = mod.lazy_import('requests==2.25.1', requires=['colored{new_ver}']); "
+                f"requests.__version__; "
+                f"import colored; print(colored.__version__)",
+            ]
+        )
+        assert _version_ok(get_colored_version(), expected)
+
+    @pytest.mark.parametrize(
+        "import_params",
+        [
+            "'colored'",
+            "'colored', pkg_name='colored'",
+            "'colored', install_flags=['--force-reinstall']",
+        ],
+    )
+    def test_autoinstall(self, poly_venv, import_params):
+        poly_venv.env["POLYGRAPHY_AUTOINSTALL_DEPS"] = "1"
+        assert "colored" not in poly_venv.installed_packages()
+
+        poly_venv.run(
+            [
+                poly_venv.python,
+                "-c",
+                f"from polygraphy import mod; colored = mod.lazy_import({import_params}); mod.autoinstall(colored)",
+            ]
+        )
+
+        assert "colored" in poly_venv.installed_packages()
+
+    @pytest.mark.parametrize(
+        "response, should_install",
+        [
+            (" ", True),
+            ("yes", True),
+            ("Y", True),
+            ("n", False),
+        ],
+    )
+    def test_ask_before_autoinstall(self, response, should_install, poly_venv):
+        poly_venv.env["POLYGRAPHY_AUTOINSTALL_DEPS"] = "1"
+        poly_venv.env["POLYGRAPHY_ASK_BEFORE_INSTALL"] = "1"
+        assert "colored" not in poly_venv.installed_packages()
+
+        process = sp.Popen(
+            [
+                poly_venv.python,
+                "-c",
+                "from polygraphy import mod; "
+                "colored = mod.lazy_import('colored'); "
+                "mod.autoinstall(colored)",
+            ],
+            env=poly_venv.env,
+            stdin=sp.PIPE,
+        )
+        process.communicate(input=response.encode())
+        process.wait()
+
+        assert ("colored" in poly_venv.installed_packages()) == should_install
+
+    # We can import inner modules, and Polygraphy should still autoinstall the outermost one.
+    def test_can_install_for_nested_import(self, poly_venv):
+        poly_venv.env["POLYGRAPHY_AUTOINSTALL_DEPS"] = "1"
+
+        poly_venv.run(
+            [
+                poly_venv.python,
+                "-c",
+                "from polygraphy import mod; "
+                "shape_inference = mod.lazy_import('onnx.shape_inference'); "
+                "print(shape_inference.infer_shapes)",
+            ]
+        )
+
+        assert "onnx" in poly_venv.installed_packages()
+
+    def test_all_lazy_imports(self):
+        # NOTE: If this test fails, it means a new lazy dependency has been
+        # introduced. Please ensure that AUTOINSTALL continues to work with the
+        # new dependency.
+        expected = [
+            "fcntl",
+            "matplotlib.pyplot",
+            "matplotlib",
+            "msvcrt",
+            "numpy",
+            "onnx_graphsurgeon",
+            "onnx.numpy_helper",
+            "onnx",
+            "onnxmltools",
+            "onnxruntime.tools.symbolic_shape_infer",
+            "onnxruntime",
+            "tensorflow",
+            "tensorrt",
+            "tf2onnx",
+            "torch",
+            "yaml",
+        ]
+        if sys.version_info < (3, 8):
+            expected.append("importlib_metadata")
+
+        assert mod.importer._all_external_lazy_imports == set(expected)

tools/Polygraphy/tests/mod/test_exporter.py

@@ -0,0 +1,228 @@
+#
+# 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 pytest
+from polygraphy import mod
+from polygraphy.backend.base import BaseLoader
+
+# For test_funcify_with_collision
+functor2 = None
+
+
+class TestExporter:
+    def test_func(self):
+        @mod.export()
+        def test_func0():
+            pass
+
+        assert "test_func0" in __all__
+
+    def test_class(self):
+        @mod.export()
+        class TestClass0:
+            pass
+
+        assert "TestClass0" in __all__
+
+    def test_funcify_func_fails(self):
+        with pytest.raises(AssertionError, match="must be a loader"):
+
+            @mod.export(funcify=True)
+            def test_func1():
+                pass
+
+    def test_funcify_non_base_loader_class(self):
+        with pytest.raises(AssertionError, match="must derive from BaseLoader"):
+
+            @mod.export(funcify=True)
+            class NonFunctor0:
+                def __init__(self, x):
+                    self.x = x
+
+    def test_funcify_duplicate_parameters_in_call_init(self):
+        with pytest.raises(
+            AssertionError, match="call_impl and __init__ have the same argument names"
+        ):
+
+            @mod.export(funcify=True)
+            class DupArgs(BaseLoader):
+                def __init__(self, x):
+                    self.x = x
+
+                def call_impl(self, x):
+                    self.x = x
+
+    def test_funcify_takes_docstring(self):
+        @mod.export(funcify=True)
+        class DocstringFunctor(BaseLoader):
+            """This is a docstring"""
+
+            def __init__(self):
+                pass
+
+            def call_impl(self):
+                pass
+
+        assert "DocstringFunctor" in __all__
+        assert "docstring_functor" in __all__
+
+        assert (
+            docstring_functor.__doc__
+            == "Immediately evaluated functional variant of :class:`DocstringFunctor` .\n"
+        )
+
+    def test_funcify_functor_no_call_args(self):
+        @mod.export(funcify=True)
+        class Functor0(BaseLoader):
+            def __init__(self, x):
+                self.x = x
+
+            def call_impl(self):
+                return self.x
+
+        assert "Functor0" in __all__
+        assert "functor0" in __all__
+        assert functor0(0) == 0
+
+    def test_funcify_functor_with_call_args(self):
+        @mod.export(funcify=True)
+        class Functor1(BaseLoader):
+            def __init__(self, x):
+                self.x = x
+
+            def call_impl(self, y, z):
+                return self.x, y, z
+
+        assert "Functor1" in __all__
+        assert "functor1" in __all__
+
+        # __init__ arguments always precede __call__ arguments
+        x, y, z = functor1(0, 1, -1)
+        assert (x, y, z) == (0, 1, -1)
+
+        # Keyword arguments should behave as expected
+        x, y, z = functor1(y=1, x=0, z=-1)
+        assert (x, y, z) == (0, 1, -1)
+
+    def test_funcify_functor_with_call_args_defaults(self):
+        @mod.export(funcify=True)
+        class FunctorWithCallArgs(BaseLoader):
+            def __init__(self, x=0):
+                self.x = x
+
+            def call_impl(self, y, z=-1):
+                return self.x, y, z
+
+        assert "FunctorWithCallArgs" in __all__
+        assert "functor_with_call_args" in __all__
+
+        # __init__ arguments always precede __call__ arguments
+        x, y, z = functor_with_call_args(y=1)
+        assert (x, y, z) == (0, 1, -1)
+
+        # Keyword arguments should behave as expected
+        x, y, z = functor_with_call_args(y=1)
+        assert (x, y, z) == (0, 1, -1)
+
+    def test_funcify_with_collision(self):
+        with pytest.raises(AssertionError, match="symbol is already defined"):
+
+            @mod.export(funcify=True)
+            class Functor2(BaseLoader):
+                def __init__(self, x):
+                    self.x = x
+
+                def call_impl(self, y, z):
+                    return self.x, y, z
+
+    def test_funcify_functor_with_dynamic_call_args_kwargs(self):
+        @mod.export(funcify=True)
+        class Functor3(BaseLoader):
+            def __init__(self, f):
+                self.f = f
+
+            def call_impl(self, *args, **kwargs):
+                return self.f(*args, **kwargs)
+
+        assert "Functor3" in __all__
+        assert "functor3" in __all__
+
+        # We should be able to pass arbitrary arguments to call now.
+        # __init__ arguments are always first.
+        def func(arg0, arg1, arg2):
+            return arg0 + arg1 + arg2
+
+        assert functor3(func, 1, 2, arg2=4) == 7
+
+    def test_funcify_with_inherited_init(self):
+        class BaseFunctor4(BaseLoader):
+            def __init__(self, x):
+                self.x = x
+
+        @mod.export(funcify=True)
+        class Functor4(BaseFunctor4):
+            def call_impl(self):
+                return self.x
+
+        assert "Functor4" in __all__
+        assert "functor4" in __all__
+
+        assert functor4(-1) == -1
+
+    def test_funcify_functor_with_default_vals(self):
+        @mod.export(funcify=True)
+        class FunctorWithDefaults(BaseLoader):
+            def __init__(self, w, x=1):
+                self.w = w
+                self.x = x
+
+            def call_impl(self, y, z=3):
+                return self.w, self.x, y, z
+
+        assert "FunctorWithDefaults" in __all__
+        assert "functor_with_defaults" in __all__
+
+        # Since x and z have default values, the arguments will be interlaced into:
+        # w, y, x, z
+        # __init__ parameters take precedence, and call_impl parameters follow.
+        w, x, y, z = functor_with_defaults(-1, -2)  # Set just w, y
+        assert (w, x, y, z) == (-1, 1, -2, 3)
+
+        w, x, y, z = functor_with_defaults(0, 1, 2, 3)  # Set all
+        assert (w, x, y, z) == (0, 2, 1, 3)
+
+    def test_funcify_functor_with_type_annotations(self):
+        @mod.export(funcify=True)
+        class FunctorWithTypeAnnotations(BaseLoader):
+            def __init__(self, w: int, x: int = 1):
+                self.w = w
+                self.x = x
+
+            def call_impl(self, y: int, z: int = 3):
+                return self.w, self.x, y, z
+
+        assert "FunctorWithTypeAnnotations" in __all__
+        assert "functor_with_type_annotations" in __all__
+
+        # Since x and z have default values, the arguments will be interlaced into:
+        # w, y, x, z
+        # __init__ parameters take precedence, and call_impl parameters follow.
+        w, x, y, z = functor_with_type_annotations(-1, -2)  # Set just w, y
+        assert (w, x, y, z) == (-1, 1, -2, 3)
+
+        w, x, y, z = functor_with_type_annotations(0, 1, 2, 3)  # Set all
+        assert (w, x, y, z) == (0, 2, 1, 3)

tools/Polygraphy/tests/mod/test_importer.py

@@ -0,0 +1,185 @@
+#
+# 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 copy
+import os
+import sys
+from textwrap import dedent
+
+import pytest
+import tempfile
+import tensorrt as trt
+from polygraphy import mod, util
+from polygraphy.exception import PolygraphyException
+from polygraphy.mod.importer import _version_ok
+
+common_backend = mod.lazy_import("polygraphy.backend.common")
+
+class TestImporter:
+    def test_import_from_script(self):
+        script = dedent(
+            """
+            from polygraphy.backend.trt import CreateNetwork
+            from polygraphy import func
+            import tensorrt as trt
+
+            @func.extend(CreateNetwork())
+            def load_network(builder, network):
+                inp = network.add_input("input", dtype=trt.float32, shape=(1, 1))
+                out = network.add_identity(inp).get_output(0)
+                network.mark_output(out)
+            """
+        )
+
+        with util.NamedTemporaryFile("w+", suffix=".py") as f:
+            f.write(script)
+            f.flush()
+            os.fsync(f.fileno())
+
+            orig_sys_path = copy.deepcopy(sys.path)
+            load_network = mod.import_from_script(f.name, "load_network")
+            assert sys.path == orig_sys_path
+            builder, network = load_network()
+            with builder, network:
+                assert isinstance(builder, trt.Builder)
+                assert isinstance(network, trt.INetworkDefinition)
+                assert network.num_layers == 1
+                assert network.get_layer(0).type == trt.LayerType.IDENTITY
+            assert sys.path == orig_sys_path
+
+    def test_import_from_script_same_method_different_modules(self):
+        module1_script = dedent(
+            """
+            def print_message():
+                print(f"msg1::print_message")
+                return "msg1"
+            """
+        )
+
+        module2_script = dedent(
+            """
+            def print_message():
+                print(f"msg2::print_message")
+                return "msg2"
+            """
+        )
+
+        with tempfile.TemporaryDirectory() as tempdir:
+            os.mkdir(os.path.join(tempdir, "msg1"))
+            with open(os.path.join(tempdir, "msg1", "msg.py"), "w+") as msg1_msg:
+                msg1_msg.write(module1_script)
+                msg1_msg.flush()
+                os.fsync(msg1_msg.fileno())
+
+                os.mkdir(os.path.join(tempdir, "msg2"))
+                with open(os.path.join(tempdir, "msg2", "msg.py"), "w+") as msg2_msg:
+                    msg2_msg.write(module2_script)
+                    msg2_msg.flush()
+                    os.fsync(msg2_msg.fileno())
+                    
+                    for msg_module in ['msg1', 'msg2']:
+                        msg_loc = os.path.join(tempdir,msg_module,'msg.py')
+                        msg = common_backend.invoke_from_script(msg_loc, "print_message")
+                        assert msg==msg_module
+
+    def test_import_non_existent(self):
+        script = dedent(
+            """
+            def example():
+                pass
+            """
+        )
+
+        with util.NamedTemporaryFile("w+", suffix=".py") as f:
+            f.write(script)
+            f.flush()
+            os.fsync(f.fileno())
+
+            orig_sys_path = copy.deepcopy(sys.path)
+            example = mod.import_from_script(f.name, "example")
+            assert sys.path == orig_sys_path
+
+            assert example is not None
+            example()
+
+            with pytest.raises(
+                PolygraphyException, match="Could not import symbol: non_existent from"
+            ):
+                mod.import_from_script(f.name, "non_existent")
+            assert sys.path == orig_sys_path
+
+    @pytest.mark.parametrize(
+        "ver, pref, expected",
+        [
+            ("0.0.0", "==0.0.0", True),
+            ("0.0.0", "== 0.0.1", False),
+            ("0.0.0", ">= 0.0.0", True),
+            ("0.0.0", ">=0.0.1", False),
+            ("0.0.0", "<= 0.0.0", True),
+            ("0.0.2", "<=0.0.1", False),
+            ("0.0.1", "> 0.0.0", True),
+            ("0.0.1", ">0.0.1", False),
+            ("0.0.0", "< 0.0.1", True),
+            ("0.0.0", "< 0.0.0", False),
+            ("0.2.0", mod.LATEST_VERSION, False),
+        ],
+    )
+    def test_version_ok(self, ver, pref, expected):
+        assert _version_ok(ver, pref) == expected
+
+    def test_is_installed_works_when_package_name_differs_from_module_name(
+        self, poly_venv
+    ):
+        assert "onnxruntime" not in poly_venv.installed_packages()
+        assert "onnxruntime-gpu" not in poly_venv.installed_packages()
+
+        poly_venv.run(
+            [poly_venv.python, "-m", "pip", "install", "onnxruntime-gpu", "--no-deps"]
+        )
+
+        # The `onnxruntime-gpu` package provides the `onnxruntime` module.
+        # `is_installed()` should be able to understand that.
+        poly_venv.run(
+            [
+                poly_venv.python,
+                "-c",
+                "from polygraphy import mod; onnxrt = mod.lazy_import('onnxruntime<0'); assert onnxrt.is_installed()",
+            ]
+        )
+
+    @pytest.mark.parametrize(
+        "mod_check", ["mod.has_mod('colored')", "colored.is_installed()"]
+    )
+    def test_has_mod(self, poly_venv, mod_check):
+        assert "colored" not in poly_venv.installed_packages()
+        poly_venv.run(
+            [
+                poly_venv.python,
+                "-c",
+                f"from polygraphy import mod; colored = mod.lazy_import('colored'); assert not {mod_check}",
+            ]
+        )
+
+        poly_venv.run([poly_venv.python, "-m", "pip", "install", "colored==1.4.0"])
+        # Make sure `has_mod` doesn't actually import the package.
+        poly_venv.run(
+            [
+                poly_venv.python,
+                "-c",
+                "from polygraphy import mod; import sys; assert mod.has_mod('colored'); assert 'colored' not in sys.modules; import colored; assert 'colored' in sys.modules",
+            ]
+        )

tools/Polygraphy/tests/mod/test_util.py

@@ -0,0 +1,39 @@
+#
+# 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 pytest
+from polygraphy import mod
+
+
+@pytest.mark.parametrize(
+    "ver0, ver1, expected",
+    [
+        ("1.0.0", "2.0.0", False),
+        ("1.0.0", "0.9.0", True),
+        ("0.0b1", "0.1b1", False),
+        ("0.1b1", "0.1b0", True),
+        ("0.12b1", "0.1b0", True),
+        ("0.1b0", "0.1a0", True),
+        ("0.1rc0", "0.1b0", True),
+        ("0.post1", "0.post0", True),
+        ("0.post1", "0.post2", False),
+        ("1.13.1+cu117", "1.13.0", True),
+        ("1.13.1+cu117", "1.13.2", False),
+    ],
+)
+def test_version(ver0, ver1, expected):
+    assert (mod.version(ver0) > mod.version(ver1)) == expected

tools/Polygraphy/tests/models/and.onnx

@@ -0,0 +1,17 @@
+backend-test:_
+
+
x
+
yand"And
+test_and2dZ
+
x
+	
+
+Z
+
y
+	
+
+b
+and
+	
+
+B	

tools/Polygraphy/tests/models/dynamic_identity.onnx

@@ -0,0 +1,15 @@
+backend_test:y
+
+
X
Y"Identityonnx_dynamic_identityZ&
+
X!
+

+

+
+height
+widthb&
+
Y!
+

+

+
+height
+widthB

tools/Polygraphy/tests/models/identity.onnx

@@ -0,0 +1,15 @@
+backend-test:[
+
+
x
y"Identity
+test_identityZ
+
x
+

+

+

+
+b
+
y
+

+

+

+

tools/Polygraphy/tests/models/instancenorm.onnx

@@ -0,0 +1,27 @@
+TensorRT-WF:é

+9
+
x
+
s
+bias
y"InstanceNormalization*
+epsilon
+×#< 

instancenorm2d_4dims_epsilon*
"™ë?ä×o?Ï…¿B
s*
"hÖ¿81=¿6þ·¿BbiasZ
+
x
+

+
+
+
+Z
+
s
+
+

+Z
+bias
+
+

+b
+
y
+

+
+
+
+B

tools/Polygraphy/tests/models/make_models.py

@@ -0,0 +1,974 @@
+#
+# 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.
+#
+"""
+Helper utility to generate models to help test the `debug reduce`
+subtool, which reduces failing ONNX models.
+"""
+import os
+
+import tempfile
+import numpy as np
+import onnx
+import subprocess
+import onnx_graphsurgeon as gs
+from meta import ONNX_MODELS
+from polygraphy.tools.sparse import SparsityPruner
+
+CURDIR = os.path.dirname(__file__)
+
+
+@gs.Graph.register()
+def identity(self, inp, **kwargs):
+    out = self.layer(op="Identity", inputs=[inp], outputs=["identity_out"], **kwargs)[0]
+    out.dtype = inp.dtype
+    return out
+
+
+@gs.Graph.register()
+def add(self, a, b, **kwargs):
+    return self.layer(op="Add", inputs=[a, b], outputs=["add_out"], **kwargs)[0]
+
+
+@gs.Graph.register()
+def div(self, a, b, **kwargs):
+    return self.layer(op="Div", inputs=[a, b], outputs=["div_out"], **kwargs)[0]
+
+
+@gs.Graph.register()
+def sub(self, a, b, **kwargs):
+    return self.layer(op="Sub", inputs=[a, b], outputs=["sub_out"], **kwargs)[0]
+
+
+@gs.Graph.register()
+def constant(self, values: gs.Constant, **kwargs):
+    return self.layer(
+        op="Constant", outputs=["constant_out"], attrs={"value": values}, **kwargs
+    )[0]
+
+
+@gs.Graph.register()
+def reshape(self, data, shape, **kwargs):
+    return self.layer(
+        op="Reshape", inputs=[data, shape], outputs=["reshape_out"], **kwargs
+    )[0]
+
+
+@gs.Graph.register()
+def matmul(self, a, b, **kwargs):
+    return self.layer(op="MatMul", inputs=[a, b], outputs=["matmul_out"], **kwargs)[0]
+
+
+@gs.Graph.register()
+def tile(self, inp, repeats):
+    return self.layer(op="Tile", inputs=[inp, repeats], outputs=["tile_out"])[0]
+
+
+@gs.Graph.register()
+def nonzero(self, inp, **kwargs):
+    return self.layer(op="NonZero", inputs=[inp], outputs=["nonzero_out"], **kwargs)[0]
+
+
+# Name range as onnx_range as range is a python built-in function.
+@gs.Graph.register()
+def onnx_range(self, start, limit, delta, **kwargs):
+    return self.layer(
+        op="Range", inputs=[start, limit, delta], outputs=["range_out"], **kwargs
+    )[0]
+
+
+@gs.Graph.register()
+def cast(self, input, type, **kwargs):
+    return self.layer(
+        op="Cast", inputs=[input], attrs={"to": type}, outputs=["cast_out"], **kwargs
+    )[0]
+
+
+@gs.Graph.register()
+def reduce_max(self, input, keep_dims, **kwargs):
+    return self.layer(
+        op="ReduceMax",
+        inputs=[input],
+        attrs={"keepdims": keep_dims},
+        outputs=["reduce_max_out"],
+        **kwargs,
+    )[0]
+
+
+@gs.Graph.register()
+def conv(self, input, weights, kernel_shape, **kwargs):
+    return self.layer(
+        op="Conv",
+        inputs=[input, weights],
+        attrs={"kernel_shape": kernel_shape},
+        outputs=["conv_out"],
+        **kwargs,
+    )[0]
+
+
+@gs.Graph.register()
+def split(self, inp, split, axis=0):
+    return self.layer(
+        op="Split",
+        inputs=[inp],
+        outputs=[f"split_out_{i}" for i in range(len(split))],
+        attrs={"axis": axis, "split": split},
+    )
+
+
+@gs.Graph.register()
+def transpose(self, inp, **kwargs):
+    return self.layer(
+        op="Transpose", inputs=[inp], outputs=["transpose_out"], **kwargs
+    )[0]
+
+
+@gs.Graph.register()
+def quantize_linear(self, inp, y_scale, y_zero_point, **kwargs):
+    return self.layer(
+        op="QuantizeLinear",
+        inputs=[inp, y_scale, y_zero_point],
+        outputs=["quantize_linear_out"],
+        **kwargs,
+    )[0]
+
+
+@gs.Graph.register()
+def dequantize_linear(self, inp, x_scale, x_zero_point, **kwargs):
+    return self.layer(
+        op="DequantizeLinear",
+        inputs=[inp, x_scale, x_zero_point],
+        outputs=["dequantize_linear_out"],
+        **kwargs,
+    )[0]
+
+
+def save(graph, model_name):
+    path = os.path.join(CURDIR, model_name)
+    print(f"Writing: {path}")
+    onnx.save(gs.export_onnx(graph), path)
+
+
+def make_sparse(graph):
+    sparsity_pruner = SparsityPruner(gs.export_onnx(graph))
+    return gs.import_onnx(sparsity_pruner.prune())
+
+
+# Generates a model with multiple inputs/outputs:
+#
+#    X0    Y0
+#    |     |
+#    X1    Y1
+#      \  /
+#       Z0
+#      /  \
+#    Z1    Z2
+#
+def make_multi_input_output():
+    DTYPE = np.float32
+    SHAPE = (1,)
+
+    X0 = gs.Variable("X0", dtype=DTYPE, shape=SHAPE)
+    Y0 = gs.Variable("Y0", dtype=DTYPE, shape=SHAPE)
+
+    graph = gs.Graph(inputs=[X0, Y0])
+
+    X1 = graph.identity(X0)
+    Y1 = graph.identity(Y0)
+
+    Z0 = graph.add(X1, Y1)
+
+    Z1 = graph.identity(Z0)
+    Z1.dtype = DTYPE
+    Z1.shape = SHAPE
+
+    Z2 = graph.identity(Z0)
+    Z2.dtype = DTYPE
+    Z2.shape = SHAPE
+
+    graph.outputs = [Z1, Z2]
+
+    save(graph, "reducable.onnx")
+
+
+make_multi_input_output()
+
+
+# Generates a linear model with a Constant node and no inputs:
+#
+#    X0 (Constant)
+#    |
+#    X1 (Identity)
+#    |
+#    X2 (Identity)
+#
+def make_constant_linear():
+    DTYPE = np.float32
+    SHAPE = (4, 4)
+
+    graph = gs.Graph()
+
+    X0 = graph.constant(gs.Constant("const", values=np.ones(SHAPE, dtype=DTYPE)))
+    # Explicitly clear shape to trigger the failure condition in reduce
+    X0.shape = None
+
+    X1 = graph.identity(X0)
+    X2 = graph.identity(X1)
+    X2.dtype = DTYPE
+    X2.shape = SHAPE
+
+    graph.outputs = [X2]
+
+    save(graph, "reducable_with_const.onnx")
+
+
+make_constant_linear()
+
+
+# Generates a model whose node uses the same tensor for multiple inputs
+#
+#    inp
+#    / \
+#    Add
+#     |
+#    out
+#
+def make_dup_input():
+    DTYPE = np.float32
+    SHAPE = (4, 4)
+
+    inp = gs.Variable("inp", dtype=DTYPE, shape=SHAPE)
+
+    graph = gs.Graph(inputs=[inp])
+    out = graph.add(inp, inp)
+    out.dtype = DTYPE
+    graph.outputs = [out]
+
+    save(graph, "add_with_dup_inputs.onnx")
+
+
+make_dup_input()
+
+
+# Generates a model with a no-op reshape
+#
+#     inp    shape
+#        \ /
+#       Reshape
+#         |
+#        out
+#
+def make_no_op_reshape():
+    DTYPE = np.float32
+    SHAPE = (4, 4)
+
+    data = gs.Variable("data", dtype=DTYPE, shape=SHAPE)
+
+    graph = gs.Graph(inputs=[data])
+    out = graph.reshape(data, np.array(SHAPE, dtype=np.int64))
+    out.dtype = DTYPE
+    graph.outputs = [out]
+
+    save(graph, "no_op_reshape.onnx")
+
+
+make_no_op_reshape()
+
+
+# Generates a model that overflows FP16
+#
+#   inp
+#    |
+#  MatMul
+#    |
+#  Add
+#   |
+#  Sub
+#   |
+#  MatMul
+#   |
+#  out
+#
+def make_needs_constraints():
+    SIZE = 256
+
+    x = gs.Variable("x", shape=(1, 1, SIZE, SIZE), dtype=np.float32)
+    I_rot90 = gs.Constant(
+        name="I_rot90",
+        values=np.rot90(
+            np.identity(SIZE, dtype=np.float32).reshape((1, 1, SIZE, SIZE))
+        ),
+    )
+    fp16_max = gs.Constant(
+        name="fp16_max",
+        values=np.array([np.finfo(np.float16).max], dtype=np.float32).reshape(
+            (1, 1, 1, 1)
+        ),
+    )
+
+    graph = gs.Graph(inputs=[x])
+    y = graph.matmul(x, I_rot90, name="MatMul_0")
+    z = graph.add(y, fp16_max, name="Add")
+    w = graph.sub(z, fp16_max, name="Sub")
+    u = graph.matmul(w, I_rot90, name="MatMul_1")
+
+    u.dtype = np.float32
+    graph.outputs = [u]
+
+    save(graph, "needs_constraints.onnx")
+
+
+make_needs_constraints()
+
+
+# Generates a model that will become very large when constant-folded
+#
+#   inp
+#    |
+#  Tile
+#    |
+#   out
+#
+def make_constant_fold_bloater():
+    graph = gs.Graph()
+    # Input is 1MiB, tiled to 10MiB
+    out = graph.tile(
+        np.ones(shape=(1024, 256), dtype=np.float32), repeats=np.array([1, 10])
+    )
+    out.dtype = np.float32
+    graph.outputs = [out]
+
+    save(graph, "constant_fold_bloater.onnx")
+
+
+make_constant_fold_bloater()
+
+
+# Generate a model with a data-dependent shape
+#
+#    inp
+#     |
+#   NonZero
+#     |
+#    out
+#
+def make_nonzero():
+    inp = gs.Variable("input", shape=(4,), dtype=np.int64)
+
+    graph = gs.Graph(inputs=[inp])
+    out = graph.nonzero(inp)
+    out.dtype = np.int64
+    graph.outputs = [out]
+
+    save(graph, "nonzero.onnx")
+
+
+make_nonzero()
+
+
+# Generate a model where a node has multiple outputs that are graph outputs
+#
+#     inp
+#      |
+#   Identity
+#      |
+#     id0
+#        \
+#       Split
+#     /       \
+# split_out0   split_out1 (graph output)
+#      |
+#   Identity
+#      |
+#  id1 (graph output)
+#
+#
+def make_multi_output():
+    inp = gs.Variable("input", shape=(4, 5), dtype=np.float32)
+
+    graph = gs.Graph(inputs=[inp])
+    id0 = graph.identity(inp)
+    [split_out0, split_out1] = graph.split(id0, split=[2, 2])
+    id1 = graph.identity(split_out0)
+    graph.outputs = [id1, split_out1]
+
+    for out in graph.outputs:
+        out.dtype = np.float32
+
+    save(graph, "multi_output.onnx")
+
+
+make_multi_output()
+
+
+# Generate a model where a tensor contains unbounded DDS.
+# Use Conv_0 and ReduceMax to generate a DDS scalar tensor, and send to Range as input `limit`.
+# The output of Range has an unbounded shape.
+#
+#     input
+#       |
+#     Conv_0
+#       |
+#    ReduceMax
+#       |
+#     Range
+#       |
+#     Conv_1
+#       |
+#     output
+#
+def make_unbounded_dds():
+    input = gs.Variable("Input", shape=(1, 3, 10, 10), dtype=np.float32)
+    graph = gs.Graph(inputs=[input], opset=13)
+    weights_0 = graph.constant(
+        gs.Constant("Weights_0", values=np.ones((3, 3, 3, 3), dtype=np.float32))
+    )
+    weights_1 = graph.constant(
+        gs.Constant("Weights_1", values=np.ones((4, 1, 1, 1), dtype=np.float32))
+    )
+
+    conv_0 = graph.conv(input, weights_0, [3, 3], name="Conv_0")
+    reduce_max_0 = graph.reduce_max(conv_0, keep_dims=0, name="ReduceMax_0")
+
+    cast_0 = graph.cast(
+        reduce_max_0, getattr(onnx.TensorProto, "INT64"), name="Cast_to_int64"
+    )
+    range_0 = graph.onnx_range(
+        np.array(0, dtype=np.int64), cast_0, np.array(1, dtype=np.int64), name="Range"
+    )
+    cast_1 = graph.cast(
+        range_0, getattr(onnx.TensorProto, "FLOAT"), name="Cast_to_float"
+    )
+
+    reshape_1 = graph.reshape(
+        cast_1, np.array([1, 1, -1, 1], dtype=np.int64), name="Reshape_1"
+    )
+    conv_1 = graph.conv(reshape_1, weights_1, [1, 1], name="Conv_1")
+
+    graph.outputs = [conv_1]
+
+    for out in graph.outputs:
+        out.dtype = np.float32
+
+    save(graph, "unbounded_dds.onnx")
+
+
+make_unbounded_dds()
+
+
+def make_small_matmul(name, dtype, save_sparse=False):
+    M = 8
+    N = 8
+    K = 16
+    a = gs.Variable("a", shape=(M, K), dtype=dtype)
+    g = gs.Graph(inputs=[a], opset=13)
+    val = np.random.uniform(-3, 3, size=K * N).astype(dtype).reshape((K, N))
+    b = gs.Constant("b", values=val)
+    c = g.matmul(a, b, name="matmul")
+    c.dtype = dtype
+    g.outputs = [c]
+
+    save(g, name)
+    if save_sparse:
+        save(make_sparse(g), "sparse." + name)
+
+
+make_small_matmul("matmul.onnx", np.float32, save_sparse=True)
+make_small_matmul("matmul.fp16.onnx", np.float16)
+
+
+def make_small_conv(name):
+    N = 1
+    C = 16
+    H = 8
+    W = 8
+    K = 4
+    F = 4
+    a = gs.Variable("a", shape=(N, C, H, W), dtype=np.float32)
+    g = gs.Graph(inputs=[a], opset=13)
+    val = (
+        np.random.uniform(-3, 3, size=K * C * F * F)
+        .reshape((K, C, F, F))
+        .astype(np.float32)
+    )
+    b = gs.Constant("b", values=val)
+    c = g.conv(a, b, (F, F), name="conv")
+    c.dtype = np.float32
+    g.outputs = [c]
+
+    save(g, name)
+    save(make_sparse(g), "sparse." + name)
+
+
+make_small_conv("conv.onnx")
+
+
+def make_unsorted():
+    inp = gs.Variable("input", shape=(1, 1), dtype=np.float32)
+    graph = gs.Graph(inputs=[inp])
+    graph.outputs = [graph.identity(graph.identity(inp))]
+
+    graph.nodes = list(reversed(graph.nodes))
+    save(graph, "unsorted.onnx")
+
+
+make_unsorted()
+
+
+def make_empty():
+    g = gs.Graph(inputs=[], opset=13)
+    g.outputs = []
+
+    save(g, "empty.onnx")
+
+
+make_empty()
+
+
+# Builds a graph that has unused nodes and inputs.
+#
+# f  e
+# |\  |
+# H  G
+# |  |
+# h  g
+# |
+# I
+# |
+# i
+#
+# e is an unused input.
+# G is an unused node.
+# This graph is useful for testing if `lint` catches unused nodes and inputs.
+def make_cleanable():
+    e = gs.Variable(name="e", dtype=np.float32, shape=(1, 1))
+    f = gs.Variable(name="f", dtype=np.float32, shape=(1, 1))
+    h = gs.Variable(name="h", dtype=np.float32, shape=(1, 1))
+    i = gs.Variable(name="i", dtype=np.float32, shape=(1, 1))
+    g = gs.Variable(name="g", dtype=np.float32, shape=(2, 1))
+
+    nodes = [
+        gs.Node(op="Concat", name="G", inputs=[e, f], outputs=[g], attrs={"axis": 0}),
+        gs.Node(op="Dropout", name="H", inputs=[f], outputs=[h]),
+        gs.Node(op="Identity", name="I", inputs=[h], outputs=[i]),
+    ]
+
+    graph = gs.Graph(nodes=nodes, inputs=[e, f], outputs=[i])
+    save(graph, "cleanable.onnx")
+
+
+make_cleanable()
+
+
+# Generates a graph with very deranged names
+# Tests that the unique renaming in lint tool works
+def make_renamable():
+    a = gs.Variable(name="a", dtype=np.float32, shape=(1, 1))
+    b = gs.Variable(name="b", dtype=np.float32, shape=(1, 1))
+    c = gs.Variable(name="c", dtype=np.float32, shape=(1, 1))
+    d = gs.Variable(name="d", dtype=np.float32, shape=(1, 1))
+    e = gs.Variable(name="e", dtype=np.float32, shape=(2, 1))
+
+    nodes = [
+        gs.Node(op="Identity", name="", inputs=[a], outputs=[b]),
+        gs.Node(
+            op="Dropout", name="polygraphy_unnamed_node_0", inputs=[b], outputs=[c]
+        ),
+        gs.Node(
+            op="Identity", name="polygraphy_unnamed_node_0_0", inputs=[c], outputs=[d]
+        ),
+        gs.Node(op="Dropout", name="", inputs=[d], outputs=[e]),
+    ]
+
+    graph = gs.Graph(nodes=nodes, inputs=[a], outputs=[e])
+    save(graph, "renamable.onnx")
+
+
+make_renamable()
+
+####### Generate some invalid models #######
+
+### Graphs whose errors are data-dependent ###
+
+
+# Generats an invalid graph with multiple parallel bad nodes.
+# The graph is invalid due to multiple parallel nodes failing.
+# This is is the graph:
+#    A    B    C    D  E    F    G
+#     \  /      \  /    \  /      \
+#    MatMul_0* Add_0*  MatMul_1 NonZero
+#        \        /        \    /
+#         MatMul_2       MatMul_3*
+#               \       /
+#                \     /
+#                Add_1
+#                  |
+#                output
+# The graph is invalid because MatMul_0, Add_0 and MatMul_3 all will fail.
+# MatMul_0 should fail because A and B are not compatible.
+# Add_0 should fail because C and D are not compatible.
+# MatMul_3 should fail because result of MatMul2 and the Data-dependent shape of output of
+# NonZero are not compatible.
+#
+# This graph is useful for testing if `lint` catches multiple parallel bad nodes that may/may not be data-dependent.
+#
+def make_bad_graph_with_parallel_invalid_nodes():
+    DTYPE = np.float32
+    BAD_DIM = 3
+
+    graph = gs.Graph(name="bad_graph_with_parallel_invalid_nodes")
+
+    A = gs.Variable("A", dtype=DTYPE, shape=(1, BAD_DIM))
+    B = gs.Variable("B", dtype=DTYPE, shape=(4, 4))
+    mm_ab_out = graph.matmul(
+        A, B, name="MatMul_0"
+    )  # This node will fail because A and B are not compatible.
+
+    C = gs.Variable("C", dtype=DTYPE, shape=(BAD_DIM, 4))
+    D = gs.Variable("D", dtype=DTYPE, shape=(4, 1))
+    add_cd_out = graph.add(
+        C, D, name="Add_0"
+    )  # This node will fail because C and D are not compatible.
+
+    pre_out_1 = graph.matmul(mm_ab_out, add_cd_out, name="MatMul_2")
+
+    E = gs.Variable("E", dtype=DTYPE, shape=(1, 4))
+    F = gs.Variable("F", dtype=DTYPE, shape=(4, 1))
+    mm_ef_out = graph.matmul(E, F, name="MatMul_1")
+    mm_ef_out_int64 = graph.cast(
+        mm_ef_out, onnx.TensorProto.INT64, name="cast_to_int64"
+    )
+
+    G = gs.Variable("G", dtype=np.int64, shape=(4, 4))
+    nz_g_out = graph.nonzero(G, name="NonZero")  # `nz_g_out` shape is data-dependent.
+
+    pre_out_2 = graph.matmul(
+        mm_ef_out_int64, nz_g_out, name="MatMul_3"
+    )  # This node will fail because `mm_ef_out_int64` and `nz_g_out` are not compatible.
+    pre_out_2_float = graph.cast(
+        pre_out_2, getattr(onnx.TensorProto, "FLOAT"), name="cast_to_float"
+    )
+
+    out = graph.add(pre_out_1, pre_out_2_float, name="Add_1")
+    out.dtype = DTYPE
+
+    graph.inputs = [A, B, C, D, E, F, G]
+    graph.outputs = [out]
+
+    save(graph, "bad_graph_with_parallel_invalid_nodes.onnx")
+
+
+make_bad_graph_with_parallel_invalid_nodes()
+
+
+# Generates the following graph:
+#                 cond
+#                  |
+#                 If
+#                  |
+#             z (x or y)
+#              \   |
+#               MatMul
+#                  |
+#               output
+# If `cond` is True, then `x` is used, otherwise `y` is used.
+# `x` is compatible with `z`, while `y` is NOT compatible with `z`.
+# Based on the value of `cond`, the graph may be valid or invalid.
+#
+# This graph is useful to check whether the error message is caught or not at runtime based on data input.
+#
+def make_bad_graph_conditionally_invalid():
+    X = [[4.0], [3.0]]  # shape (2, 1), compatible with Z for MatMul
+    Y = [2.0, 4.0]  # shape (2,), incompatible with Z for MatMul
+    Z = [[2.0, 4.0]]  # shape (1, 2)
+
+    cond = gs.Variable(
+        "cond", dtype=np.bool_, shape=(1,)
+    )  # input to If, True or False based on user input.
+
+    graph = gs.Graph(name="bad_graph_conditionally_invalid")
+
+    x = gs.Constant("x", values=np.array(X, dtype=np.float32))
+    y = gs.Constant("y", values=np.array(Y, dtype=np.float32))
+
+    then_out = gs.Variable("then_out", dtype=np.float32, shape=None)
+    else_out = gs.Variable("else_out", dtype=np.float32, shape=None)
+
+    then_const_node = gs.Node(
+        op="Constant", inputs=[], outputs=[then_out], attrs={"value": x}
+    )  # node for `then_branch` Graph
+    else_const_node = gs.Node(
+        op="Constant", inputs=[], outputs=[else_out], attrs={"value": y}
+    )  # node for `else_branch` Graph
+
+    then_body = gs.Graph(
+        nodes=[then_const_node], name="then_body", inputs=[], outputs=[then_out]
+    )  # Graph for `then_branch`
+    else_body = gs.Graph(
+        nodes=[else_const_node], name="else_body", inputs=[], outputs=[else_out]
+    )  # Graph for `else_branch`
+
+    res = gs.Variable("res", dtype=np.float32, shape=None)  # shape is data-dependent
+
+    if_node = gs.Node(
+        op="If",
+        name="If_Node",
+        inputs=[cond],
+        outputs=[res],
+        attrs={"then_branch": then_body, "else_branch": else_body},
+    )
+    graph.nodes = [if_node]
+
+    out = graph.matmul(
+        res, gs.Constant("z", values=np.array(Z, dtype=np.float32)), name="MatMul"
+    )
+    out.dtype = np.float32
+
+    graph.inputs = [cond]
+    graph.outputs = [out]
+
+    save(graph, "bad_graph_conditionally_invalid.onnx")
+
+
+make_bad_graph_conditionally_invalid()
+
+
+### Bad GraphProto ###
+### Graphs that break the ONNX Specification for GraphProto ###
+
+
+# Generates a model where the GraphProto has no name.
+#
+# This is invalid as ONNX Specification requires that the GraphProto has a name.
+#
+def make_bad_graph_with_no_name():
+    DTYPE = np.float32
+    SHAPE = (4, 4)
+
+    inp = gs.Variable("inp", dtype=DTYPE, shape=SHAPE)
+
+    graph = gs.Graph(inputs=[inp], name="")
+    out = graph.add(inp, inp)
+    out.dtype = DTYPE
+    graph.outputs = [out]
+
+    save(graph, "bad_graph_with_no_name.onnx")
+
+
+make_bad_graph_with_no_name()
+
+
+# Generates a model where the GraphProto has no imports.
+#
+# This is invalid as ONNX Specification requires that the GraphProto has at least one import.
+#
+def make_bad_graph_with_no_import_domains():
+    DTYPE = np.float32
+    SHAPE = (4, 4)
+
+    inp = gs.Variable("inp", dtype=DTYPE, shape=SHAPE)
+
+    graph = gs.Graph(inputs=[inp], import_domains=[])
+    out = graph.add(inp, inp)
+    out.dtype = DTYPE
+    graph.outputs = [out]
+
+    save(graph, "bad_graph_with_no_import_domains.onnx")
+
+
+make_bad_graph_with_no_import_domains()
+
+
+# Generates a model where the inputs (value info) of graph are duplicates.
+#
+# This is invalid as ONNX Specification requires that the (value info) inputs of a graph are unique.
+#
+#    inp
+#    / \
+#    Add
+#     |
+#    out
+#
+def make_bad_graph_with_dup_value_info():
+    DTYPE = np.float32
+    SHAPE = (4, 4)
+
+    inp = gs.Variable("inp", dtype=DTYPE, shape=SHAPE)
+
+    graph = gs.Graph(inputs=[inp, inp])
+    out = graph.add(inp, inp)
+    out.dtype = DTYPE
+    graph.outputs = [out]
+
+    save(graph, "bad_graph_with_dup_value_info.onnx")
+
+
+make_bad_graph_with_dup_value_info()
+
+
+# Generates a model with mult-level errors.
+# The model is invalid because of graph-level error (no name) and node-level error (incompatible inputs).
+def make_bad_graph_multi_level_errors():
+    DTYPE = np.float32
+    SHAPE = (4, 5)
+
+    inp1 = gs.Variable("inp1", dtype=DTYPE, shape=SHAPE)
+    inp2 = gs.Variable("inp2", dtype=DTYPE, shape=SHAPE)
+
+    graph = gs.Graph(inputs=[inp1, inp2], name="")  # graph-level error: empty name
+    out = graph.matmul(inp1, inp2)  # node-level error: incompatible inputs
+    out.dtype = DTYPE
+    out.shape = []  # we need to specify this so GS creates valid ONNX model.
+    graph.outputs = [out]
+
+    save(graph, "bad_graph_with_multi_level_errors.onnx")
+
+
+make_bad_graph_multi_level_errors()
+
+
+# Generates a model where graph has multiple node names with same non-empty string.
+def make_bad_graph_with_duplicate_node_names():
+    DTYPE = np.float32
+    SHAPE = (4, 5)
+
+    inp = gs.Variable("inp", dtype=DTYPE, shape=SHAPE)
+
+    graph = gs.Graph(inputs=[inp], name="bad_graph_with_duplicate_node_names")
+    inter1 = graph.identity(inp, name="identical")
+    out = graph.identity(
+        inter1, name="identical"
+    )  # node-level error: duplicate node names
+    graph.outputs = [out]
+
+    save(graph, "bad_graph_with_duplicate_node_names.onnx")
+
+
+make_bad_graph_with_duplicate_node_names()
+
+
+# Generates a model where the graph has a subgraph matching toyPlugin's graph pattern
+def make_graph_with_subgraph_matching_toy_plugin():
+    i0 = gs.Variable(name="i0", dtype=np.float32)
+    i1 = gs.Variable(name="i1", dtype=np.float32)
+    i2 = gs.Variable(name="i2", dtype=np.float32)
+    i3 = gs.Variable(name="i3", dtype=np.float32)
+    i4 = gs.Variable(name="i4", dtype=np.float32)
+
+    o1 = gs.Variable(name="o1", dtype=np.float32)
+    o2 = gs.Variable(name="o2", dtype=np.float32)
+
+    O_node = gs.Node(op="O", inputs=[i0], outputs=[i1], name="n1")
+    A_node = gs.Node(op="A", inputs=[i1], outputs=[i2], name="n2")
+    B_node = gs.Node(op="B", inputs=[i1], outputs=[i3], name="n3")
+    C_node = gs.Node(op="C", inputs=[i2, i3], outputs=[i4], attrs={"x": 1}, name="n4")
+    D_node = gs.Node(op="D", inputs=[i4], outputs=[o1], name="n5")
+    E_node = gs.Node(op="E", inputs=[i4], outputs=[o2], name="n6")
+
+    graph = gs.Graph(
+        nodes=[O_node, A_node, B_node, C_node, D_node, E_node],
+        inputs=[i0],
+        outputs=[o1, o2],
+    )
+
+    save(graph, "toy_subgraph.onnx")
+
+
+make_graph_with_subgraph_matching_toy_plugin()
+
+
+# Generates the following Graph
+#
+# The input to the Transpose op is an initializer
+#
+#    Transpose
+#       |
+#      MatMul
+#       |
+#      out
+#
+def make_transpose_matmul():
+    M = 8
+    N = 8
+    K = 16
+    a = gs.Variable("a", shape=(M, K), dtype=np.float32)
+    g = gs.Graph(inputs=[a], opset=13)
+    val = np.random.uniform(-3, 3, size=K * N).astype(np.float32).reshape((N, K))
+    b = gs.Constant("b", values=val)
+    b_transpose = g.transpose(b, name="transpose")
+    c = g.matmul(a, b_transpose, name="matmul")
+    c.dtype = np.float32
+    g.outputs = [c]
+
+    save(g, "transpose_matmul.onnx")
+
+
+make_transpose_matmul()
+
+
+# Generates the following Graph
+#
+# The input to the QuantizeLinear op is an initializer
+#
+#    QuantizeLinear
+#       |
+#    DequantizeLinear
+#       |
+#      Conv
+#       |
+#      out
+#
+def make_qdq_conv():
+    x = (
+        np.random.uniform(-3, 3, size=3 * 3 * 130)
+        .astype(np.float32)
+        .reshape((1, 3, 3, 130))
+    )
+    y_scale = np.array([2, 4, 5], dtype=np.float32)
+    y_zero_point = np.array([84, 24, 196], dtype=np.uint8)
+    x_const = gs.Constant("x", values=x)
+    y_scale_const = gs.Constant("y_scale", values=y_scale)
+    y_zero_point_const = gs.Constant("y_zero_point", values=y_zero_point)
+
+    weight = gs.Constant("Weights_0", values=np.ones((3, 3, 3, 3), dtype=np.float32))
+
+    g = gs.Graph(inputs=[], opset=13)
+    q_layer = g.quantize_linear(x_const, y_scale_const, y_zero_point_const)
+    dq_layer = g.dequantize_linear(q_layer, y_scale_const, y_zero_point_const)
+    out = g.conv(dq_layer, weight, [3, 3], name="Conv_0")
+    out.dtype = np.float32
+    g.outputs = [out]
+
+    save(g, "qdq_conv.onnx")
+
+
+make_qdq_conv()
+
+
+def make_weightless_network(model_name):
+    ipath = ONNX_MODELS[model_name].path
+    opath = os.path.join(CURDIR, "weightless." + model_name + ".onnx")
+    cmd = [f"polygraphy surgeon weight-strip {ipath} -o {opath}"]
+    subprocess.run(cmd, shell=True)
+
+
+make_weightless_network("matmul.fp16")
+make_weightless_network("matmul.bf16")
+make_weightless_network("sparse.matmul")
+make_weightless_network("conv")
+make_weightless_network("sparse.conv")
+make_weightless_network("transpose_matmul")
+make_weightless_network("qdq_conv")

tools/Polygraphy/tests/models/matmul.bf16.i32data.onnx

@@ -0,0 +1,17 @@
+onnx-example:�
+
+
A
+
B
C"MatMul
+test-model*ר@*לה‏נxך}ֳ~ת��|´~₪~�{�‏³‏¶~ק~כ~��ּz‹‏מ}�|נ‎�~±‏ע~×�­|¥{�‏ֲ~¨}«�ֲ‎¯�¼~³תˆת�{…|�{י‏ָ}ֽ}ר‏¸‏§‏„~ױ~­}¥‎�~ֻ‏ת}ש‎ֵ}ױ~צ‎�‏י�¯~‘‏‘}ֳ‏¦‏ˆ‏ן‎¬‎»|£‏ּ}ִ�ֵ‏��£~“z�~ד�צ‎ִ‏�‏ל‎ה}׳‎ם‎פ‏‘‎�‏ד~¹‏�‏²‏נ|�‏ֿ~ׂyה{ּ~™}±|�‎‏‏‎‎ה‏‡|ו~ֽ‎¼‏ח�ֱ‎ױy“}­‏©‎ֿ~‹}´‎»�ƒ�×}¸‏י~›‏’}™‏�|�|›zק‏…‏©ת‘~ן‎ע‎ˆ‏ב~”‎ר|ײ~ב‎�yי~ַ~ן~ַ‏ק‎ט~ֱ~®�²�ל‏’‏ס‏ח‏�‎ֶ~�|�~€}ֱ�׀}£|ֳ}�‎¹}�‎א~₪�ק~½�¦w²‏§}¥}ז‏›~’‏²‏€‎�}נ~�~יתק~¸‏®~�‏¥‏¿‏���|ַ~ }�‏�‏�תˆ��‏–‏½}‚�¨‎ִ|ד‎·|•zס}׳|ײ~װ}“‎»�ס‏‰�‡‏ּ‏§�»��~³‏�ת›‏כ�ג}��ע‏¸‎״~†{¡~°|³y׃‎—{א~�‎‰|�~ת‏ׂ‎ָ‏�תח‏ֱz¿‏�}ש‏¨~ֻz‘}�~�}—wע~€}§z‰‏ִ~¬|ִ��|��…‏¦�ּ|‰~ם~ ‎¸‏ֻ~¾�³‏”~»{ר{ׂ~¥�בzִ|�~׀~¿}ג‏»�™~כz�‎”}�~�~§‎ת‏ץ~א~�~÷‏ֺ~ˆ|א�׀�®‎©�ו~½{ˆ}³‏¹}²~–‎�~ש~ח~‘��}ִ‏ע‎²{�}ר‏ײ}³��}�‏×‏µ}ת|§}¨zן‎ֱ‏‡|¡}׃ת«‏¢}®‏��ֽ~ב‎‰‏ֲ~ש‏��’‏�}´‏•‏…�ז~ֳ|¥ת¦�³‎²~ֵתה‏פ‏¡�‎~ש‎װ{�~ר~׃‎ף�ד~£|ֱ‎�‏ֺ~¬ת£~ײ|׳‏»‏ֱ~×‏°~€�’}¹|ײ}�‎₪‎ִ‏ּ}װ}„‏ה~�‏�~‰~†‏ן‏�~�~ֱ~ָ‎‘‏ז~ƒ‏±|ח}™|׃~„‎ש‎†}²~�}•}�~ר~ˆ~´ת�‏ױ|צ~ל~�|ׂ‏ƒ‏�‎�~�}ה~ֲy‘|₪‏ג~ג��|ג~³}ײ‏ֺ‏ג{ע}ˆ}�|”~´~ד~�‎ק‎׃}�~ְz®‏ֹ~��ז‏¹~�~·‏ׂ}־|ׂ‏ֵ‏¡�ׂ‏ƒ~‎~ֱ‎�‏§‎¿‎מ~ֳ~…zד‎‚}ֽ}”�¨~¦‏ח‎ו~ׁ‎½~™‎–~÷‏ֽ‎ƒ‏›~°�ק}‰}�‏ס}÷}°‎ה{ד‎¹|׀��‎�zװ~‚|§}ב�µ~†‏ץ~¿�µ�¦‎�|„‏��ו‏»|ׂ~�y�zג}װ��ת�‏ל}בzב{ץתה‏ע|„z ‏’}ם‎לר׃ת¥~µ}µ‏»‏₪~¢‏ׁ~ש}�‏°~�‎¼�ׂ{™‎¶|´‎�ק��½}»‏”}�~´‎¡‏´‏מ‏‚‎“~¢zף|־‏„z¡}²‏�‏’‎ƒ~װ}�|·~€}ַ~ֵ~•‏ַz�ת¡~§‏ƒ|ה|ן�¿|ם~¢~₪~¢‏¥‏ֶ‏ִ‏ח�ץ{�‎ׁ~נ‏ָ}��ֺ~‚~א}ל}�}®y�~ֶ~¸‏¹‎‹~�~²‏±}ץ|�‎ררע}¼~×~™zֿש¨‎ַt¡{ֱ‏­{¨}ַ{�~ּ~�~ˆ}�~×‏·~¿|כ}§~€ִ‏±~א‎צ~�~ו‏¹‏�‏צ�א}ק��‏¯�‏‎ך~÷‏�~ז}°~­‎װ‎נ}ֲ~•~†~›~§�±|�‏¨~�‏�‎«~‏}�}ה‏ˆ~�‎¨‎ט~�}ח|ה‎¿~’‎©�°�±z‰~ם‎�{™‏†‎א~א~÷~ו‎�}ך|�~ˆ~�|´}ת}±~‏‎�‎‚|ֱ}˜‎׃�ֳ‏˜~©~«~‏|כ}„‎�‎“~”~ב~ֺ~”‏¹‏ם~ׁ�½‏£}»~ˆ‎�‎¦‎ס~‏}�‎י‎�~¢|“}�~ֵ‏¡‏י~¸~ל}–ר¹}ת�™�ֲ|ֵש„‏‚|װ‎ך‏¹‏«�¾|ֺ�×‎ח{ך}ֺ‎…��‏‹~ח}«‏ד�ע|ל‎«}’zנ‏ֵ~ק~¹‏ֵ‎´‏‘}”{¸ש׳}�~·‏£‏�}כ‏•‏₪‏ֿ}—‏�‏ֿ‎ֻ�€‏°}ת‎²}„~�‎¨}ׂ}”}¶�‚{�~µ~ת‏¥}ץ}½~מzר}ף~ז��‎�~»~²‏‚‏צ�˜‏�~ƒ~¶‏§�ב‏ִ}י‏�~¦~ֲ‎¥}„{ֵ‎¹‎��®}ַ‎¹}·}„‏±~¹}¥~¢‎ס‏ײ|ֺ~ו~”‏¢‏–}�‏†~‹~ָ‏¥‏�‏ת~‘‎ו��~ט}��²|ל‎ֵ‏‰}�‏ל~�‏�‎�‏�}”}ח‎ְ~¨‎ג‏ן‎פ}־}«‏§‏�{¨{�}“‎��«~ֶ~�~ש‎¼‏׃~־~׀~ו‏²‏�~–|�}ֱ~א�ֱ‎»}�~‘‏ ‏¶‏ֻ~־~¾‏װ~¯zְ~’‏ּ|¸{�~ע‏ֵ~ֵ~ץ~¡‎״‏�~�~·yִ{��‘~�~¡�¾~�ש¨~‚�­‏…}ֲ‏ׂ�פ‎�~�‏ּ‏ש‎¸‏¥|×‏�‏‰‎כ��‏ױ‏ש{ן~“‎¥~‎~ג~¸�ּ~ע~°‏�‏‚‏‡‏��פ|’}ד‏¯~ֹ‏‚‎—|•‎ֵ�“‎�‎…‎¼ת�~ָ‏�{»~‘z¯��~ײ‎ת~״‏¼|־‏†‏‚ק¶�ז}ױ�־‎ֵ~ק‏ה|ֵ‎״‏־‏•|ע‎ןy×~�‏ֻ}ƒz��ֶ‏�{ת‎א}ל}ׂ�ֱ‏•ת׃}»{©}B
BZ
+
A
+
+
+@b
+
C
+
+
+j
+
B
+
+@
+B

tools/Polygraphy/tests/models/matmul.exclude_list.txt

@@ -0,0 +1 @@
+b

tools/Polygraphy/tests/models/meta.py

@@ -0,0 +1,465 @@
+#
+# 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 os
+
+import numpy as np
+import tensorrt as trt
+
+from polygraphy import mod, util
+from polygraphy.backend.common import BytesFromPath
+from polygraphy.backend.onnx import OnnxFromPath
+from polygraphy.backend.tf import GraphFromFrozen
+from polygraphy.common import TensorMetadata
+from polygraphy.datatype import DataType
+
+
+def model_path(name=None):
+    path = os.path.abspath(os.path.dirname(__file__))
+    if name is not None:
+        path = os.path.join(path, name)
+    return path
+
+
+class Model:
+    def __init__(
+        self, path, LoaderType, check_runner, input_metadata=None, ext_data=None
+    ):
+        self.path = path
+        self.loader = LoaderType(self.path)
+        self.check_runner = check_runner
+        self.input_metadata = input_metadata
+        self.ext_data = ext_data
+
+
+def check_tf_identity(runner):
+    feed_dict = {
+        "Input:0": np.random.random_sample(size=(1, 15, 25, 30)).astype(np.float32)
+    }
+    outputs = runner.infer(feed_dict)
+    assert np.all(outputs["Identity_2:0"] == feed_dict["Input:0"])
+
+
+MODELS_DIR = os.path.join(os.path.dirname(__file__))
+
+TF_MODELS = {
+    "identity": Model(
+        path=model_path("tf_identity.pb"),
+        LoaderType=GraphFromFrozen,
+        check_runner=check_tf_identity,
+    ),
+}
+
+
+def check_identity(runner):
+    feed_dict = {"x": np.random.random_sample(size=(1, 1, 2, 2)).astype(np.float32)}
+    outputs = runner.infer(feed_dict)
+    assert np.all(outputs["y"] == feed_dict["x"])
+
+
+def check_identity_identity(runner):
+    feed_dict = {"X": np.random.random_sample(size=(64, 64)).astype(np.float32)}
+    outputs = runner.infer(feed_dict)
+    assert np.all(outputs["identity_out_2"] == feed_dict["X"])
+
+
+def check_dynamic_identity(runner, shapes):
+    feed_dict = {"X": np.random.random_sample(size=shapes["X"]).astype(np.float32)}
+    outputs = runner.infer(feed_dict)
+    assert np.array_equal(outputs["Y"], feed_dict["X"])
+
+
+def check_empty_tensor_expand(runner, shapes):
+    shape = shapes["new_shape"]
+    feed_dict = {
+        "data": np.zeros(shape=(2, 0, 3, 0), dtype=np.float32),
+        "new_shape": np.array(
+            shape,
+            dtype=(
+                np.int32
+                if mod.version(trt.__version__) < mod.version("9.0")
+                else np.int64
+            ),
+        ),
+    }
+    outputs = runner.infer(feed_dict)
+    # Empty tensor will still be empty after broadcast
+    assert outputs["expanded"].shape == shape
+    assert util.volume(outputs["expanded"].shape) == 0
+
+
+def check_reshape(runner):
+    feed_dict = {"data": np.random.random_sample(size=(1, 3, 5, 5)).astype(np.float32)}
+    outputs = runner.infer(feed_dict)
+    assert np.all(outputs["output"] == feed_dict["data"].ravel())
+
+
+def check_residual_block(runner, shapes):
+    feed_dict = {
+        "gpu_0/data_0": np.random.random_sample(size=shapes["gpu_0/data_0"]).astype(
+            np.float32
+        )
+    }
+    # Confirm inference can go through without error
+    outputs = runner.infer(feed_dict)
+
+
+def check_matmul_2layer(runner, shape=(2, 8)):
+    feed_dict = {
+        "onnx::MatMul_0": np.random.random_sample(size=shape).astype(np.float32)
+    }
+    # Confirm inference can go through without error
+    outputs = runner.infer(feed_dict)
+
+
+def no_check_implemented(runner):
+    raise NotImplementedError("No check_runner implemented for this model")
+
+
+ONNX_MODELS = {
+    "identity": Model(
+        path=model_path("identity.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=check_identity,
+        input_metadata=TensorMetadata().add(
+            "x", dtype=DataType.FLOAT32, shape=(1, 1, 2, 2)
+        ),
+    ),
+    "identity_identity": Model(
+        path=model_path("identity_identity.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=check_identity_identity,
+    ),
+    "dynamic_identity": Model(
+        path=model_path("dynamic_identity.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=check_dynamic_identity,
+        input_metadata=TensorMetadata().add(
+            "X", dtype=DataType.FLOAT32, shape=(1, 1, -1, -1)
+        ),
+    ),
+    "identity_multi_ch": Model(
+        path=model_path("identity_multi_ch.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+        input_metadata=TensorMetadata().add(
+            "x", dtype=DataType.FLOAT32, shape=(2, 4, 3, 3)
+        ),
+    ),
+    "empty_tensor_expand": Model(
+        path=model_path("empty_tensor_expand.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=check_empty_tensor_expand,
+    ),
+    "and": Model(
+        path=model_path("and.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "scan": Model(
+        path=model_path("scan.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "pow_scalar": Model(
+        path=model_path("pow_scalar.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "dim_param": Model(
+        path=model_path("dim_param.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "tensor_attr": Model(
+        path=model_path("tensor_attr.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "identity_with_initializer": Model(
+        path=model_path("identity_with_initializer.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "const_foldable": Model(
+        path=model_path("const_foldable.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "reshape": Model(
+        path=model_path("reshape.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=check_reshape,
+    ),
+    "reducable": Model(
+        path=model_path("reducable.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+        input_metadata=TensorMetadata()
+        .add("X0", shape=(1,), dtype=DataType.FLOAT32)
+        .add("Y0", shape=(1,), dtype=DataType.FLOAT32),
+    ),
+    "reducable_with_const": Model(
+        path=model_path("reducable_with_const.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "ext_weights": Model(
+        path=model_path("ext_weights.onnx"),
+        LoaderType=OnnxFromPath,
+        check_runner=no_check_implemented,
+        ext_data=model_path("data"),
+    ),
+    "ext_weights_same_dir": Model(
+        path=model_path(os.path.join("ext_weights_same_dir", "ext_weights.onnx")),
+        LoaderType=OnnxFromPath,
+        check_runner=no_check_implemented,
+        ext_data=model_path("ext_weights_same_dir"),
+    ),
+    "capability": Model(
+        path=model_path("capability.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "instancenorm": Model(
+        path=model_path("instancenorm.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "add_with_dup_inputs": Model(
+        path=model_path("add_with_dup_inputs.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "needs_constraints": Model(
+        path=model_path("needs_constraints.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+        input_metadata=TensorMetadata().add(
+            "x", dtype=DataType.FLOAT32, shape=(1, 1, 256, 256)
+        ),
+    ),
+    "constant_fold_bloater": Model(
+        path=model_path("constant_fold_bloater.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "renamable": Model(
+        path=model_path("renamable.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "cleanable": Model(
+        path=model_path("cleanable.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "nonzero": Model(
+        path=model_path("nonzero.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "inp_dim_val_not_set": Model(
+        path=model_path("inp_dim_val_not_set.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "multi_output": Model(
+        path=model_path("multi_output.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "unbounded_dds": Model(
+        path=model_path("unbounded_dds.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "loop": Model(
+        path=model_path("loop.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "matmul.fp16": Model(
+        path=model_path("matmul.fp16.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "matmul": Model(
+        path=model_path("matmul.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "sparse.matmul": Model(
+        path=model_path("sparse.matmul.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "matmul.bf16": Model(
+        path=model_path("matmul.bf16.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "matmul.bf16.i32data": Model(
+        path=model_path("matmul.bf16.i32data.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "matmul_2layer": Model(
+        path=model_path("matmul_2layer.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=check_matmul_2layer,
+    ),
+    "unsorted": Model(
+        path=model_path("unsorted.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "conv": Model(
+        path=model_path("conv.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "sparse.conv": Model(
+        path=model_path("sparse.conv.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "no_op_reshape": Model(
+        path=model_path("no_op_reshape.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "bad_graph_with_dup_value_info": Model(
+        path=model_path("bad_graph_with_dup_value_info.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "bad_graph_with_no_name": Model(
+        path=model_path("bad_graph_with_no_name.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "bad_graph_with_no_import_domains": Model(
+        path=model_path("bad_graph_with_no_import_domains.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "bad_graph_with_parallel_invalid_nodes": Model(
+        path=model_path("bad_graph_with_parallel_invalid_nodes.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "bad_graph_conditionally_invalid": Model(
+        path=model_path("bad_graph_conditionally_invalid.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "custom_op_node": Model(
+        path=model_path("custom_op_node.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "bad_graph_with_duplicate_node_names": Model(
+        path=model_path("bad_graph_with_duplicate_node_names.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "bad_graph_with_multi_level_errors": Model(
+        path=model_path("bad_graph_with_multi_level_errors.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "empty": Model(
+        path=model_path("empty.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "residual_block": Model(
+        path=model_path("residual_block.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=check_residual_block,
+    ),
+    "graph_with_subgraph_matching_toy_plugin": Model(
+        path=model_path("toy_subgraph.onnx"), 
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "transpose_matmul": Model(
+        path=model_path("transpose_matmul.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "qdq_conv": Model(
+        path=model_path("qdq_conv.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "weightless.matmul.fp16": Model(
+        path=model_path("weightless.matmul.fp16.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "weightless.matmul.bf16": Model(
+        path=model_path("weightless.matmul.bf16.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "weightless.conv": Model(
+        path=model_path("weightless.conv.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "weightless.sparse.matmul": Model(
+        path=model_path("weightless.sparse.matmul.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "weightless.sparse.conv": Model(
+        path=model_path("weightless.sparse.conv.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "weightless.transpose_matmul": Model(
+        path=model_path("weightless.transpose_matmul.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "weightless.qdq_conv": Model(
+        path=model_path("weightless.qdq_conv.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "roialign": Model(
+        path=model_path("roialign.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "attention": Model(
+        path=model_path("attention.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "multi_attention": Model(
+        path=model_path("multi_attention.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+    "attention_same_qkv": Model(
+        path=model_path("attention_same_qkv.onnx"),
+        LoaderType=BytesFromPath,
+        check_runner=no_check_implemented,
+    ),
+}

tools/Polygraphy/tests/models/plugins/toyPlugin/pattern.py

@@ -0,0 +1,48 @@
+from polygraphy import mod
+from typing import List,Dict
+gs = mod.lazy_import("onnx_graphsurgeon>=0.5.0")
+
+def get_plugin_pattern():
+    """
+    Toy plugin pattern:
+        A     B
+        \   /
+          C, attrs['x'] < 2.0
+        /   \
+        D     E
+    """
+    pattern = gs.GraphPattern()
+    in_0 = pattern.variable()
+    in_1 = pattern.variable()
+    a_out = pattern.add("Anode", "A", inputs=[in_0])
+    b_out = pattern.add("Bnode", "B", inputs=[in_1])
+    check_function = lambda node : node.attrs["x"] < 2.0
+    c_out = pattern.add("Cnode", "C", inputs=[a_out, b_out], check_func=check_function)
+    d_out = pattern.add("Dnode", "D", inputs=[c_out])
+    e_out = pattern.add("Enode", "E", inputs=[c_out])
+    pattern.set_output_tensors([d_out, e_out])
+
+    return pattern
+
+def get_matching_subgraphs(graph) -> List[Dict[str,str]]:
+    gp = get_plugin_pattern()
+    matches = gp.match_all(graph)
+    ans = []
+    for m in matches:
+        # save the input and output tensor names of the matching subgraph(s)
+        input_tensors = list(set([ip_tensor.name for ip_tensor in m.inputs]))
+        output_tensors = list(set([op_tensor.name for op_tensor in m.outputs]))
+
+        attrs = {"ToyX": int(m.get("Cnode").attrs["x"]) * 2}
+        ioa = {
+            'inputs':input_tensors,
+            'outputs':output_tensors,
+            'attributes':attrs
+        }
+        ans.append(ioa)
+    return ans
+
+def get_plugin_metadata() -> Dict[str,str]:
+    return {'name':'toyPlugin',
+            'op':'CustomToyPlugin',
+            }