chore: import upstream snapshot with attribution

samples/python/README.md

@@ -0,0 +1,64 @@
+General Setup Guide for Samples
+==============================
+
+
+## Download Sample Data
+
+Install the tool dependencies via `python3 -m pip install -r requirements.txt`.
+
+Invoke [downloader.py](downloader.py) to download the data with
+a command like the one below if `download.yml` is present in the
+sample directory ([example](onnx_packnet/download.yml)).
+
+```sh
+downloader.py -d /path/to/data/dir -f /path/to/download.yml
+```
+
+The data directory i.e. `/path/to/data/dir` is a centralized directory
+to store data of all samples. So you can use same one for all samples.
+It can be provided by either `-d /path/to/data/dir` or the environment variable
+`$TRT_DATA_DIR`, where the `-d` has higher priority.
+
+Remember to use `-d` or `$TRT_DATA_DIR` when running sample scripts
+that rely on downloaded data. Scripts will abort if no downloaded data
+is found in data directory. (`$TRT_DATA_DIR` will be much simplier.)
+An error will be thrown if the data is not properly setup.
+
+The `download.yml` file is owned by the sample which describes the sample
+name, the path, URL and checksum of the data files that are required by the sample.
+
+
+**Notes for sample developers**
+
+To use the downloaded data files, integrate the code segment like below into
+the sample code, and obtain the path to the data file by passing the `path`
+as specified in the associated `download.yml` file of the sample.
+
+```py
+TRT_DATA_DIR = None
+
+def getFilePath(path):
+    global TRT_DATA_DIR
+    if not TRT_DATA_DIR:
+        parser = argparse.ArgumentParser(description="Convert PackNet to ONNX")
+        parser.add_argument('-d', '--data', help="Specify the data directory where it is saved in. $TRT_DATA_DIR will be overwritten by this argument.")
+        args, _ = parser.parse_known_args()
+        TRT_DATA_DIR = os.environ.get('TRT_DATA_DIR', None) if args.data is None else args.data
+    if TRT_DATA_DIR is None:
+        raise ValueError("Data directory must be specified by either `-d $DATA` or environment variable $TRT_DATA_DIR.")
+
+    fullpath = os.path.join(TRT_DATA_DIR, path)
+    if not os.path.exists(fullpath):
+        raise ValueError("Data file %s doesn't exist!" % fullpath)
+
+    return fullpath
+```
+
+**Python Version Support**
+
+All Python samples are expected to be run with Python>=3.10. It is not recommended to use any lower version as there may be compatibility issues.
+
+# Changelog
+
+August 2025
+Removed support for Python versions < 3.10.

samples/python/aliased_io_plugin/README.md

@@ -0,0 +1,91 @@
+# Utilizing a plugin with aliased I/O to realize in-place updates
+
+## Description
+
+This sample, `aliased_io_plugin`, implements a Python-based plugin for an in-place scatter-add operation.
+
+Scatter-add "scatters" a set of source values into memory locations based on a given set of indices and adds together those values mapped to the same location.
+
+## How does this sample work?
+
+This sample creates and runs a TensorRT engine demonstrating an example commonly encountered with Graph Neural Networks (GNNs). In GNNs, the features associated with the neighbors of each node is aggregated with an order-independent operation (e.g. sum, product), averaged by the size of the neighborhood, and then run through a classifier to determine a property of interest; example applications of GNNs include the modeling of social networks and building recommendation systems.
+
+Here, we use an addition as the aggregation function; therefore, we build a network containing a Scatter-add plugin node. It receives a "source" tensor containing the features of the neighbors of each node, and an "index" tensor denoting the index of each such node. For example, consider the following graph:
+
+![alt text](aliased_io_gnn.png "GNN example")
+
+For simplicity, in this example, and in the sample in general, we utilize scalar features at each node. The "source" could be represented as a flattened tensor `[1.0, 3.0, 5.0, 7.0, 1.0, 3.0]` while the corresponding source nodes are `[1, 2, 3, 0, 2, 3]`. It is clear that the Scatter-add should yield `[7.0, 1.0, 4.0, 8.0]`. This result is then normalized by the number of neighbors of each node and then fed into a simple dense layer followed by ReLU activation.
+
+### Implementing an in-place Scatter-add plugin using `IPluginV3OneBuildV2` interface
+
+Before the introduction of `IPluginV3OneBuildV2` interface, TensorRT plugin inputs were to be treated as read-only. In-place optimizations (output written to an input) and operations that inherently required an input to be modified, were kept out-of-reach due to this limitation.
+
+In the Scatter-add operation, an in-place operation is useful because a node of interest may have some pre-conditions that require the neighborhood aggregation to be combined with a bias. Another use case is in hierarchical aggregation where higher-layer features may have to be integrated as well.
+
+To allow writes to the input, `IPluginV3OneBuildV2` interface provides an API to declare certain input-output pairs as being aliased. In this case, the first output of the plugin and the first input are aliased, so we may declare:
+```py
+def get_aliased_input(self, output_index: int):
+	if output_index == 0:
+		return 0
+
+	return -1
+```
+A return value of `-1` indicates that that `output_index` is not aliased to any input.
+
+This new method `get_aliased_input` is the only difference between `IPluginV3OneBuildV2` and `IPluginV3OneBuild`. As part of the `V3_ONE` set of capability interfaces, `IPluginV3OneBuildV2` may be used in conjunction with `IPluginV3OneCore` and `IPluginV3OneRuntime`.
+
+### Creating network and building the engine
+
+To add the plugin to the network, the `INetworkDefinition::add_plugin_v3()` method is used.
+
+For subsequent averaging and classification steps, TensorRT ElementWise, MatrixMultiply, Activation and SoftMax layers are used.
+
+## Running the sample
+
+1.  Run the sample to create a TensorRT inference engine and run inference:
+    `python3 aliased_io_plugin.py [-h] [--precision {fp32,fp16}] [--node_features NODE_FEATURES] [--edges EDGES] [--num_classes NUM_CLASSES] [--validate] [--seed SEED]`
+
+2.  If the `--validate` flag was passed, verify that the sample ran successfully. If the sample runs successfully, you should see the following message:
+     ```
+    Validation against reference successful!
+    ```
+
+### Sample `--help` options
+
+To see the full list of available options and their descriptions, use the `-h` or `--help` command line option.
+
+
+# Additional resources
+
+The following resources provide a deeper understanding about the V3 TensorRT plugins and the Scatter-Add operation:
+
+**ScatterElements**
+- [ONNX: ScatterElements](https://onnx.ai/onnx/operators/onnx__ScatterElements.html)
+
+**TensorRT plugins**
+- [Extending TensorRT with Custom Layers](https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#extending)
+- [TensorRT Python-based Plugins](https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/#add_custom_layer_python)
+
+**Other documentation**
+- [Introduction To NVIDIA’s TensorRT Samples](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sample-support-guide/index.html#samples)
+- [Working With TensorRT Using The Python API](https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/#python_topics)
+- [NVIDIA’s TensorRT Documentation Library](https://docs.nvidia.com/deeplearning/sdk/tensorrt-archived/index.html)
+
+# License
+
+For terms and conditions for use, reproduction, and distribution, see the [TensorRT Software License Agreement](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sla/index.html) documentation.
+
+# Changelog
+
+October 2025
+Migrate to strongly typed APIs.
+
+August 2025
+Removed support for Python versions < 3.10.
+
+August 2024
+This is the first version of this `README.md` file.
+
+# Known issues
+
+There are no known issues in this sample.

samples/python/aliased_io_plugin/aliased_io_plugin.py

@@ -0,0 +1,452 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025 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 torch
+import triton
+import triton.language as tl
+
+import tensorrt as trt
+import cupy as cp
+import numpy as np
+import ast
+
+from polygraphy.backend.trt import (
+    CreateConfig,
+    TrtRunner,
+    create_network,
+    engine_from_network,
+)
+
+import argparse
+
+import logging
+
+logging.basicConfig(level=logging.INFO)
+logging.getLogger("AliasedIOPlugin").setLevel(logging.INFO)
+log = logging.getLogger("AliasedIOPlugin")
+
+import sys
+
+# An OpenAI Triton kernel to both perform the scatter-add and counts of each index
+@triton.jit
+def scatter_add_kernel(
+    self_ptr,
+    src_ptr,  # Source array
+    index_ptr,  # Indices
+    n_elements,  # Number of elements in the source/indices array
+    n_labels,  # Number of labels (distinct indices)
+    counts,  # Output counts of each distinct index
+    BLOCK_SIZE: tl.constexpr,
+    BLOCK_SIZE_C: tl.constexpr,
+):
+    pid = tl.program_id(axis=0)
+    block_start = pid * BLOCK_SIZE
+    offsets = block_start + tl.arange(0, BLOCK_SIZE)
+
+    mask = offsets < n_elements
+
+    # Load the source values and indices
+    src = tl.load(src_ptr + offsets, mask=mask)
+    indices = tl.load(index_ptr + offsets, mask=mask)
+
+    # Iterate over n_labels
+    for i in range(0, BLOCK_SIZE_C):
+        idx = i + tl.program_id(1) * BLOCK_SIZE_C + 1
+        if idx <= n_labels:
+            l_mask = indices == idx
+            # Perform the scatter-add operation
+            tl.atomic_add(self_ptr + idx - 1, tl.sum(tl.where(l_mask, src, 0)))
+            # Update count for idx
+            tl.atomic_add(counts + idx - 1, tl.sum(tl.where(l_mask, 1, 0)))
+
+
+def volume(d):
+    return np.prod(d)
+
+
+class UnownedMemory:
+    def __init__(self, ptr, shape, dtype):
+        mem = cp.cuda.UnownedMemory(ptr, volume(shape) * cp.dtype(dtype).itemsize, self)
+        cupy_ptr = cp.cuda.MemoryPointer(mem, 0)
+        self.d = cp.ndarray(shape, dtype=dtype, memptr=cupy_ptr)
+
+
+class ScatterAddPlugin(
+    trt.IPluginV3,
+    trt.IPluginV3OneCore,
+    trt.IPluginV3OneBuildV2,
+    trt.IPluginV3OneRuntime,
+):
+    def __init__(self, fc=None):
+        trt.IPluginV3.__init__(self)
+        trt.IPluginV3OneCore.__init__(self)
+        trt.IPluginV3OneBuildV2.__init__(self)
+        trt.IPluginV3OneRuntime.__init__(self)
+
+        self.plugin_namespace = ""
+        self.plugin_name = "ScatterAddPlugin"
+        self.plugin_version = "1"
+        self.num_outputs = 2
+
+    def get_capability_interface(self, type):
+        return self
+
+    def get_output_data_types(self, input_types):
+        self.type = input_types[0]
+        return [input_types[0], trt.int64]
+
+    def get_fields_to_serialize(self):
+        return trt.PluginFieldCollection([])
+
+    def get_output_shapes(self, inputs, shape_inputs, exprBuilder):
+        output_dims = [
+            inputs[0],
+            trt.DimsExprs([inputs[0][0], exprBuilder.constant(1)]),
+        ]
+
+        return output_dims
+
+    def configure_plugin(self, inp, out):
+        pass
+
+    def on_shape_change(self, inp, out):
+        pass
+
+    def supports_format_combination(
+        self, pos: int, in_out: "list[trt.PluginTensorDesc]", num_inputs: int
+    ):
+        assert num_inputs == 3
+        assert pos < len(in_out)
+
+        desc = in_out[pos].desc
+        if desc.format != trt.TensorFormat.LINEAR:
+            return False
+
+        # self, src and output have the same type
+        if pos in [0, 1, 3]:
+            return desc.type == self.type
+
+        # indices anc the counts output are int64
+        return desc.type == trt.int64
+
+    def enqueue(self, input_desc, output_desc, inputs, outputs, workspace, stream):
+
+        # No-copy operations to setup torch tensors over the I/O buffers
+        inp_mem = UnownedMemory(
+            inputs[0], input_desc[0].dims, trt.nptype(input_desc[0].type)
+        )
+        src_mem = UnownedMemory(
+            inputs[1], input_desc[1].dims, trt.nptype(input_desc[1].type)
+        )
+        idx_mem = UnownedMemory(
+            inputs[2], input_desc[2].dims, trt.nptype(input_desc[2].type)
+        )
+        counts_mem = UnownedMemory(
+            outputs[1], output_desc[1].dims, trt.nptype(output_desc[1].type)
+        )
+
+        inp = torch.as_tensor(inp_mem.d, device="cuda")
+        src = torch.as_tensor(src_mem.d, device="cuda")
+        idx = torch.as_tensor(idx_mem.d, device="cuda")
+        counts = torch.as_tensor(counts_mem.d, device="cuda")
+
+        # Zero out the counts before passing to kernel
+        counts.zero_()
+
+        n_classes = inp.shape[0]
+        n_elements = src.numel()
+
+        # Block size definitions
+        BLOCK_SIZE = 1024
+        BLOCK_SIZE_C = 32
+
+        # Calculate grid size
+        grid_x = (n_elements + BLOCK_SIZE - 1) // BLOCK_SIZE
+        grid_y = (n_classes + BLOCK_SIZE_C - 1) // BLOCK_SIZE_C
+
+        scatter_add_kernel[(grid_x, grid_y)](
+            inp, src, idx, n_elements, n_classes, counts, BLOCK_SIZE, BLOCK_SIZE_C
+        )
+
+    def attach_to_context(self, context):
+        return self.clone()
+
+    def set_tactic(self, tactic):
+        pass
+
+    def get_aliased_input(self, output_index: int):
+        if output_index == 0:
+            return 0
+
+        return -1
+
+    def clone(self):
+        cloned_plugin = ScatterAddPlugin()
+        cloned_plugin.__dict__.update(self.__dict__)
+        return cloned_plugin
+
+
+class ScatterAddPluginCreator(trt.IPluginCreatorV3One):
+    def __init__(self):
+        trt.IPluginCreatorV3One.__init__(self)
+        self.name = "ScatterAddPlugin"
+        self.plugin_namespace = ""
+        self.plugin_version = "1"
+        self.field_names = trt.PluginFieldCollection([])
+
+    def create_plugin(self, name, fc, phase):
+        return ScatterAddPlugin()
+
+
+def torch_ref(node_features, edges, W, precision):
+    # Initialize an output tensor for aggregation
+    aggregated = torch.zeros_like(node_features, dtype=precision, device="cuda")
+
+    # Perform aggregation using scatter_add_
+    aggregated.scatter_add_(0, edges[:, 1].unsqueeze(1), node_features[edges[:, 0]])
+
+    # Get the counts of each distinct index
+    bincounts = torch.bincount(edges[:, 1].contiguous())
+
+    # Normalize and classify
+    Y = W * (aggregated / bincounts.unsqueeze(1)).transpose(1, 0)
+    return torch.softmax(torch.relu(Y), dim=0)
+
+
+numpy_to_torch_dtype = {
+    np.int32: torch.int32,
+    np.int64: torch.int64,
+    np.float16: torch.float16,
+    np.float32: torch.float32,
+}
+
+
+def parse_edges_string(input_string):
+    try:
+        # Parse the string into a list of integer pairs
+        raw_edges = ast.literal_eval(input_string)
+
+        # Check if the parsed object is a list
+        if not isinstance(raw_edges, list):
+            return None, "The input string does not represent a list."
+
+        edges = []
+        for edge in raw_edges:
+            if (
+                not isinstance(edge, list)
+                or len(edge) != 2
+                or not all(isinstance(x, int) for x in edge)
+            ):
+                return (
+                    None,
+                    f"Each edge must be a list of two integers. Invalid edge: {edge}",
+                )
+            edges.append(edge)
+
+        return edges, None
+    except (SyntaxError, ValueError) as e:
+        return None, f"Error parsing string: {e}"
+
+
+def validate_edges(edges, n_nodes):
+    for edge in edges:
+        src, target = edge
+        if not (0 <= src < n_nodes) or not (0 <= target < n_nodes):
+            return f"Edge ({src}, {target}) is out of bounds. Must be in range [0, {n_nodes - 1}]."
+
+    # check incoming edges
+    incoming_edges_count = [0] * n_nodes
+    for _, target in edges:
+        incoming_edges_count[target] += 1
+
+    for idx in range(n_nodes):
+        if incoming_edges_count[idx] == 0:
+            return f"Index {idx} has no incoming edges."
+    return None
+
+
+def parse_edges(input_string, n_nodes):
+    parsed_edges, parse_error = parse_edges_string(input_string)
+    if parse_error:
+        return None, parse_error
+    else:
+        # Validate the edges
+        validation_error = validate_edges(parsed_edges, n_nodes)
+        if validation_error is not None:
+            return None, validation_error
+        else:
+            return parsed_edges, None
+
+
+# Print adjacency matrix
+def print_graph(edges, n_nodes):
+    adjacency_matrix = [[0] * n_nodes for _ in range(n_nodes)]
+
+    for src, tgt in edges:
+        adjacency_matrix[src][tgt] = 1
+
+    for row in adjacency_matrix:
+        print(row)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--precision",
+        type=str,
+        default="fp32",
+        choices=["fp32", "fp16"],
+        help="Precision for node features",
+    )
+    parser.add_argument(
+        "--node_features",
+        type=str,
+        default="[1.0,3.0,5.0,7.0]",
+        help="List of node features as a comma-separated list. e.g. [1.0,2.0,3.0].",
+    )
+    parser.add_argument(
+        "--edges",
+        type=str,
+        default="[[0,1],[1,2],[2,3],[3,0],[0,2],[1,3]]",
+        help="Pairs of source->target directed edges. Every node must have at least one incoming edge. e.g. [[0,1],[1,0]].",
+    )
+    parser.add_argument(
+        "--num_classes", type=int, default=3, help="Number of classes in the classifier"
+    )
+    parser.add_argument(
+        "--validate", action="store_true", help="Validate result with reference"
+    )
+    parser.add_argument("--seed", type=int, help="Seed to use for weights generation")
+
+    args = parser.parse_args()
+
+    if args.seed is not None:
+        print("Setting seed to:", args.seed)
+        torch.manual_seed(args.seed)
+    else:
+        print("Setting seed to:", torch.seed())
+
+    precision = trt.float32 if args.precision == "fp32" else trt.float16
+    n_classes = args.num_classes
+
+    numpy_precision = trt.nptype(precision)
+    torch_precision = numpy_to_torch_dtype[numpy_precision]
+
+    if args.num_classes < 1:
+        parser.print_help()
+        log.error("num_classes must be a positive integer")
+        sys.exit(1)
+
+    try:
+        float_list = ast.literal_eval(args.node_features)
+        if not isinstance(float_list, list):
+            parser.print_help()
+            log.error("The node_features string does not represent a list")
+            sys.exit(1)
+
+        # Check if all elements in the list are floats/ints
+        if not all(isinstance(x, (float, int)) for x in float_list):
+            parser.print_help()
+            log.error("The node_features list must contain only numbers")
+            sys.exit(1)
+    except (SyntaxError, ValueError) as e:
+        parser.print_help()
+        log.error(f"The node_features string could not be parsed as a list: {e}")
+        sys.exit(1)
+
+    node_features = torch.tensor(float_list, dtype=torch_precision, device="cuda").view(
+        -1, 1
+    )
+
+    n_nodes = node_features.shape[0]
+
+    parsed_edges, parse_error = parse_edges(args.edges, n_nodes)
+    if parse_error:
+        parser.print_help()
+        log.error(parse_error)
+        sys.exit(1)
+
+    edges = torch.tensor(parsed_edges, device="cuda", dtype=torch.int64)
+
+    print()
+    print("Adjacency matrix for graph:")
+    print_graph(edges, n_nodes)
+    print()
+
+    target = torch.zeros_like(node_features, device="cuda")
+
+    input_x = target.clone()
+    input_src = node_features[edges[:, 0]].flatten()
+    input_idx = edges[:, 1].contiguous() + 1
+
+    W = torch.randn((n_classes, 1), dtype=torch_precision, device="cuda")
+
+    plg_registry = trt.get_plugin_registry()
+    my_plugin_creator = ScatterAddPluginCreator()
+    plg_registry.register_creator(my_plugin_creator, "")
+
+    builder, network = create_network(strongly_typed=True)
+    input_x_T = network.add_input(name="X", dtype=precision, shape=input_x.shape)
+    input_src_T = network.add_input(name="src", dtype=precision, shape=input_src.shape)
+    input_idx_T = network.add_input(name="idx", dtype=trt.int64, shape=input_idx.shape)
+    w_T = network.add_input(name="W", dtype=precision, shape=W.shape)
+    out = network.add_plugin_v3(
+        [input_x_T, input_src_T, input_idx_T], [], ScatterAddPlugin()
+    )
+    cast_layer = network.add_cast(out.get_output(1), precision)
+    div_layer = network.add_elementwise(
+        out.get_output(0),
+        cast_layer.get_output(0),
+        op=trt.ElementWiseOperation.FLOOR_DIV,
+    )
+    matmul_layer = network.add_matrix_multiply(
+        w_T,
+        trt.MatrixOperation.NONE,
+        div_layer.get_output(0),
+        trt.MatrixOperation.TRANSPOSE,
+    )
+    relu_layer = network.add_activation(
+        matmul_layer.get_output(0), type=trt.ActivationType.RELU
+    )
+    softmax_layer = network.add_softmax(relu_layer.get_output(0))
+    softmax_layer.get_output(0).name = "softmax"
+    network.mark_output(tensor=softmax_layer.get_output(0))
+    build_engine = engine_from_network(
+        (builder, network),
+        CreateConfig(
+            preview_features=[trt.PreviewFeature.ALIASED_PLUGIN_IO_10_03],
+        ),
+    )
+
+    with TrtRunner(build_engine, "trt_runner") as runner:
+        outputs = runner.infer(
+            {"X": input_x, "src": input_src, "idx": input_idx, "W": W},
+            copy_outputs_to_host=False,
+        )
+
+        print()
+        print("Classifier output:")
+        print(outputs["softmax"])
+        print()
+
+        if args.validate:
+            tref = torch_ref(node_features, edges, W, torch_precision)
+            if torch.allclose(outputs["softmax"], tref, 1e-2):
+                print("Validation against reference successful!")
+            else:
+                print("Validation against reference failed!")

samples/python/aliased_io_plugin/requirements.txt

@@ -0,0 +1,10 @@
+cupy-cuda12x
+triton==3.2.0; (platform_system != "Windows")
+torch
+--extra-index-url https://pypi.ngc.nvidia.com
+polygraphy
+colored
+numpy==1.26.4
+pyyaml==6.0.3
+requests==2.32.4
+tqdm==4.66.4

samples/python/attention_mdtrt/README.md

@@ -0,0 +1,133 @@
+# Multi-Device Attention Inference with TensorRT
+
+This sample demonstrates how to run a self-attention model across multiple GPUs using TensorRT's multi-device inference feature. It covers single-GPU execution as a baseline and multi-GPU execution using MPI for process management and NCCL for GPU-to-GPU communication.
+
+## Introduction
+
+TensorRT supports splitting a single model across multiple GPUs for inference. This is useful when a model is too large to fit on a single GPU, or when you want to reduce latency by parallelizing computation across devices.
+
+TensorRT supports multiple parallelism strategies for multi-device inference, including tensor parallelism (TP) and context parallelism (CP). This sample focuses on **context parallelism (CP)**, where the input sequence is split across GPUs along the sequence dimension. Each GPU processes its portion independently for most operations (linear projections, normalization), but attention requires cross-device communication because every token must attend to every other token. TensorRT handles this communication transparently using collective operations embedded in the model. For a detailed explanation of context parallelism, see [Context Parallelism for Scalable Million-Token Inference](https://arxiv.org/abs/2411.01783).
+
+### How Does it Work?
+
+To run a model on multiple GPUs, the model must be **sharded** (split into pieces that each GPU can execute independently). For context parallelism, sharding means dividing the input sequence across GPUs and inserting communication ops so that each GPU can still compute correct attention over the full sequence.
+
+When a model is sharded for multi-device execution, special **DistCollective** operations are inserted into the ONNX graph:
+
+- **ReduceScatter**: Splits and reduces input data across GPUs. Used at the start to distribute the input sequence.
+- **AllGather**: Collects data from all GPUs into a full tensor. Used before attention so each GPU sees the full K/V tensors, and after the output projection to reconstruct the full output.
+
+These ops are not present in the single-device ONNX model. They are inserted automatically by `polygraphy multi-device shard` using sharding hints that describe how the model should be split.
+
+### What is hint.json?
+
+The `hint.json` file tells the polygraphy sharding tool how to partition the model. The sharding tool reads the single-device ONNX graph, identifies the attention layers and I/O tensors specified in the hints, and inserts DistCollective ops at the appropriate points.
+
+```json
+{
+    "parallelism": "CP",
+    "attention_layers": [
+        {
+            "q": "q_scaled",          // Name of the Q tensor feeding into QK^T matmul
+            "gather_kv": true,        // Gather K/V across devices before attention
+            "gather_q": false         // Q stays local (each GPU has its own chunk)
+        }
+    ],
+    "dist_collectives": {
+        "nb_rank": 2,                 // Number of GPUs to shard across
+        "reduce_op": "max"            // Reduction operation for ReduceScatter
+    },
+    "inputs": [
+        {
+            "name": "input",          // Input tensor name in the ONNX model
+            "seq_len_idx": 0,         // Which dimension is the sequence length
+            "rank": 3                 // Number of dimensions
+        }
+    ],
+    "outputs": [
+        {
+            "name": "output",
+            "seq_len_idx": 0,
+            "rank": 3
+        }
+    ]
+}
+```
+
+For full documentation on the sharding tool and hint format, see the [Polygraphy multi-device documentation](https://github.com/NVIDIA/TensorRT/tree/main/tools/Polygraphy/polygraphy/tools/multi_device).
+
+## Prerequisites
+
+- A machine with multi-GPU
+- `polygraphy` >= 0.49.25 (for `multi-device shard` support)
+
+## Install Dependencies
+
+```bash
+pip3 install -r requirements.txt
+```
+
+## Generating the ONNX Models
+
+### Step 1: Generate the single-device model
+
+```bash
+python3 create_onnx.py --output attention_sd.onnx
+```
+
+This creates a self-attention model with:
+- Input/Output: `(sequence_length, batch_size, 4096)` in float16
+- 32 attention heads, 128-dim per head
+- Q/K/V projections, RMSNorm, scaled dot-product attention, output projection
+
+### Step 2: Shard for multi-device
+
+```bash
+polygraphy multi-device shard attention_sd.onnx -s hint.json -o attention_md.onnx
+```
+
+This inserts DistCollective ops (ReduceScatter, AllGather) into the model based on `hint.json`. The resulting `attention_md.onnx` is designed to run on 2 GPUs.
+
+## Running the Sample
+
+### Single-GPU
+
+```bash
+python3 attention_mdtrt.py \
+  --onnx-path attention_sd.onnx \
+  --sequence-length 56320 \
+  --batch-size 1 \
+  --num-iterations 50
+```
+
+### Multi-GPU (2 GPUs)
+
+```bash
+mpirun -np 2 python3 attention_mdtrt.py \
+  --onnx-path attention_md.onnx \
+  --sequence-length 56320 \
+  --batch-size 1 \
+  --num-iterations 50
+```
+
+### Using a specific libnccl.so
+
+```bash
+LD_PRELOAD=/path/to/libnccl.so mpirun -np 2 python3 attention_mdtrt.py \
+  --onnx-path attention_md.onnx
+```
+
+## License
+
+For terms and conditions for use, reproduction, and distribution, see the [TensorRT Software License Agreement](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sla/index.html) documentation.
+
+## Changelog
+
+April 2026
+Added `create_onnx.py` for ONNX model generation using the GraphSurgeon layer API. Added `--save-output` flag for saving inference output. Updated documentation with DistCollective and sharding explanations.
+
+January 2026
+Initial release of this sample.
+
+## Known Issues
+None

samples/python/attention_mdtrt/attention_mdtrt.py

@@ -0,0 +1,491 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 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 tensorrt as trt
+import torch
+import sys
+import os
+import time
+import argparse
+from cuda.bindings import runtime as cudart
+from ctypes import py_object, pythonapi, c_void_p, c_char_p
+from typing import Optional
+
+try:
+    from mpi4py import MPI
+except ImportError:
+    MPI = None
+
+try:
+    import nccl.core as nccl
+except ImportError:
+    nccl = None
+
+sys.path.insert(1, os.path.join(os.path.dirname(os.path.realpath(__file__)), os.pardir))
+import common
+
+
+def communicator_to_capsule(comm):
+    """
+    Convert nccl.core.Communicator to PyCapsule for TensorRT compatibility.
+    
+    Args:
+        comm: nccl.core.Communicator instance with .ptr attribute set to ncclComm_t handle
+        
+    Returns:
+        PyCapsule wrapping the communicator pointer, suitable for set_communicator()
+        
+    Raises:
+        ValueError: If comm.ptr is invalid (0 or None), indicating destroyed communicator
+        TypeError: If comm doesn't have a .ptr attribute
+    """
+    # Validate input
+    if comm is None:
+        raise TypeError("Communicator cannot be None")
+    
+    if not hasattr(comm, 'ptr'):
+        raise TypeError(f"Object {type(comm)} does not have 'ptr' attribute. "
+                       "Expected nccl.core.Communicator instance.")
+    
+    # Get the raw pointer from the Communicator object
+    ptr = comm.ptr
+    
+    # Validate that communicator is still alive (ptr != 0)
+    if ptr == 0:
+        raise ValueError("NCCL Communicator has been destroyed (ptr=0). "
+                        "Cannot create capsule for destroyed communicator.")
+    
+    # Convert to PyCapsule using ctypes.pythonapi
+    PyCapsule_New = pythonapi.PyCapsule_New
+    PyCapsule_New.restype = py_object
+    PyCapsule_New.argtypes = [c_void_p, c_char_p, c_void_p]
+    
+    capsule = PyCapsule_New(c_void_p(ptr), b"ncclComm_t", None)
+    
+    if capsule is None:
+        raise RuntimeError("Failed to create PyCapsule from communicator pointer")
+    
+    return capsule
+
+
+def allocate_buffers(engine: trt.ICudaEngine, profile_idx: Optional[int] = None, output_shape: Optional[tuple] = None):
+    """Allocate host and device buffers for TensorRT engine."""
+    inputs = []
+    outputs = []
+    bindings = []
+    tensor_names = [engine.get_tensor_name(i) for i in range(engine.num_io_tensors)]
+    for binding in tensor_names:
+        # Pick out the max shape to allocate enough memory for the binding.
+        shape = engine.get_tensor_shape(binding) if profile_idx is None else engine.get_tensor_profile_shape(binding, profile_idx)[-1]
+        shape_valid = np.all([s >= 0 for s in shape])
+        if not shape_valid and profile_idx is None:
+            raise ValueError(f"Binding {binding} has dynamic shape, " +\
+                "but no profile was specified.")
+        # For dynamic shapes, use fixed output shape
+        if output_shape is not None:
+            shape = output_shape
+        size = trt.volume(shape)
+        trt_type = engine.get_tensor_dtype(binding)
+
+        # Allocate host and device buffers
+        if trt_type == trt.DataType.BF16:
+            dtype = np.dtype(np.uint16)
+            bindingMemory = common.HostDeviceMem(size, dtype)
+        elif trt_type == trt.DataType.HALF:
+            dtype = np.dtype(np.uint16)
+            bindingMemory = common.HostDeviceMem(size, dtype)
+        elif trt_type == trt.DataType.FLOAT:
+            dtype = np.dtype(np.float32)
+            bindingMemory = common.HostDeviceMem(size, dtype)
+        else:
+            try:
+                dtype = np.dtype(trt.nptype(trt_type))
+                bindingMemory = common.HostDeviceMem(size, dtype)
+            except TypeError:
+                size = int(size * trt_type.itemsize)
+                bindingMemory = common.HostDeviceMem(size)
+
+        # Append the device buffer to device bindings.
+        bindings.append(int(bindingMemory.device_ptr))
+
+        # Append to the appropriate list.
+        if engine.get_tensor_mode(binding) == trt.TensorIOMode.INPUT:
+            inputs.append(bindingMemory)
+        else:
+            outputs.append(bindingMemory)
+    return inputs, outputs, bindings
+
+
+class AttentionSD:
+    """Base class for Attention model using TensorRT (Single Device)"""
+    
+    def __init__(self, mpi_comm, rank, onnx_path):
+        """
+        Initialize the Attention class
+        
+        Args:
+            mpi_comm: MPI communicator
+            rank: Current instance ID
+            onnx_path: Path to the ONNX model
+        """
+        self.onnx_path = onnx_path
+        self.logger = trt.Logger(trt.Logger.WARNING)
+        self.engine = None
+        self.context = None
+        self.inputs = None
+        self.outputs = None
+        self.bindings = None
+        self.mpi_comm = mpi_comm
+        self.rank = rank
+
+    def setup(self, actual_input_shape, output_shape):
+        """
+        Set up everything before doing inference.
+        """
+        engine_string = self.build_serialized_network()
+        self.engine = trt.Runtime(self.logger).deserialize_cuda_engine(engine_string)
+        if self.engine is None:
+            print("Failed deserializing engine!")
+            exit(-1)
+        print("Succeeded deserializing engine!")
+
+        self.context = self.engine.create_execution_context()
+        
+        # For dynamic shapes, we need to specify the actual input shape we want to use
+        input_name = self.engine.get_tensor_name(0)
+        self.context.set_input_shape(input_name, actual_input_shape)
+        
+        # Allocate buffers
+        self.inputs, self.outputs, self.bindings = allocate_buffers(
+            self.engine, profile_idx=0, output_shape=output_shape
+        )
+        
+        num_io = self.engine.num_io_tensors
+        tensor_names = [self.engine.get_tensor_name(i) for i in range(num_io)]
+        for i in range(num_io):
+            self.context.set_tensor_address(tensor_names[i], self.bindings[i])
+
+    def build_serialized_network(self):
+        """Create and serialize a network from the ONNX model."""
+        # Create builder and empty network
+        builder = trt.Builder(self.logger)
+        network = builder.create_network(flags=1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED))
+        
+        # Setup parser and parse the ONNX model
+        print(f"Parsing ONNX model from {self.onnx_path}")
+        parser = trt.OnnxParser(network, self.logger)
+
+        with open(self.onnx_path, "rb") as f:
+            if not parser.parse(f.read()):
+                print("Failed to parse ONNX model")
+                for error in range(parser.num_errors):
+                    print(parser.get_error(error))
+                return None
+        
+        # Get input dimensions and data type
+        input_tensor = network.get_input(0)
+        input_shape = input_tensor.shape
+        input_name = input_tensor.name
+        input_dtype = input_tensor.dtype
+        
+        print(f"[Rank {self.rank}] Input shape: {input_shape}")
+        print(f"[Rank {self.rank}] Input name: {input_name}")
+        print(f"[Rank {self.rank}] Input data type: {input_dtype}")
+
+        # Create a builder config
+        config = builder.create_builder_config()
+        
+        config.profiling_verbosity = trt.ProfilingVerbosity.DETAILED
+        
+        config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 16 * 1024 * 1024 * 1024)  # 16GB workspace
+        config.set_memory_pool_limit(trt.MemoryPoolType.TACTIC_SHARED_MEMORY, 1 * 1024 * 1024 * 1024)  # 1GB shared mem
+        
+        profile = builder.create_optimization_profile()
+        
+        # Set the shape range for the input tensor
+        min_shape = (1, 1, 4096) 
+        opt_shape = (56320, 1, 4096)
+        max_shape = (56320, 1, 4096)
+        
+        profile.set_shape(input_name, min_shape, opt_shape, max_shape)
+        config.add_optimization_profile(profile)
+        
+        # Build the serialized network
+        serialized_engine = builder.build_serialized_network(network, config)
+        if serialized_engine is None:
+            print(f"[Rank {self.rank}] Failed building serialized engine!")
+            exit(-1)
+        print(f"[Rank {self.rank}] Succeeded building serialized engine!")
+        
+        return serialized_engine
+    
+    def infer(self, input_data, output_shape, num_iterations):
+        """
+        Execute inference on the input data.
+        
+        Args:
+            input_data: Input data for inference
+            output_shape: Expected output shape for reshaping
+            num_iterations: Number of inference iterations for averaging timing results
+            
+        Returns:
+            output_data: List of output tensors
+        """
+        print(f"[Rank {self.rank}] Input shape: {input_data.shape}")
+        
+        # Copy input data to device
+        for input_buffer in self.inputs:
+            common.memcpy_host_to_device(input_buffer.device_ptr, input_data)
+        
+        # Warmup
+        with common.CudaStreamContext() as stream:
+            self.context.execute_async_v3(stream.stream)
+            common.cuda_call(cudart.cudaStreamSynchronize(stream.stream))
+        
+            # Run inference
+            start = time.time()
+            for _ in range(num_iterations):
+                self.context.execute_async_v3(stream.stream)
+                common.cuda_call(cudart.cudaStreamSynchronize(stream.stream))
+            end = time.time()
+        
+            print(f"[Rank {self.rank}] Time spent in TRT attention: {(end-start)/num_iterations * 1000} ms")
+        
+        # Get output
+        output_data = []
+        for output in self.outputs:
+            common.memcpy_device_to_host(output.host, output.device_ptr)
+
+            # Process based on data type
+            if self.engine.get_tensor_dtype(self.engine.get_tensor_name(1)) == trt.DataType.BF16:
+                numpy_output = np.frombuffer(output.host, dtype=np.uint16).reshape(output_shape)
+                torch_output = torch.from_numpy(numpy_output).view(torch.bfloat16)
+                torch_output = torch_output.reshape(output_shape)
+            elif self.engine.get_tensor_dtype(self.engine.get_tensor_name(1)) == trt.DataType.HALF:
+                numpy_output = np.frombuffer(output.host, dtype=np.float16).reshape(output_shape)
+                torch_output = torch.from_numpy(numpy_output)
+            else:
+                numpy_output = np.frombuffer(output.host, dtype=np.float32).reshape(output_shape)
+                torch_output = torch.from_numpy(numpy_output)
+            
+            output_data.append(torch_output)
+                
+        return output_data
+    
+    def cleanup(self):
+        """
+        Free the buffer resources.
+        """
+        common.free_buffers(self.inputs, self.outputs)
+
+
+class AttentionMD(AttentionSD):
+    """Multi-device Attention model using TensorRT with NCCL for communication"""
+    
+    def __init__(self, mpi_comm, num_ranks, rank, onnx_path):
+        """
+        Initialize the multi-device Attention class
+        
+        Args:
+            mpi_comm: MPI communicator
+            num_ranks: Number of instances/devices
+            rank: Current instance ID
+            onnx_path: Path to the ONNX model
+        """
+        super(AttentionMD, self).__init__(mpi_comm, rank, onnx_path)
+        self.num_ranks = num_ranks
+        self.nccl_comm = None
+
+    def setup_multidevice(self, root):
+        """
+        Set up CUDA devices and initialize NCCL communicator.
+        
+        Args:
+            root: Root rank for communication
+        """
+        assert nccl is not None
+        assert root <= self.num_ranks - 1
+        assert self.rank <= self.num_ranks - 1
+
+        num_devices = common.cuda_call(cudart.cudaGetDeviceCount())
+        assert num_devices >= self.num_ranks
+
+        common.cuda_call(cudart.cudaSetDevice(self.rank))
+
+        if self.rank == root:
+            nccl_comm_id = nccl.get_unique_id()
+        else:
+            nccl_comm_id = None
+
+        nccl_comm_id = self.mpi_comm.bcast(nccl_comm_id, root=root)
+        self.nccl_comm = nccl.Communicator.init(nranks=self.num_ranks, rank=self.rank, unique_id=nccl_comm_id)
+
+    def setup(self, actual_input_shape, output_shape, root=0):
+        """
+        Set up the multi-device environment and build/load the engine
+        
+        Args:
+            root: Root rank for communication
+        """
+        self.setup_multidevice(root)
+        
+        # Load or build TRT engine
+        if self.rank == root:
+            engine_bin = bytes(self.build_serialized_network())
+        else:
+            engine_bin = None
+        
+        # Broadcast the serialized engine from root to all ranks
+        engine_bin = self.mpi_comm.bcast(engine_bin, root=root)
+        
+        # Deserialize the engine
+        self.engine = trt.Runtime(self.logger).deserialize_cuda_engine(engine_bin)
+        if self.engine is None:
+            print(f"[Rank {self.rank}] Failed deserializing engine!")
+            exit(-1)
+        print(f"[Rank {self.rank}] Succeeded deserializing engine!")
+        
+        # Create an execution context
+        self.context = self.engine.create_execution_context()
+        
+        # Set the NCCL communicator for multi-device communication
+        capsule = communicator_to_capsule(self.nccl_comm)
+        if not self.context.set_communicator(capsule):
+            print(f"[Rank {self.rank}] Failed to set communicator")
+            exit(-1)
+
+        # For dynamic shapes, we need to specify the actual input shape we want to use
+        input_name = self.engine.get_tensor_name(0)
+        self.context.set_input_shape(input_name, actual_input_shape)
+        
+        # Allocate buffers for local portion of data
+        self.inputs, self.outputs, self.bindings = allocate_buffers(
+            self.engine, profile_idx=0, output_shape=output_shape
+        )
+        
+        num_io = self.engine.num_io_tensors
+        tensor_names = [self.engine.get_tensor_name(i) for i in range(num_io)]
+        for i in range(num_io):
+            self.context.set_tensor_address(tensor_names[i], self.bindings[i])
+
+
+def generate_random_input(sequence_length, batch_size):
+    """Generate random float16 input data with the given shape."""
+    torch.manual_seed(42)
+    torch_input = torch.rand((sequence_length, batch_size, 4096)).to(torch.float16)
+    input_data = np.ascontiguousarray(torch_input.cpu().numpy())
+    return input_data, (sequence_length, batch_size, 4096)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Sample script for Attention MDTRT")
+    parser.add_argument("--onnx-path", type=str, required=True, help="Path to ONNX model")
+    parser.add_argument("--sequence-length", type=int, default=56320, help="Sequence length for input")
+    parser.add_argument("--batch-size", type=int, default=1, help="Batch size")
+    parser.add_argument("--num-iterations", type=int, default=50, help="Number of inference iterations for timing")
+    parser.add_argument("--save-output", type=str, default=None, help="Save output tensor to .npy file (root rank only)")
+
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    
+    # Initialize MPI if available
+    if MPI is not None:
+        mpi_comm = MPI.COMM_WORLD
+        num_ranks = mpi_comm.Get_size()
+        rank = mpi_comm.Get_rank()
+        root = 0
+    else:
+        # Fallback for single-process execution
+        mpi_comm = None
+        num_ranks = 1
+        rank = 0
+        root = 0
+    
+    actual_input_shape = (args.sequence_length, args.batch_size, 4096)
+    output_shape = (args.sequence_length, args.batch_size, 4096)
+    
+    # Print configuration
+    if rank == root:
+        print(f"[setup] Configuration:")
+        print(f"[setup]   Number of GPUs: {num_ranks}")
+        print(f"[setup]   Sequence Length: {args.sequence_length}")
+        print(f"[setup]   Batch Size: {args.batch_size}")
+        print(f"[setup]   Data Type: float16")
+        print(f"[setup]   Input Shape: {actual_input_shape}")
+        print(f"[setup]   Output Shape: {output_shape}")
+    
+    # Generate random input data with FULL sequence length (only on root rank)
+    if rank == root:
+        input_data, input_shape = generate_random_input(args.sequence_length, args.batch_size)
+        print(f"[Rank {rank}] Generated random input data with shape: {input_shape}")
+        
+        if num_ranks == 1:
+            print(f"[Rank {rank}] Running single-device inference...")
+            try:
+                attention_sd = AttentionSD(mpi_comm, rank, args.onnx_path)
+                attention_sd.setup(actual_input_shape, output_shape)
+                sd_output = attention_sd.infer(input_data, output_shape, args.num_iterations)[0]
+                print(f"[Rank {rank}] Single-device inference completed")
+                print(f"[Rank {rank}] Output shape: {sd_output.shape}")
+                if args.save_output:
+                    np.save(args.save_output, sd_output.float().cpu().numpy())
+                    print(f"[Rank {rank}] Output saved to {args.save_output}")
+                attention_sd.cleanup()
+            except Exception as e:
+                print(f"[Rank {rank}] Error in single-device inference: {e}")
+                sys.exit(1)
+    else:
+        input_data = None
+    
+    # Broadcast full input data to all ranks for multi-device inference
+    if MPI is not None and num_ranks > 1:
+        input_data = mpi_comm.bcast(input_data, root=root)
+    
+    # Run multi-device inference if num_gpus > 1
+    if num_ranks > 1:
+        if MPI is None:
+            print(f"Error: MPI is required for multi-GPU tests but not available. Ensure you run with mpirun.")
+            sys.exit(1)
+        
+        if nccl is None:
+            print(f"Error: nccl is required for multi-GPU tests but not available.")
+            sys.exit(1)
+        
+        print(f"[Rank {rank}] Running multi-device inference...")
+        try:
+            attention_md = AttentionMD(mpi_comm, num_ranks, rank, args.onnx_path)
+            attention_md.setup(actual_input_shape, output_shape, root)
+            md_output = attention_md.infer(input_data, output_shape, args.num_iterations)[0]
+            print(f"[Rank {rank}] Multi-device inference completed")
+            print(f"[Rank {rank}] Output shape: {md_output.shape}")
+            if rank == root and args.save_output:
+                np.save(args.save_output, md_output.float().cpu().numpy())
+                print(f"[Rank {rank}] Output saved to {args.save_output}")
+            attention_md.cleanup()
+        except Exception as e:
+            print(f"[Rank {rank}] Error in multi-device inference: {e}")
+            sys.exit(1)
+    
+    print(f"[Rank {rank}] Test completed successfully!")
+
+
+if __name__ == "__main__":
+    main()

samples/python/attention_mdtrt/create_onnx.py

@@ -0,0 +1,258 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 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.
+#
+
+"""
+Construct a self-attention ONNX model using the ONNX GraphSurgeon layer API.
+
+The model implements:
+  1. Q/K/V linear projections (MatMul with 4096x4096 weights)
+  2. Reshape to multi-head layout (seq, batch, 4096) -> (seq, batch, 32, 128)
+  3. RMSNorm on Q and K
+  4. Scaled dot-product attention (QK^T / sqrt(d), softmax, attn * V)
+  5. Reshape back and output projection
+
+Input/Output: (sequence_length, batch_size, 4096), float16
+"""
+
+import argparse
+import math
+
+import numpy as np
+import onnx
+import onnx_graphsurgeon as gs
+
+NUM_HEADS = 32
+HEAD_DIM = 128
+HIDDEN_DIM = NUM_HEADS * HEAD_DIM  # 4096
+OPSET = 17
+
+
+# Register ONNX ops as methods on gs.Graph using the layer API.
+# Each returns the output tensor(s) directly for easy chaining.
+
+@gs.Graph.register()
+def matmul(self, a, b):
+    return self.layer(op="MatMul", inputs=[a, b], outputs=["matmul_out"])[0]
+
+
+@gs.Graph.register()
+def transpose(self, a, perm):
+    return self.layer(op="Transpose", inputs=[a], attrs={"perm": perm}, outputs=["transpose_out"])[0]
+
+
+@gs.Graph.register()
+def reshape(self, data, shape):
+    return self.layer(op="Reshape", inputs=[data, shape], attrs={"allowzero": 0}, outputs=["reshape_out"])[0]
+
+
+@gs.Graph.register()
+def softmax(self, a, axis=-1):
+    return self.layer(op="Softmax", inputs=[a], attrs={"axis": axis}, outputs=["softmax_out"])[0]
+
+
+@gs.Graph.register()
+def cast(self, a, to):
+    return self.layer(op="Cast", inputs=[a], attrs={"to": to}, outputs=["cast_out"])[0]
+
+
+@gs.Graph.register()
+def sqrt(self, a):
+    return self.layer(op="Sqrt", inputs=[a], outputs=["sqrt_out"])[0]
+
+
+@gs.Graph.register()
+def add(self, a, b):
+    return self.layer(op="Add", inputs=[a, b], outputs=["add_out"])[0]
+
+
+@gs.Graph.register()
+def mul(self, a, b):
+    return self.layer(op="Mul", inputs=[a, b], outputs=["mul_out"])[0]
+
+
+@gs.Graph.register()
+def div(self, a, b):
+    return self.layer(op="Div", inputs=[a, b], outputs=["div_out"])[0]
+
+
+@gs.Graph.register()
+def pow(self, a, b):
+    return self.layer(op="Pow", inputs=[a, b], outputs=["pow_out"])[0]
+
+
+@gs.Graph.register()
+def reduce_mean(self, a, axes, keepdims=1):
+    return self.layer(
+        op="ReduceMean", inputs=[a], attrs={"axes": axes, "keepdims": keepdims},
+        outputs=["reduce_mean_out"],
+    )[0]
+
+
+@gs.Graph.register()
+def shape_op(self, a):
+    return self.layer(op="Shape", inputs=[a], outputs=["shape_out"])[0]
+
+
+@gs.Graph.register()
+def gather(self, data, indices):
+    return self.layer(op="Gather", inputs=[data, indices], attrs={"axis": 0}, outputs=["gather_out"])[0]
+
+
+@gs.Graph.register()
+def unsqueeze(self, a, axes):
+    return self.layer(op="Unsqueeze", inputs=[a, axes], outputs=["unsqueeze_out"])[0]
+
+
+@gs.Graph.register()
+def concat(self, inputs, axis=0):
+    return self.layer(op="Concat", inputs=inputs, attrs={"axis": axis}, outputs=["concat_out"])[0]
+
+
+def build_attention_graph():
+    """Build the full self-attention ONNX graph."""
+    rng = np.random.default_rng(42)
+    graph = gs.Graph(opset=OPSET)
+
+    def fp16_weights(shape):
+        return rng.standard_normal(shape).astype(np.float16)
+
+    def fp32_scalar(val):
+        return np.array([val], dtype=np.float32)
+
+    axes_0 = np.array([0], dtype=np.int64)
+
+    # Input: (seq, batch, 4096) fp16
+    graph_input = gs.Variable("input", dtype=np.float16, shape=["sequence_length", "batch_size", HIDDEN_DIM])
+    graph.inputs = [graph_input]
+
+    # Q/K/V projections
+    q_proj = graph.matmul(graph_input, fp16_weights((HIDDEN_DIM, HIDDEN_DIM)))
+    k_proj = graph.matmul(graph_input, fp16_weights((HIDDEN_DIM, HIDDEN_DIM)))
+    v_proj = graph.matmul(graph_input, fp16_weights((HIDDEN_DIM, HIDDEN_DIM)))
+
+    # Dynamic reshape: (seq, batch, 4096) -> (seq, batch, 32, 128)
+    # Build target shape [seq_dim, batch_dim, 32, 128] from input shape
+    def reshape_to_heads(proj):
+        inp_shape = graph.shape_op(proj)
+        seq_dim = graph.unsqueeze(graph.gather(inp_shape, np.array(0, dtype=np.int64)), axes_0)
+        batch_dim = graph.unsqueeze(graph.gather(inp_shape, np.array(1, dtype=np.int64)), axes_0)
+        target_shape = graph.concat([
+            seq_dim, batch_dim,
+            np.array([NUM_HEADS], dtype=np.int64),
+            np.array([HEAD_DIM], dtype=np.int64),
+        ])
+        return graph.reshape(proj, target_shape)
+
+    q_4d = reshape_to_heads(q_proj)
+    k_4d = reshape_to_heads(k_proj)
+    v_4d = reshape_to_heads(v_proj)
+
+    # RMSNorm: x * rsqrt(mean(x^2) + eps) * weight
+    def rmsnorm(x):
+        x_fp32 = graph.cast(x, onnx.TensorProto.FLOAT)
+        sq = graph.pow(x_fp32, fp32_scalar(2.0))
+        mean = graph.reduce_mean(sq, axes=[-1])
+        rms = graph.sqrt(graph.add(mean, fp32_scalar(1e-6)))
+        inv_rms = graph.div(fp32_scalar(1.0), rms)
+        normed = graph.mul(x_fp32, inv_rms)
+        normed_fp16 = graph.cast(normed, onnx.TensorProto.FLOAT16)
+        weight = rng.standard_normal((1, 1, 1, HEAD_DIM)).astype(np.float16)
+        return graph.mul(weight, normed_fp16)
+
+    q_norm = rmsnorm(q_4d)
+    k_norm = rmsnorm(k_4d)
+
+    # Transpose to attention layout: (seq, batch, heads, hdim) -> (batch, heads, seq, hdim)
+    q_attn = graph.transpose(q_norm, perm=[1, 2, 0, 3])
+    k_attn = graph.transpose(k_norm, perm=[1, 2, 0, 3])
+    v_attn = graph.transpose(v_4d, perm=[1, 2, 0, 3])
+
+    # Dynamic reshape Q/K/V to (batch, heads, -1, hdim) using shape extraction
+    def reshape_attn(x):
+        s = graph.shape_op(x)
+        batch = graph.unsqueeze(graph.gather(s, np.array(0, dtype=np.int64)), axes_0)
+        heads = graph.unsqueeze(graph.gather(s, np.array(1, dtype=np.int64)), axes_0)
+        hdim = graph.unsqueeze(graph.gather(s, np.array(3, dtype=np.int64)), axes_0)
+        target = graph.concat([batch, heads, np.array([-1], dtype=np.int64), hdim])
+        return graph.reshape(x, target)
+
+    q_r = reshape_attn(q_attn)
+    k_r = reshape_attn(k_attn)
+    v_r = reshape_attn(v_attn)
+
+    # Scale: split sqrt(1/sqrt(head_dim)) across Q and K
+    scale_val = math.sqrt(math.sqrt(1.0 / HEAD_DIM))
+    scale_fp16 = np.array([scale_val], dtype=np.float16)
+    q_scaled = graph.mul(q_r, scale_fp16)
+    q_scaled.name = "q_scaled"
+
+    # Transpose K: (batch, heads, seq, hdim) -> (batch, heads, hdim, seq)
+    k_t = graph.transpose(k_r, perm=[0, 1, 3, 2])
+    k_scaled = graph.mul(k_t, scale_fp16)
+
+    # QK^T -> Softmax -> Attn*V
+    qk = graph.matmul(q_scaled, k_scaled)
+    attn_weights = graph.softmax(qk, axis=-1)
+    attn_out = graph.matmul(attn_weights, v_r)
+
+    # Reshape back: (batch, heads, seq, hdim) -> (seq, batch, 4096)
+    attn_t = graph.transpose(attn_out, perm=[2, 0, 1, 3])
+    attn_shape = graph.shape_op(attn_t)
+    seq_dim = graph.unsqueeze(graph.gather(attn_shape, np.array(0, dtype=np.int64)), axes_0)
+    batch_dim = graph.unsqueeze(graph.gather(attn_shape, np.array(1, dtype=np.int64)), axes_0)
+    # Compute heads * hdim
+    heads_dim = graph.gather(attn_shape, np.array(2, dtype=np.int64))
+    hdim_dim = graph.gather(attn_shape, np.array(3, dtype=np.int64))
+    hidden = graph.unsqueeze(graph.mul(heads_dim, hdim_dim), axes_0)
+    flat_shape = graph.concat([seq_dim, batch_dim, hidden])
+    attn_flat = graph.reshape(attn_t, flat_shape)
+
+    # Output projection
+    output = graph.matmul(attn_flat, fp16_weights((HIDDEN_DIM, HIDDEN_DIM)))
+    output.name = "output"
+    output.dtype = np.float16
+    output.shape = ["sequence_length", "batch_size", HIDDEN_DIM]
+    graph.outputs = [output]
+
+    graph.cleanup().toposort()
+    model = gs.export_onnx(graph)
+    model.ir_version = 8
+
+    return model
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Generate attention_sd.onnx model using the ONNX GraphSurgeon layer API"
+    )
+    parser.add_argument(
+        "--output", type=str, default="attention_sd.onnx",
+        help="Output ONNX file path (default: attention_sd.onnx)",
+    )
+    args = parser.parse_args()
+
+    model = build_attention_graph()
+    onnx.save(model, args.output)
+
+    print(f"Saved model to {args.output}")
+    print(f"  Nodes: {len(model.graph.node)}")
+    print(f"  Initializers: {len(model.graph.initializer)}")
+    print(f"  Opset: {model.opset_import[0].version}")
+
+
+if __name__ == "__main__":
+    main()

samples/python/attention_mdtrt/hint.json

@@ -0,0 +1,40 @@
+{
+    "parallelism": "CP",
+    "attention_layers": [
+        {
+            "q": "q_scaled",
+            "gather_kv": true,
+            "gather_q": false,
+            "replace": null,
+            "polygraphy_class": "AttentionLayerHint"
+        }
+    ],
+    "dist_collectives": {
+        "group_size": 0,
+        "root": -1,
+        "nb_rank": 2,
+        "reduce_op": "max",
+        "groups": [],
+        "polygraphy_class": "DistCollective"
+    },
+    "inputs": [
+        {
+            "name": "input",
+            "seq_len_idx": 0,
+            "rank": 3,
+            "polygraphy_class": "ShardTensor"
+        }
+    ],
+    "outputs": [
+        {
+            "name": "output",
+            "seq_len_idx": 0,
+            "rank": 3,
+            "polygraphy_class": "ShardTensor"
+        }
+    ],
+    "k_seq_len_idx": 3,
+    "v_seq_len_idx": 2,
+    "kv_rank": 4,
+    "polygraphy_class": "ShardHints"
+}

samples/python/attention_mdtrt/requirements.txt

@@ -0,0 +1,7 @@
+numpy==1.26.4
+cuda-bindings
+mpi4py==4.1.1
+nccl4py==0.1.1
+onnx_graphsurgeon
+polygraphy>=0.49.25
+torch

samples/python/common.py

@@ -0,0 +1,142 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2025 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 argparse
+import os
+
+import tensorrt as trt
+from common_runtime import *
+
+# FileNotFoundError is available in Python 3.3+
+
+
+def GiB(val):
+    return val * 1 << 30
+
+
+def add_help(description):
+    parser = argparse.ArgumentParser(
+        description=description, formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+    args, _ = parser.parse_known_args()
+
+
+def find_sample_data(
+    description="Runs a TensorRT Python sample", subfolder="", find_files=[], err_msg=""
+):
+    """
+    Parses sample arguments.
+
+    Args:
+        description (str): Description of the sample.
+        subfolder (str): The subfolder containing data relevant to this sample
+        find_files (str): A list of filenames to find. Each filename will be replaced with an absolute path.
+
+    Returns:
+        str: Path of data directory.
+    """
+
+    # Standard command-line arguments for all samples.
+    kDEFAULT_DATA_ROOT = os.path.join(os.sep, "usr", "src", "tensorrt", "data")
+    parser = argparse.ArgumentParser(
+        description=description, formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+    parser.add_argument(
+        "-d",
+        "--datadir",
+        help="Location of the TensorRT sample data directory, and any additional data directories.",
+        action="append",
+        default=[kDEFAULT_DATA_ROOT],
+    )
+    args, _ = parser.parse_known_args()
+
+    def get_data_path(data_dir):
+        # If the subfolder exists, append it to the path, otherwise use the provided path as-is.
+        data_path = os.path.join(data_dir, subfolder)
+        if not os.path.exists(data_path):
+            if data_dir != kDEFAULT_DATA_ROOT:
+                print(
+                    "WARNING: "
+                    + data_path
+                    + " does not exist. Trying "
+                    + data_dir
+                    + " instead."
+                )
+            data_path = data_dir
+        # Make sure data directory exists.
+        if not (os.path.exists(data_path)) and data_dir != kDEFAULT_DATA_ROOT:
+            print(
+                "WARNING: {:} does not exist. Please provide the correct data path with the -d option.".format(
+                    data_path
+                )
+            )
+        return data_path
+
+    data_paths = [get_data_path(data_dir) for data_dir in args.datadir]
+    return data_paths, locate_files(data_paths, find_files, err_msg)
+
+
+def locate_files(data_paths, filenames, err_msg=""):
+    """
+    Locates the specified files in the specified data directories.
+    If a file exists in multiple data directories, the first directory is used.
+
+    Args:
+        data_paths (List[str]): The data directories.
+        filename (List[str]): The names of the files to find.
+
+    Returns:
+        List[str]: The absolute paths of the files.
+
+    Raises:
+        FileNotFoundError if a file could not be located.
+    """
+    found_files = [None] * len(filenames)
+    for data_path in data_paths:
+        # Find all requested files.
+        for index, (found, filename) in enumerate(zip(found_files, filenames)):
+            if not found:
+                file_path = os.path.abspath(os.path.join(data_path, filename))
+                if os.path.exists(file_path):
+                    found_files[index] = file_path
+
+    # Check that all files were found
+    for f, filename in zip(found_files, filenames):
+        if not f or not os.path.exists(f):
+            raise FileNotFoundError(
+                "Could not find {:}. Searched in data paths: {:}\n{:}".format(
+                    filename, data_paths, err_msg
+                )
+            )
+    return found_files
+
+
+# Sets up the builder to use the timing cache file, and creates it if it does not already exist
+def setup_timing_cache(config: trt.IBuilderConfig, timing_cache_path: os.PathLike):
+    buffer = b""
+    if os.path.exists(timing_cache_path):
+        with open(timing_cache_path, mode="rb") as timing_cache_file:
+            buffer = timing_cache_file.read()
+    timing_cache: trt.ITimingCache = config.create_timing_cache(buffer)
+    config.set_timing_cache(timing_cache, True)
+
+
+# Saves the config's timing cache to file
+def save_timing_cache(config: trt.IBuilderConfig, timing_cache_path: os.PathLike):
+    timing_cache: trt.ITimingCache = config.get_timing_cache()
+    with open(timing_cache_path, "wb") as timing_cache_file:
+        timing_cache_file.write(memoryview(timing_cache.serialize()))

samples/python/common_runtime.py

@@ -0,0 +1,474 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2026 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 ctypes
+from typing import Optional, List, Union
+
+import numpy as np
+import tensorrt as trt
+
+from cuda.bindings import driver as cuda, runtime as cudart, nvrtc
+
+
+class ArrayWithOwner(np.ndarray):
+    """Numpy array that holds a reference to its owner object"""
+    def __new__(cls, input_array, owner):
+        obj = np.asarray(input_array).view(cls)
+        obj._owner = owner
+        return obj
+
+    def __array_finalize__(self, obj):
+        if obj is None:
+            return
+        self._owner = getattr(obj, '_owner', None)
+
+
+
+def cuda_call(call):
+    """Helper function to make CUDA calls and check for errors"""
+    def _cudaGetErrorEnum(error):
+        if isinstance(error, cuda.CUresult):
+            err, name = cuda.cuGetErrorName(error)
+            return name if err == cuda.CUresult.CUDA_SUCCESS else "<unknown>"
+        elif isinstance(error, cudart.cudaError_t):
+            return cudart.cudaGetErrorName(error)[1]
+        elif isinstance(error, nvrtc.nvrtcResult):
+            return nvrtc.nvrtcGetErrorString(error)[1]
+        else:
+            raise RuntimeError("Unknown error type: {}".format(error))
+
+    err, res = call[0], call[1:]
+    if err.value:
+        raise RuntimeError(
+            "CUDA error code={}({})".format(
+                err.value, _cudaGetErrorEnum(err)
+            )
+        )
+    if len(res) == 1:
+        return res[0]
+    elif len(res) == 0:
+        return None
+    else:
+        return res
+
+
+def create_cuda_context(device):
+    """
+    Create CUDA context with version-aware API handling.
+
+    Handles different CUDA API versions based on actual documented signatures:
+    - CUDA 11.8-12.9: cuCtxCreate(flags, device) - 2 arguments
+    - CUDA 13.0+: cuCtxCreate(ctxCreateParams, flags, device) - 3 arguments
+
+    Args:
+        device: CUDA device handle from cuDeviceGet
+
+    Returns:
+        CUDA context handle
+    """
+    # Try different API versions
+    try:
+        # Try CUDA 13.0+ API first (3 arguments with ctxCreateParams)
+        # cuCtxCreate(ctxCreateParams, flags, device)
+        return cuda_call(cuda.cuCtxCreate(None, 0, device))
+    except TypeError:
+        # CUDA 11.8-12.9 API: cuCtxCreate(flags, device)
+        return cuda_call(cuda.cuCtxCreate(0, device))
+
+
+
+class HostDeviceMem:
+    """Pair of host and device memory using RAII composition"""
+    def __init__(self, size: int, dtype: Optional[np.dtype] = None):
+        if dtype is None:
+            dtype = np.dtype(np.uint8)
+        else:
+            dtype = np.dtype(dtype)
+        self._size = size
+        self._dtype = dtype
+
+        # Use RAII classes for memory management
+        self._host_mem = PinnedHostMem(size, dtype)
+        self._device_mem = DeviceMem(size * dtype.itemsize)
+
+    @property
+    def host(self) -> np.ndarray:
+        # Return the array directly - ArrayWithOwner ensures proper lifetime management
+        return self._host_mem.array
+
+    @host.setter
+    def host(self, data: Union[np.ndarray, bytes]):
+        # Delegate to PinnedHostMem for proper data handling
+        self._host_mem.array = data
+
+    @property
+    def device_ptr(self) -> int:
+        """Device memory pointer"""
+        return self._device_mem.device_ptr
+
+    @property
+    def nbytes(self) -> int:
+        return self._host_mem.nbytes
+
+
+    def __str__(self):
+        return f"Host:\n{self.host}\nDevice:\n{self.device_ptr}\nSize:\n{self.nbytes}\n"
+
+    def __repr__(self):
+        return self.__str__()
+
+
+class DeviceMem:
+    """Device-only memory allocation for cases where host memory is not needed"""
+    def __init__(self, size: int):
+        self._device_ptr = cuda_call(cudart.cudaMalloc(size))
+        self._nbytes = size
+
+    @property
+    def device_ptr(self) -> int:
+        """Device memory pointer"""
+        return self._device_ptr
+
+    @property
+    def nbytes(self) -> int:
+        return self._nbytes
+
+    def free(self):
+        """Explicitly free device memory"""
+        if self._device_ptr is not None:
+            try:
+                cuda_call(cudart.cudaFree(self._device_ptr))
+                self._device_ptr = None
+            except Exception:
+                # Log but don't raise - cleanup should be best effort
+                pass
+
+    def __str__(self):
+        return f"Device:\n{self.device_ptr}\nSize:\n{self.nbytes}\n"
+
+    def __repr__(self):
+        return self.__str__()
+
+    def __del__(self):
+        # Fallback cleanup - not guaranteed to be called
+        self.free()
+
+
+class PinnedHostMem:
+    """Pinned host memory allocation for faster GPU transfers"""
+    def __init__(self, size: int, dtype: Optional[np.dtype] = None):
+        if dtype is None:
+            dtype = np.dtype(np.uint8)
+        else:
+            dtype = np.dtype(dtype)
+        nbytes = size * dtype.itemsize
+        host_mem = cuda_call(cudart.cudaMallocHost(nbytes))
+
+        self._host_ptr = host_mem
+        self._host_size = size
+        self._nbytes = nbytes
+        self._dtype = dtype
+
+    @property
+    def array(self) -> np.ndarray:
+        # Wrap the host buffer as uint8 first, then view as the real dtype.
+        # np.ctypeslib.as_ctypes_type has no mapping for dtypes without a ctypes
+        # equivalent (float16/bfloat16, FP8, INT4), and would raise NotImplementedError
+        # there; the byte-view round-trip avoids that call entirely.
+        ptr = ctypes.cast(ctypes.c_void_p(self._host_ptr), ctypes.POINTER(ctypes.c_ubyte))
+        host_u8 = np.ctypeslib.as_array(ptr, (self._nbytes,))
+        host_array = host_u8.view(self._dtype).reshape(self._host_size)
+        return ArrayWithOwner(host_array, self)
+
+    @array.setter
+    def array(self, data: Union[np.ndarray, bytes]):
+        """Set the array data with proper bounds checking"""
+        host_array = self.array  # Get the numpy array view
+        if isinstance(data, np.ndarray):
+            if data.size > self._host_size:
+                raise ValueError(
+                    f"Tried to fit an array of size {data.size} into host memory of size {self._host_size}"
+                )
+            np.copyto(host_array[:data.size], data.flat, casting='safe')
+        else:
+            assert self._dtype == np.uint8
+            host_array[:self.nbytes] = np.frombuffer(data, dtype=np.uint8)
+
+
+    @property
+    def nbytes(self) -> int:
+        return self._nbytes
+
+    def free(self):
+        """Explicitly free pinned host memory"""
+        if self._host_ptr is not None:
+            try:
+                cuda_call(cudart.cudaFreeHost(self._host_ptr))
+                self._host_ptr = None
+            except Exception:
+                # Log but don't raise - cleanup should be best effort
+                pass
+
+    def __str__(self):
+        return f"PinnedHost:\n{self.array}\nSize:\n{self.nbytes}\n"
+
+    def __repr__(self):
+        return self.__str__()
+
+    def __del__(self):
+        # Fallback cleanup - not guaranteed to be called
+        self.free()
+
+class CudaStreamContext:
+    """CUDA stream lifecycle management with context manager support"""
+    def __init__(self):
+        """Initialize CUDA stream"""
+        self._stream = cuda_call(cudart.cudaStreamCreate())
+
+    def __enter__(self):
+        """Create CUDA stream when entering context (if not already created)"""
+        if self._stream is None:
+            self._stream = cuda_call(cudart.cudaStreamCreate())
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Destroy CUDA stream when exiting context"""
+        if self._stream is not None:
+            try:
+                cuda_call(cudart.cudaStreamDestroy(self._stream))
+            except Exception:
+                # Silently handle cleanup failures
+                pass
+            self._stream = None
+
+    @property
+    def stream(self) -> cudart.cudaStream_t:
+        if self._stream is None:
+            raise RuntimeError("Stream not created. Use 'with' statement.")
+        return self._stream
+
+    def synchronize(self):
+        """Synchronize the stream"""
+        if self._stream is None:
+            raise RuntimeError("Stream not created. Use 'with' statement.")
+        cuda_call(cudart.cudaStreamSynchronize(self._stream))
+
+    def free(self):
+        """Explicitly free the CUDA stream"""
+        if self._stream is not None:
+            try:
+                cuda_call(cudart.cudaStreamDestroy(self._stream))
+                self._stream = None
+            except Exception:
+                # Log but don't raise - cleanup should be best effort
+                pass
+
+    def __del__(self):
+        """Cleanup stream on destruction"""
+        if hasattr(self, '_stream') and self._stream is not None:
+            self.free()
+
+    def __str__(self):
+        return f"CudaStreamContext: {self._stream}"
+
+    def __repr__(self):
+        return self.__str__()
+
+
+# Allocates all buffers required for an engine, i.e. host/device inputs/outputs.
+# Shape precedence: context.infer_shapes() (when context is given and yields fully
+# concrete shapes) > profile_idx max shape > engine.get_tensor_shape(). Using the
+# context branch avoids allocating profile-max-sized output buffers when the actual
+# input shapes are smaller.
+#
+# engine.get_tensor_profile_shape() is an INPUT-only API; for outputs it returns a
+# Dims sentinel with nbDims=-1, which would crash tuple()/__len__() downstream.
+# When the caller passes profile_idx without a configured context, this function
+# builds a temporary IExecutionContext, pins each input to its profile MAX, runs
+# infer_shapes() to derive output MAX shapes, allocates buffers, then discards the
+# temporary context.
+def allocate_buffers(
+    engine: trt.ICudaEngine,
+    profile_idx: Optional[int] = None,
+    context: Optional[trt.IExecutionContext] = None,
+):
+    inputs = []
+    outputs = []
+    bindings = []
+    tensor_names = [engine.get_tensor_name(i) for i in range(engine.num_io_tensors)]
+
+    def is_input(n):
+        return engine.get_tensor_mode(n) == trt.TensorIOMode.INPUT
+
+    # Decide a shape source. Prefer caller's context (if it yields concrete shapes);
+    # otherwise build a private one pinned to profile MAX inputs.
+    shape_ctx = None
+    owns_ctx = False
+    owns_stream = None
+
+    if context is not None:
+        try:
+            if not context.infer_shapes():
+                shape_ctx = context
+        except Exception:
+            # Any failure to infer shapes (RuntimeError, AttributeError, TRT-specific
+            # exceptions) just means we fall back to profile/static shapes below.
+            shape_ctx = None
+
+    try:
+        if shape_ctx is None and profile_idx is not None:
+            shape_ctx = engine.create_execution_context()
+            owns_ctx = True
+            if engine.num_optimization_profiles > 1:
+                owns_stream = cuda_call(cudart.cudaStreamCreate())
+                shape_ctx.set_optimization_profile_async(profile_idx, owns_stream)
+                cuda_call(cudart.cudaStreamSynchronize(owns_stream))
+            for name in tensor_names:
+                if is_input(name):
+                    max_shape = engine.get_tensor_profile_shape(name, profile_idx)[-1]
+                    shape_ctx.set_input_shape(name, max_shape)
+            if shape_ctx.infer_shapes():
+                raise RuntimeError(
+                    "infer_shapes failed; cannot size output buffers from profile_idx."
+                )
+
+        for binding in tensor_names:
+            if shape_ctx is not None:
+                shape = tuple(shape_ctx.get_tensor_shape(binding))
+            else:
+                shape = tuple(engine.get_tensor_shape(binding))
+                if any(s < 0 for s in shape):
+                    raise ValueError(
+                        f"Binding {binding} has dynamic shape, but no profile was specified."
+                    )
+            size = trt.volume(shape)
+            trt_type = engine.get_tensor_dtype(binding)
+
+            # Allocate host and device buffers
+            try:
+                dtype = np.dtype(trt.nptype(trt_type))
+                binding_memory = HostDeviceMem(size, dtype)
+            except TypeError: # no numpy support: create a byte array instead (BF16, FP8, INT4)
+                nbytes = int(size * trt_type.itemsize)
+                binding_memory = HostDeviceMem(nbytes)
+
+            # Append the device buffer to device bindings.
+            bindings.append(int(binding_memory.device_ptr))
+
+            # Append to the appropriate list.
+            if is_input(binding):
+                inputs.append(binding_memory)
+            else:
+                outputs.append(binding_memory)
+    finally:
+        if owns_stream is not None:
+            with contextlib.suppress(Exception):
+                cuda_call(cudart.cudaStreamDestroy(owns_stream))
+        if owns_ctx:
+            del shape_ctx
+
+    return inputs, outputs, bindings
+
+
+# Frees the resources allocated in allocate_buffers
+def free_buffers(inputs: List[HostDeviceMem], outputs: List[HostDeviceMem]):
+    """
+    Explicitly free CUDA memory resources.
+
+    While __del__ methods provide automatic cleanup, they are not guaranteed to be called.
+    This function provides explicit resource management for critical applications.
+    """
+    for inp in inputs:
+        if hasattr(inp, '_device_mem') and hasattr(inp._device_mem, 'free'):
+            inp._device_mem.free()
+        if hasattr(inp, '_host_mem') and hasattr(inp._host_mem, 'free'):
+            inp._host_mem.free()
+
+    for out in outputs:
+        if hasattr(out, '_device_mem') and hasattr(out._device_mem, 'free'):
+            out._device_mem.free()
+        if hasattr(out, '_host_mem') and hasattr(out._host_mem, 'free'):
+            out._host_mem.free()
+
+
+# Wrapper for cudaMemcpy which infers copy size and does error checking
+def memcpy_host_to_device(device_ptr: int, host_arr: np.ndarray):
+    cuda_call(cudart.cudaMemcpy(device_ptr, host_arr.ctypes.data, host_arr.nbytes, cudart.cudaMemcpyKind.cudaMemcpyHostToDevice))
+
+# Wrapper for cudaMemcpy which infers copy size and does error checking
+def memcpy_device_to_host(host_arr: np.ndarray, device_ptr: int):
+    cuda_call(cudart.cudaMemcpy(host_arr.ctypes.data, device_ptr, host_arr.nbytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost))
+
+
+# Additional CUDA wrapper functions for common operations
+
+
+def cuda_init():
+    """Initialize CUDA driver API with error checking."""
+    cuda_call(cuda.cuInit(0))
+
+
+def cuda_get_device(device_id: int = 0):
+    """Get CUDA device handle with error checking."""
+    return cuda_call(cuda.cuDeviceGet(device_id))
+
+
+# CUDA Runtime API functions (preferred over driver API when available)
+
+
+def cuda_memcpy_htod(device_ptr: int, host_data: np.ndarray):
+    """Copy data from host to device using CUDA runtime API with error checking.
+
+    Note: Consider using HostDeviceMem.host setter for integrated memory management.
+    """
+    cuda_call(cudart.cudaMemcpy(device_ptr, host_data, host_data.nbytes, cudart.cudaMemcpyKind.cudaMemcpyHostToDevice))
+
+
+def _do_inference_base(inputs, outputs, stream, execute_async_func):
+    # Transfer input data to the GPU.
+    kind = cudart.cudaMemcpyKind.cudaMemcpyHostToDevice
+    [cuda_call(cudart.cudaMemcpyAsync(inp.device_ptr, inp.host.ctypes.data, inp.nbytes, kind, stream)) for inp in inputs]
+    # Run inference.
+    execute_async_func()
+    # Transfer predictions back from the GPU.
+    kind = cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost
+    [cuda_call(cudart.cudaMemcpyAsync(out.host.ctypes.data, out.device_ptr, out.nbytes, kind, stream)) for out in outputs]
+    # Synchronize the stream
+    cuda_call(cudart.cudaStreamSynchronize(stream))
+    # Return only the host outputs.
+    return [out.host.copy() for out in outputs]
+
+
+# This function is generalized for multiple inputs/outputs.
+# inputs and outputs are expected to be lists of HostDeviceMem objects.
+def do_inference(context, engine, bindings, inputs, outputs, stream):
+    """
+    Perform inference using the provided context and stream.
+
+    Usage with context manager:
+        with stream:  # Ensures proper stream lifecycle
+            outputs = do_inference(context, engine, bindings, inputs, outputs, stream)
+    """
+    stream_handle = stream.stream
+    def execute_async_func():
+        context.execute_async_v3(stream_handle=stream_handle)
+    # Setup context tensor address.
+    num_io = engine.num_io_tensors
+    for i in range(num_io):
+        context.set_tensor_address(engine.get_tensor_name(i), bindings[i])
+    return _do_inference_base(inputs, outputs, stream_handle, execute_async_func)

samples/python/cute_dsl_plugin/README.md

@@ -0,0 +1,206 @@
+# RMSNorm PluginV3 with CuteDSL
+
+This sample shows how to author a TensorRT `IPluginV3` whose `enqueue()` invokes a kernel written in [CuteDSL](https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/overview.html), CUTLASS's Python DSL. The operator is **RMSNorm**, a normalization used in essentially every modern LLM (Llama, Mistral, Qwen, etc.).
+
+New to TensorRT plugins? Start with [`non_zero_plugin`](../non_zero_plugin) for a minimal end-to-end PluginV3 walkthrough. This sample assumes that flow and focuses on the **CuteDSL ↔ TensorRT** handoff rather than the kernel itself. For CuteDSL specifics, see the [CUTLASS Python DSL overview](https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/overview.html), the [`cutlass.cute` API reference](https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/cute_dsl_api/cute.html), and the [`cutlass/examples/python`](https://github.com/NVIDIA/cutlass/tree/main/examples/python) kernels.
+
+## What the sample demonstrates
+
+- Plugging a RMSNorm CuteDSL kernel into TensorRT through `IPluginV3` + `IPluginV3OneCore` + `IPluginV3OneBuild` + `IPluginV3OneRuntime`.
+- Sharing GPU buffers with the kernel zero-copy via `cupy.cuda.UnownedMemory` → `torch.as_tensor` → `cute.runtime.from_dlpack`.
+- Launching the kernel on the CUDA stream TRT passes into `enqueue()` (not the default stream), by declaring a `CUstream` parameter on the JIT-compiled launcher.
+- Reducing in FP32 for numerical stability on large hidden dims, then casting back to FP16 on store.
+
+## Requirements
+
+```bash
+cd $TRT_OSSPATH/samples/python/cute_dsl_plugin
+pip3 install -r requirements.txt
+```
+
+CuteDSL requires Ampere (SM80) or newer.
+
+## Running the sample
+
+```bash
+python3 rms_norm_plugin_cutedsl.py
+```
+
+The sample builds one engine with `num_tokens` declared dynamic over `[1, 512]` (opt=128) and `hidden_dim=1024`, then runs `num_tokens=128` through it. It prints `Inference result correct!` or `Inference result incorrect!`. Exit code is 0 on success.
+
+## Plugin interface
+
+- Inputs:
+  - `X`: rank-2 tensor, shape `(num_tokens, hidden_dim)`, FP16. `num_tokens` is dynamic (set by an optimization profile at build time); `hidden_dim` is static.
+  - `weight`: rank-1 tensor, shape `(hidden_dim,)`, FP32.
+- Output:
+  - `Y`: rank-2 tensor, shape `(num_tokens, hidden_dim)`, FP16.
+- Attribute:
+  - `epsilon`: scalar FP32. Required; the plugin asserts it is set.
+
+## Kernel design
+
+One CUDA block per token, 256 threads per block. Each block:
+
+1. Loads `X[token]` and accumulates `sum(x * x)` in FP32 across its threads.
+2. Reduces the partial sums through shared memory and a final warp shuffle.
+3. Computes `rms = rsqrt(sum_sq / H + eps)` once per block.
+4. Writes `Y[token, i] = X[token, i] * weight[i] * rms`.
+
+The reduction is always done in FP32. The result is cast back to FP16 only when storing `Y`.
+
+## How CuteDSL is wired into the plugin
+
+The CuteDSL side is **two free functions** plus a launch helper. The TensorRT side is the usual `IPluginV3` class. They meet inside `enqueue()`. Concretely:
+
+### 1. Author the kernel as a `@cute.kernel`
+
+```python
+@cute.kernel
+def rms_norm_kernel(mX: cute.Tensor, mW: cute.Tensor, mY: cute.Tensor,
+                    threads_per_block: cutlass.Constexpr,
+                    hidden_dim: cutlass.Constexpr,
+                    epsilon: cutlass.Constexpr):
+    ...
+```
+
+Tensor arguments are typed `cute.Tensor`. Values you want baked in at compile time are typed `cutlass.Constexpr` (here: block size, hidden dim, epsilon).
+
+### 2. Wrap the launch in a `@cute.jit`
+
+```python
+from cuda.bindings.driver import CUstream
+
+@cute.jit
+def rms_norm_launch(mX, mW, mY,
+                    num_tokens: cutlass.Int32,
+                    hidden_dim: cutlass.Constexpr,
+                    epsilon: cutlass.Constexpr,
+                    stream: CUstream):
+    rms_norm_kernel(mX, mW, mY, THREADS_PER_BLOCK, hidden_dim, epsilon).launch(
+        grid=(num_tokens, 1, 1),
+        block=(THREADS_PER_BLOCK, 1, 1),
+    )
+```
+
+Two things to notice:
+
+- `num_tokens` is typed `cutlass.Int32`, **not** `cutlass.Constexpr`. That means it is a runtime value that changes per call (it controls the grid dimension), so one compiled kernel handles every sequence length the optimization profile allows. `hidden_dim` stays `Constexpr` because the kernel's inner unroll trip count depends on it.
+- The `stream: CUstream` parameter is what tells the CuteDSL runtime which CUDA stream to launch on. It does not need to be passed to `.launch()` directly. The runtime picks it up automatically from any argument typed as `CUstream`.
+
+### 3. Subclass `IPluginV3` and the three capability interfaces
+
+```python
+class RmsNormPlugin(
+    trt.IPluginV3,            # top-level plugin interface
+    trt.IPluginV3OneCore,     # capability: core metadata
+    trt.IPluginV3OneBuild,    # capability: shape/format/serialization at build time
+    trt.IPluginV3OneRuntime,  # capability: enqueue and runtime hooks
+):
+```
+
+`IPluginV3` is the top-level interface (it owns `get_capability_interface()`, `clone()`, `destroy()`). The three `IPluginV3One*` mixins are the **capabilities** TRT will ask for at build and runtime. Inheriting from all four lets a single Python object answer to every request.
+
+Call **each** base `__init__` explicitly in your own `__init__`, and set the four book-keeping attributes that TensorRT reads: `plugin_namespace`, `plugin_name`, `plugin_version`, `num_outputs`. Set `timing_cache_id = ""` so TRT doesn't try to time per-instance.
+
+### 4. Methods you must implement
+
+These are the methods this sample overrides. Everything else falls through to defaults.
+
+| Method | Capability | Purpose |
+|---|---|---|
+| `get_capability_interface(type)` | `IPluginV3` | Return `self`. TRT calls this with `type ∈ {CORE, BUILD, RUNTIME}`. Our class inherits all three capability mixins, so the same object serves any of them. |
+| `get_output_data_types(input_types)` | `Build` | `Y` is FP16, matching `X`. |
+| `get_output_shapes(inputs, shape_inputs, exprBuilder)` | `Build` | `Y` has the same shape as `X` (`trt.DimsExprs(inputs[0])`). |
+| `get_fields_to_serialize()` | `Build` | Tells TRT which plugin attributes to save into the engine file. TRT calls this when building the engine, then hands the same fields back to the plugin creator at load time. Here we save `epsilon`. |
+| `configure_plugin(inp, out)` | `Build` | No-op for this kernel. Override if your kernel needs to precompute something from the I/O descriptors. |
+| `supports_format_combination(pos, in_out, num_inputs)` | `Build` | `X` and `Y` are FP16, `weight` is FP32, all `LINEAR`. |
+| `on_shape_change(inp, out)` | `Runtime` | No-op here. Override if you need to invalidate caches when shapes change. |
+| `enqueue(input_desc, output_desc, inputs, outputs, workspace, stream)` | `Runtime` | Where the kernel launch happens during inference. See below. |
+| `attach_to_context(context)` | `Runtime` | Return `self.clone()` so each execution context owns its own kernel cache. |
+| `set_tactic(tactic)` | `Runtime` | No-op since this plugin is single-tactic. |
+| `clone()` | `IPluginV3` | Return a fresh `RmsNormPlugin` with `__dict__` copied. Clear the JIT cache on the clone. |
+| `destroy()` | `IPluginV3` | Release (clear) the JIT cache when TRT is done with the plugin. |
+
+### 5. The handoff in `enqueue()`
+
+This is the only place CuteDSL touches TensorRT.
+
+```python
+def enqueue(self, input_desc, output_desc, inputs, outputs, workspace, stream):
+    # Read runtime shape and dtype from the descriptors.
+    num_tokens, hidden_dim = int(input_desc[0].dims[0]), int(input_desc[0].dims[1])
+
+    # Step A: raw device ptr -> cupy ndarray (zero copy, just a view).
+    x_cp = UnownedMemory(inputs[0], (num_tokens, hidden_dim), trt.nptype(input_desc[0].type)).d
+    # Step B: cupy ndarray -> torch tensor (still zero copy, shares __cuda_array_interface__).
+    x_t = torch.as_tensor(x_cp, device="cuda")
+    # Step C: torch tensor -> CuteDSL tensor via dlpack.
+    mX = from_dlpack(x_t, assumed_align=16)
+    # (same three steps for `weight` and `Y`)
+
+    # Step D: compile once per (hidden_dim, epsilon) and cache.
+    key = (hidden_dim, self.epsilon)
+    if key not in self._compiled:
+        self._compiled[key] = cute.compile(
+            rms_norm_launch, mX, mW, mY,
+            num_tokens, hidden_dim, self.epsilon, make_fake_stream())
+
+    # Step E: launch on TRT's stream.
+    self._compiled[key](mX, mW, mY, num_tokens, CUstream(stream))
+```
+
+A few things to note:
+
+- `cute.compile()` is the expensive step. The `self._compiled` dict makes sure it runs only once per `(hidden_dim, epsilon)` for the life of the plugin instance. The first call is effectively a warmup.
+- The `stream` argument that TRT passes into `enqueue()` is the CUDA stream the engine wants the work scheduled on. Always launch on that stream, not the default one, or you break TRT's stream ordering and risk hangs or races when an engine is run concurrently with other CUDA work.
+
+#### Why the cupy → torch → dlpack chain?
+
+TRT gives the plugin **raw integer device pointers** (`inputs[0]`, `outputs[0]`, ...) along with shape and dtype descriptors. CuteDSL wants a `cute.Tensor`. There is no single API that converts the first to the second, so the sample bridges them through two protocol hops. Each hop is **zero-copy**: we never touch the bytes, we only re-type them.
+
+| Hop | What it does | Why it's there |
+|---|---|---|
+| `cupy.cuda.UnownedMemory` | Wraps the raw integer pointer plus shape/dtype into a `cupy.ndarray`. | The only public Python API that wraps a **foreign** device pointer (one we did not allocate and must not free). Plain `torch.as_tensor` does not accept a raw `int` + shape + dtype. |
+| `torch.as_tensor(cp_array, device="cuda")` | Reads CuPy's `__cuda_array_interface__` and returns a torch view of the same memory. | We need a producer that exposes `__dlpack__`, and the existing TRT plugin samples already use torch for this; CuteDSL's `from_dlpack` consumes torch's DLPack capsule directly. |
+| `cute.runtime.from_dlpack(torch_tensor, assumed_align=16)` | Reads torch's `__dlpack__()` capsule and produces a `cute.Tensor`. | The published CuteDSL entry point for "take an external GPU buffer and use it as a `cute.Tensor`". |
+
+In short: TRT speaks raw pointers (because the C++ plugin ABI is framework-agnostic), CuteDSL speaks `cute.Tensor`, and the cupy → torch → dlpack chain is the shortest path between them that avoids both a copy and ownership confusion.
+
+### 6. Plugin creator
+
+```python
+class RmsNormPluginCreator(trt.IPluginCreatorV3One):
+    def __init__(self):
+        trt.IPluginCreatorV3One.__init__(self)
+        self.name = "RmsNormPlugin"
+        self.plugin_namespace = ""
+        self.plugin_version = "1"
+        self.field_names = trt.PluginFieldCollection([
+            trt.PluginField("epsilon", np.array([], dtype=np.float32),
+                            trt.PluginFieldType.FLOAT32),
+        ])
+
+    def create_plugin(self, name, fc, phase):
+        return RmsNormPlugin(fc)
+```
+
+Nothing CuteDSL-specific here. The `field_names` list the attributes that `create_plugin()` expects to find in the `PluginFieldCollection`.
+
+### 7. Building the network
+
+`network.add_plugin_v3([X, W], [], plugin)`. The first list is the data inputs, the second is shape-tensor inputs (none here), the third is the plugin object returned by the creator. The rest of the engine build (registering the creator, calling `engine_from_network`, running with `TrtRunner`) is the same as for any other PluginV3.
+
+## Limitations
+
+- `hidden_dim` is static at engine build time (the CuteDSL kernel is JIT-compiled per `hidden_dim`). `num_tokens` is dynamic within the optimization profile's `[min, max]` range.
+- The kernel requires `hidden_dim >= 256` (the same as `THREADS_PER_BLOCK`), which is checked in `build_engine()`.
+- As with all Python plugins, the engine cannot be deserialized outside a Python interpreter that has the plugin classes available.
+
+## License
+
+For terms and conditions for use, reproduction, and distribution, see the [TensorRT Software License Agreement](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sla/index.html) documentation.
+
+## Changelog
+
+May 2026: Initial release.

samples/python/cute_dsl_plugin/requirements.txt

@@ -0,0 +1,10 @@
+--extra-index-url https://pypi.ngc.nvidia.com
+nvidia-cutlass-dsl==4.5.0
+cupy-cuda12x==13.6.0
+torch==2.9.1
+polygraphy==0.50.1
+colored==2.3.2
+numpy==1.26.4
+pyyaml==6.0.3
+requests==2.33.0
+tqdm==4.66.4

samples/python/cute_dsl_plugin/rms_norm_plugin_cutedsl.py

@@ -0,0 +1,376 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 2026 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 cupy as cp
+import numpy as np
+import tensorrt as trt
+import torch
+import torch.nn.functional as F
+
+import cutlass
+import cutlass.cute as cute
+from cutlass.cute.runtime import from_dlpack, make_fake_stream
+from cuda.bindings.driver import CUstream
+
+from polygraphy.backend.trt import (
+    CreateConfig,
+    Profile,
+    TrtRunner,
+    create_network,
+    engine_from_network,
+)
+
+THREADS_PER_BLOCK = 256
+
+
+def volume(d):
+    return int(np.prod(d))
+
+
+class UnownedMemory:
+    def __init__(self, ptr, shape, dtype):
+        mem = cp.cuda.UnownedMemory(ptr, volume(shape) * cp.dtype(dtype).itemsize, self)
+        memptr = cp.cuda.MemoryPointer(mem, 0)
+        self.d = cp.ndarray(shape, dtype=dtype, memptr=memptr)
+
+
+@cute.kernel
+def rms_norm_kernel(
+    mX: cute.Tensor,
+    mW: cute.Tensor,
+    mY: cute.Tensor,
+    threads_per_block: cutlass.Constexpr,
+    hidden_dim: cutlass.Constexpr,
+    epsilon: cutlass.Constexpr,
+):
+    # One block per token. Threads cooperatively sum-of-squares across the
+    # hidden dim into shared memory, reduce to a single RMS value, then scale.
+    smem = cutlass.utils.SmemAllocator()
+    sdata = smem.allocate_tensor(
+        cutlass.Float32,
+        layout=cute.make_layout(threads_per_block),
+        byte_alignment=16,
+    )
+    rms = smem.allocate_tensor(cutlass.Float32, layout=cute.make_layout(1))
+
+    tidx, _, _ = cute.arch.thread_idx()
+    bidx, _, _ = cute.arch.block_idx()
+
+    # Per-thread partial sum-of-squares. Accumulate in fp32 regardless of
+    # input dtype to avoid catastrophic cancellation on large hidden dims.
+    local_sum = cutlass.Float32(0.0)
+    for i in cutlass.range(tidx, hidden_dim, threads_per_block):
+        x = cutlass.Float32(mX[bidx, i])
+        local_sum += x * x
+    sdata[tidx] = local_sum
+    cute.arch.sync_threads()
+
+    # Tree reduction in shared memory down to one warp, then warp shuffle.
+    if tidx < 128:
+        sdata[tidx] += sdata[tidx + 128]
+    cute.arch.sync_threads()
+
+    if tidx < 64:
+        sdata[tidx] += sdata[tidx + 64]
+    cute.arch.sync_threads()
+
+    if tidx < 32:
+        v = sdata[tidx] + sdata[tidx + 32]
+        v = cute.arch.warp_reduction_sum(v, threads_in_group=32)
+        if tidx == 0:
+            rms[0] = cute.math.rsqrt(v / hidden_dim + epsilon, fastmath=True)
+    cute.arch.sync_threads()
+
+    scale = rms[0]
+    for i in cutlass.range(tidx, hidden_dim, threads_per_block):
+        y = cutlass.Float32(mX[bidx, i]) * cutlass.Float32(mW[i]) * scale
+        mY[bidx, i] = y.to(mY.element_type)
+
+
+@cute.jit
+def rms_norm_launch(
+    mX: cute.Tensor,
+    mW: cute.Tensor,
+    mY: cute.Tensor,
+    num_tokens: cutlass.Int32,
+    hidden_dim: cutlass.Constexpr,
+    epsilon: cutlass.Constexpr,
+    stream: CUstream,
+):
+    # `num_tokens` is a runtime value (grid dim); only `hidden_dim` is baked
+    # in at compile time. That way one compiled kernel handles every sequence
+    # length the optimization profile allows.
+    # The `stream` parameter is consumed by the CuteDSL runtime and used as
+    # the launch stream; it does not need to be passed to .launch() directly.
+    rms_norm_kernel(mX, mW, mY, THREADS_PER_BLOCK, hidden_dim, epsilon).launch(
+        grid=(num_tokens, 1, 1),
+        block=(THREADS_PER_BLOCK, 1, 1),
+    )
+
+
+class RmsNormPlugin(
+    trt.IPluginV3,
+    trt.IPluginV3OneCore,
+    trt.IPluginV3OneBuild,
+    trt.IPluginV3OneRuntime,
+):
+    def __init__(self, fc=None):
+        trt.IPluginV3.__init__(self)
+        trt.IPluginV3OneCore.__init__(self)
+        trt.IPluginV3OneBuild.__init__(self)
+        trt.IPluginV3OneRuntime.__init__(self)
+
+        self.plugin_namespace = ""
+        self.plugin_name = "RmsNormPlugin"
+        self.plugin_version = "1"
+        self.num_outputs = 1
+        self.timing_cache_id = ""
+
+        # JIT-compiled kernel cache keyed by (hidden_dim, epsilon). Both are
+        # baked into the kernel as `Constexpr`, so a change in either yields
+        # a different binary.
+        self._compiled = {}
+
+        # `fc is None` is the path clone() takes; the cloned plugin then
+        # picks up epsilon via __dict__.update(). For the build path TRT
+        # always hands us a populated fc, and we require epsilon in it.
+        if fc is not None:
+            fields = {f.name: f for f in fc}
+            assert "epsilon" in fields, "RmsNormPlugin requires an 'epsilon' field"
+            self.epsilon = float(fields["epsilon"].data[0])
+
+    def get_capability_interface(self, type):
+        """Return the object that implements the requested capability.
+
+        TRT calls this with `type` set to CORE, BUILD, or RUNTIME. Because
+        this class inherits all three capability mixins, the same `self` can
+        serve any of them.
+        """
+        return self
+
+    def get_output_data_types(self, input_types):
+        """Output dtype is FP16 (matches X)."""
+        return [input_types[0]]
+
+    def get_output_shapes(self, inputs, shape_inputs, exprBuilder):
+        """Y has the same shape as X (N, H)."""
+        return [trt.DimsExprs(inputs[0])]
+
+    def get_fields_to_serialize(self):
+        """Tell TRT which plugin attributes to save into the engine.
+
+        TRT calls this at build time and hands the same fields back to the
+        creator's `create_plugin()` at engine load time. We only need to
+        persist `epsilon`; everything else is either fixed or recomputed.
+        """
+        return trt.PluginFieldCollection(
+            [
+                trt.PluginField(
+                    "epsilon",
+                    np.array([self.epsilon], dtype=np.float32),
+                    trt.PluginFieldType.FLOAT32,
+                )
+            ]
+        )
+
+    def configure_plugin(self, inp, out):
+        """No-op. We precompute nothing from the I/O descriptors at build time."""
+        pass
+
+    def on_shape_change(self, inp, out):
+        """No-op. The only dimension that can change between calls is
+        `num_tokens` (`hidden_dim` is static). The kernel cache is keyed on
+        `(hidden_dim, epsilon)`, so a change in `num_tokens` doesn't
+        invalidate anything.
+        """
+        pass
+
+    def supports_format_combination(self, pos, in_out, num_inputs):
+        """Accept LINEAR layout for all three positions, with `X` and `Y` FP16
+        and `weight` FP32.
+        """
+        assert num_inputs == 2
+        assert pos < len(in_out)
+
+        desc = in_out[pos].desc
+        if desc.format != trt.TensorFormat.LINEAR:
+            return False
+        if pos == 0 or pos == 2:
+            return desc.type == trt.DataType.HALF
+        if pos == 1:
+            return desc.type == trt.DataType.FLOAT
+        raise AssertionError(f"unexpected pos={pos}")
+
+    def enqueue(self, input_desc, output_desc, inputs, outputs, workspace, stream):
+        """Launch the CuteDSL kernel on TRT's stream.
+
+        Read shape and dtype from the descriptors, wrap the raw device
+        pointers as CuteDSL tensors zero-copy via cupy + dlpack, JIT-compile
+        once per `(hidden_dim, epsilon)` (cached), then launch.
+        """
+        x_dims = input_desc[0].dims
+        w_dims = input_desc[1].dims
+        assert len(x_dims) == 2, "X must be rank-2 (num_tokens, hidden_dim)"
+        assert len(w_dims) == 1, "weight must be rank-1 (hidden_dim,)"
+        num_tokens, hidden_dim = int(x_dims[0]), int(x_dims[1])
+        assert int(w_dims[0]) == hidden_dim
+
+        x_np = trt.nptype(input_desc[0].type)
+        w_np = trt.nptype(input_desc[1].type)
+        y_np = trt.nptype(output_desc[0].type)
+
+        # TRT hands us raw int pointers into device memory it already owns.
+        # CuteDSL wants `cute.Tensor` objects. We bridge the two without
+        # copying via two protocol conversions:
+        #   raw int ptr
+        #     -> cupy.UnownedMemory: the only public Python API that wraps a
+        #        foreign device pointer (we do NOT own it; cupy must not free
+        #        it) into a typed, shaped ndarray.
+        #     -> torch.as_tensor: reads cupy's __cuda_array_interface__ and
+        #        gives us a torch view over the same bytes.
+        #     -> cute.runtime.from_dlpack: reads torch's __dlpack__ capsule
+        #        and produces the cute.Tensor the kernel actually sees.
+        # Each hop is zero-copy; we're just re-typing the same GPU bytes.
+        x_t = torch.as_tensor(UnownedMemory(inputs[0], (num_tokens, hidden_dim), x_np).d, device="cuda")
+        w_t = torch.as_tensor(UnownedMemory(inputs[1], (hidden_dim,), w_np).d, device="cuda")
+        y_t = torch.as_tensor(UnownedMemory(outputs[0], (num_tokens, hidden_dim), y_np).d, device="cuda")
+
+        mX = from_dlpack(x_t, assumed_align=16)
+        mW = from_dlpack(w_t, assumed_align=16)
+        mY = from_dlpack(y_t, assumed_align=16)
+
+        # Launch the kernel on the CUDA stream TRT gave us, not the default
+        # stream. CuteDSL picks the stream from a `CUstream`-typed argument
+        # of the @cute.jit launcher; at compile time we hand it a placeholder.
+        key = (hidden_dim, self.epsilon)
+        if key not in self._compiled:
+            self._compiled[key] = cute.compile(
+                rms_norm_launch, mX, mW, mY, num_tokens, hidden_dim, self.epsilon, make_fake_stream()
+            )
+        self._compiled[key](mX, mW, mY, num_tokens, CUstream(stream))
+
+    def attach_to_context(self, context):
+        """Give each execution context its own plugin copy so the JIT cache
+        is not shared across concurrent inferences.
+        """
+        return self.clone()
+
+    def set_tactic(self, tactic):
+        """No-op. This plugin advertises a single tactic, so there is nothing
+        to select.
+        """
+        pass
+
+    def clone(self):
+        """Return a copy of this plugin with an empty JIT cache.
+
+        TRT clones the plugin per execution context. State is copied via
+        `__dict__.update`; the compiled-kernel dict is intentionally reset
+        so each context owns its own cubins.
+        """
+        cloned = RmsNormPlugin()
+        cloned.__dict__.update(self.__dict__)
+        cloned._compiled = {}
+        return cloned
+
+    def destroy(self):
+        """Release the JIT cache when TRT is done with the plugin."""
+        self._compiled.clear()
+
+
+class RmsNormPluginCreator(trt.IPluginCreatorV3One):
+    def __init__(self):
+        trt.IPluginCreatorV3One.__init__(self)
+        self.name = "RmsNormPlugin"
+        self.plugin_namespace = ""
+        self.plugin_version = "1"
+        self.field_names = trt.PluginFieldCollection(
+            [trt.PluginField("epsilon", np.array([], dtype=np.float32), trt.PluginFieldType.FLOAT32)]
+        )
+
+    def create_plugin(self, name, fc, phase):
+        return RmsNormPlugin(fc)
+
+
+def build_engine(hidden_dim, min_tokens, opt_tokens, max_tokens, epsilon):
+    """Build an engine where num_tokens is dynamic in [min, max]."""
+    # Fails fast on `hidden_dim < THREADS_PER_BLOCK` since the kernel's tree
+    # reduction uses hard-coded offsets (128, 64, 32) over shared memory,
+    # so anything smaller would read uninitialized slots.
+    if hidden_dim < THREADS_PER_BLOCK:
+        raise ValueError(f"RmsNormPlugin requires hidden_dim >= {THREADS_PER_BLOCK}, got {hidden_dim}")
+
+    plg_registry = trt.get_plugin_registry()
+    if plg_registry.get_creator("RmsNormPlugin", "1", "") is None:
+        plg_registry.register_creator(RmsNormPluginCreator(), "")
+
+    creator = plg_registry.get_creator("RmsNormPlugin", "1", "")
+    pfc = trt.PluginFieldCollection(
+        [trt.PluginField("epsilon", np.array([epsilon], dtype=np.float32), trt.PluginFieldType.FLOAT32)]
+    )
+    plugin = creator.create_plugin("RmsNormPlugin", pfc, trt.TensorRTPhase.BUILD)
+
+    builder, network = create_network(strongly_typed=True)
+    # X has dynamic num_tokens (axis 0 = -1) and static hidden_dim.
+    X = network.add_input(name="X", dtype=trt.float16, shape=(-1, hidden_dim))
+    W = network.add_input(name="weight", dtype=trt.float32, shape=(hidden_dim,))
+    out = network.add_plugin_v3([X, W], [], plugin)
+    out.get_output(0).name = "Y"
+    network.mark_output(out.get_output(0))
+
+    profile = Profile().add(
+        "X",
+        min=(min_tokens, hidden_dim),
+        opt=(opt_tokens, hidden_dim),
+        max=(max_tokens, hidden_dim),
+    )
+    return engine_from_network((builder, network), CreateConfig(profiles=[profile]))
+
+
+def main():
+    cc_major, cc_minor = torch.cuda.get_device_capability()
+    assert cc_major >= 8, f"CuteDSL requires SM80+; this GPU is SM{cc_major}{cc_minor}."
+
+    epsilon = 1e-5
+    hidden_dim = 1024
+    num_tokens = 128
+
+    # Build with num_tokens dynamic in [1, 512]. The kernel is compiled once
+    # for hidden_dim=1024 and serves every num_tokens within that range.
+    engine = build_engine(hidden_dim, min_tokens=1, opt_tokens=128, max_tokens=512, epsilon=epsilon)
+
+    # fp16 cast sets the tolerance floor.
+    rtol, atol = 1e-2, 1e-2
+    torch.manual_seed(0)
+    x = torch.randn(num_tokens, hidden_dim, device="cuda", dtype=torch.float16)
+    w = torch.randn(hidden_dim, device="cuda", dtype=torch.float32)
+    y_ref = F.rms_norm(x.float(), (hidden_dim,), w, eps=epsilon).to(torch.float16)
+
+    with TrtRunner(engine, "trt_runner") as runner:
+        outputs = runner.infer({"X": x.cpu().numpy().astype(np.float16), "weight": w.cpu().numpy()})
+        y = torch.as_tensor(outputs["Y"]).to(torch.float16)
+
+    max_diff = (y - y_ref.cpu()).abs().max().item()
+    label = f"N={num_tokens}, H={hidden_dim}"
+    assert torch.allclose(
+        y, y_ref.cpu(), rtol=rtol, atol=atol
+    ), f"[{label}] Inference result incorrect! (max abs diff={max_diff:.3e})"
+    print(f"[{label}] Inference result correct! (max abs diff={max_diff:.3e})")
+
+
+if __name__ == "__main__":
+    main()

samples/python/dds_faster_rcnn/README.md

@@ -0,0 +1,105 @@
+# DDS Faster R-CNN Object Detection in TensorRT
+## Introduction
+The `dds_faster_rcnn` sample demonstrates the usage of [tensorrt.IOutputAllocator](https://docs.nvidia.com/deeplearning/tensorrt/api/python_api/infer/Core/ExecutionContext.html#tensorrt.IOutputAllocator) in TensorRT to execute networks with data-dependent shape (DDS) outputs. In this sample, we showcase an end-to-end workflow for building and running an object detection model [Faster-RCNN](https://arxiv.org/abs/1506.01497).
+
+### What are Data-Dependent Shapes (DDS)?
+Data-Dependent Shapes (DDS) refer to shapes of layer outputs in a neural network which depend on the input data to the layer; in other words, it cannot be inferred solely by inspecting the shapes of the layer's input tensors.  An example of this is the output shape of the `INonZeroLayer`, which is determined by the number of non-zero elements in the input tensor.
+
+DDS outputs are common in models that involve dynamic processing, such as object detection, segmentation, and natural language processing.
+
+### What is an `IOutputAllocator`?
+An `IOutputAllocator` is an interface in TensorRT that defines a class responsible for dynamically allocating and managing the device memory for output tensors of a TensorRT engine. The class implementing this interface must provide a way to allocate and deallocate memory for output tensors, which can vary in size depending on the input data.
+
+### Why do we need to implement `IOutputAllocator`
+In traditional models, the output shapes are typically fixed and known at build time. However, in the case of data-dependent shaped (DDS) outputs, the output size is only known at inference time. This means that the memory allocation for output tensors cannot be determined until the model is actually run with a specific input. To handle this situation, TensorRT provides the `IOutputAllocator` interface, which allows developers to implement a custom memory allocation strategy for DDS outputs. By implementing this interface, developers can ensure that the output tensors are properly allocated and deallocated during inference, avoiding potential memory issues and improving the overall performance of the model.
+
+### How does `IOutputAllocator` work?
+To implement the `IOutputAllocator` interface, you need to provide implementations for the following two key methods:
+
+- `reallocate_output_async(self, tensor_name, memory, size, alignment, stream)`: This method is responsible for allocating or reallocating memory for an output tensor. It is called during the inference phase when the output tensor size is known. The method takes in parameters such as the tensor name, current memory address, new size, alignment, and CUDA stream, and returns the new memory address.
+- `notify_shape(self, tensor_name, shape)`: This method is used to notify the allocator of a change in the shape of an output tensor. It is typically called after reallocate_output_async() to update the allocator's internal state with the new shape information.
+During inference, the TensorRT engine will call these methods to manage the memory allocation for DDS output tensors. The `IOutputAllocator` implementation is responsible for ensuring that the memory allocation is properly handled, taking into account factors such as memory fragmentation, alignment, and performance optimization.
+
+Here is a high-level overview of the workflow:
+
+1. Instantiate the output allocator and attach to TensorRT with [IExecutionContext.set_output_allocator()](https://docs.nvidia.com/deeplearning/tensorrt/api/python_api/infer/Core/ExecutionContext.html#tensorrt.IExecutionContext.set_output_allocator)
+1. The TensorRT engine determines that an output tensor needs to be allocated or reallocated.
+1. `reallocate_output_async` is called to allocate or reallocate memory for the output tensor.
+1. The allocator updates its internal state and returns the new memory address.
+1. The TensorRT engine uses the new memory address to store the output tensor data.
+1. `notify_shape()` method is called to update the allocator's internal state with the new shape information.
+
+By implementing the `IOutputAllocator` interface, developers can create custom memory allocation strategies that optimize performance, reduce memory fragmentation, and improve the overall efficiency of their model inference.
+
+## Setup
+We recommend running these scripts on an environment with TensorRT >= 10.8.0.
+
+Install TensorRT as per the [TensorRT Install Guide](https://docs.nvidia.com/deeplearning/tensorrt/latest/installing-tensorrt/installing.html). You will need to make sure the Python bindings for TensorRT are also installed correctly, these are available by installing the `python3-libnvinfer` and `python3-libnvinfer-dev` packages on your TensorRT download.
+
+To simplify TensorRT installation, use an NGC Docker Image, such as:
+
+```bash
+docker pull nvcr.io/nvidia/tensorrt:25.01-py3
+```
+
+Install all dependencies listed in requirements.txt:
+
+```bash
+pip3 install -r requirements.txt
+```
+
+## Model Conversion
+To start, download the pre-trained Faster R-CNN model in ONNX format using the following command:
+
+```bash
+wget https://github.com/onnx/models/raw/refs/heads/main/validated/vision/object_detection_segmentation/faster-rcnn/model/FasterRCNN-12.onnx
+```
+
+With the ONNX model downloaded, run the following command to prepare it for TensorRT engine conversion:
+
+```bash
+python3 modify_onnx.py \
+    --input ./FasterRCNN-12.onnx \
+    --output ./fasterrcnn12_trt.onnx
+```
+
+This will create a modified ONNX graph file that is ready for conversion to a TensorRT engine.
+
+## Build TensorRT Engine
+
+To build the TensorRT engine, run the following command:
+
+```bash
+python3 build_engine.py \
+    --onnx ./fasterrcnn12_trt.onnx \
+    --engine ./fasterrcnn12_trt.engine
+```
+
+## Inference
+To test the built TensorRT engine, download a test image using the following command:
+
+```bash
+wget https://onnxruntime.ai/images/demo.jpg
+```
+
+Then, run the inference script using the following command:
+
+```
+python3 infer.py \
+    --engine ./fasterrcnn12_trt.engine \
+    --input ./demo.jpg \
+    --output ./output_dir \
+    --labels labels_coco_80.txt
+```
+This will perform object detection on the test image and save the output to the specified directory (`output_dir` in this case).
+
+# Changelog
+
+October 2025
+Migrate to strongly typed APIs.
+
+August 2025
+Removed support for Python versions < 3.10.
+
+February 2025
+Initial release

samples/python/dds_faster_rcnn/build_engine.py

@@ -0,0 +1,142 @@
+# SPDX-FileCopyrightText: Copyright (c) 1993-2025 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 sys
+import logging
+import argparse
+
+import tensorrt as trt
+
+sys.path.insert(1, os.path.join(os.path.dirname(os.path.realpath(__file__)), os.pardir))
+import common
+
+logging.basicConfig(level=logging.INFO)
+logging.getLogger("EngineBuilder").setLevel(logging.INFO)
+log = logging.getLogger("EngineBuilder")
+
+
+class EngineBuilder:
+    """
+    Parses an ONNX graph and builds a TensorRT engine from it.
+    """
+
+    def __init__(self, verbose=False, workspace=8):
+        """
+        :param verbose: If enabled, a higher verbosity level will be set on the TensorRT logger.
+        :param workspace: Max memory workspace to allow, in Gb.
+        """
+        self.trt_logger = trt.Logger(trt.Logger.INFO)
+        if verbose:
+            self.trt_logger.min_severity = trt.Logger.Severity.VERBOSE
+
+        trt.init_libnvinfer_plugins(self.trt_logger, namespace="")
+
+        self.builder = trt.Builder(self.trt_logger)
+        self.config = self.builder.create_builder_config()
+        one_GiB = 2**30
+        self.config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace * one_GiB)
+        self.network = None
+        self.parser = None
+
+    def create_network(self, onnx_path):
+        """
+        Parse the ONNX graph and create the corresponding TensorRT network definition.
+        :param onnx_path: The path to the ONNX graph to load.
+        """
+
+        self.network = self.builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED))
+        self.parser = trt.OnnxParser(self.network, self.trt_logger)
+
+        onnx_path = os.path.realpath(onnx_path)
+        with open(onnx_path, "rb") as f:
+            if not self.parser.parse(f.read()):
+                for error in range(self.parser.num_errors):
+                    log.error(self.parser.get_error(error))
+                raise RuntimeError(
+                    f"Failed to load ONNX file: {onnx_path}. Check the logs for more details or run with --verbose."
+                )
+
+        log.info("Network Description")
+
+        profile = self.builder.create_optimization_profile()
+        profile.set_shape("image", min=(3, 1, 1), opt=(3, 800, 800), max=(3, 800, 1312))
+        self.config.add_optimization_profile(profile)
+
+        inputs = [self.network.get_input(i) for i in range(self.network.num_inputs)]
+        for input in inputs:
+            log.info(f"Input '{input.name}' with shape {input.shape} and dtype {input.dtype}")
+        outputs = [self.network.get_output(i) for i in range(self.network.num_outputs)]
+        for output in outputs:
+            log.info(f"Output '{output.name}' with shape {output.shape} and dtype {output.dtype}")
+
+    def create_engine(
+        self,
+        engine_path,
+    ):
+        """
+        Build the TensorRT engine and serialize it to disk.
+        :param engine_path: The path where to serialize the engine to.
+        """
+        engine_path = os.path.realpath(engine_path)
+        engine_dir = os.path.dirname(engine_path)
+        os.makedirs(engine_dir, exist_ok=True)
+        log.info(f"Building Engine in {engine_path}")
+
+        inputs = [self.network.get_input(i) for i in range(self.network.num_inputs)]
+
+        log.info(f"Reading timing cache from file: {args.timing_cache}")
+        common.setup_timing_cache(self.config, args.timing_cache)
+
+        engine_bytes = self.builder.build_serialized_network(self.network, self.config)
+        if engine_bytes is None:
+            raise RuntimeError("Failed to create engine. Check the logs for more details or run with --verbose.")
+
+        log.info(f"Serializing timing cache to file: {args.timing_cache}")
+        common.save_timing_cache(self.config, args.timing_cache)
+
+        with open(engine_path, "wb") as f:
+            log.info(f"Serializing engine to file: {engine_path}")
+            f.write(engine_bytes)
+
+
+def main(args):
+    builder = EngineBuilder(args.verbose, args.workspace)
+    builder.create_network(args.onnx)
+    builder.create_engine(
+        args.engine,
+    )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-o", "--onnx", required=True, help="The input ONNX model file to load")
+    parser.add_argument("-e", "--engine", required=True, help="The output path for the TRT engine")
+    parser.add_argument("-v", "--verbose", action="store_true", help="Enable more verbose log output")
+    parser.add_argument(
+        "-w",
+        "--workspace",
+        default=8,
+        type=int,
+        help="The max memory workspace size to allow in Gb, default: 8",
+    )
+    parser.add_argument(
+        "--timing_cache",
+        default="./timing.cache",
+        help="The file path for timing cache, default: ./timing.cache",
+    )
+    args = parser.parse_args()
+    main(args)

samples/python/dds_faster_rcnn/infer.py

@@ -0,0 +1,485 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2025 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 sys
+import time
+import argparse
+import numpy as np
+import tensorrt as trt
+from cuda.bindings import driver as cuda, runtime as cudart
+from PIL import Image
+from pathlib import Path
+import threading
+from visualize import visualize_detections
+
+sys.path.insert(1, os.path.join(os.path.dirname(os.path.realpath(__file__)), os.pardir))
+import common
+from common import cuda_call
+
+
+class AllocatorState:
+    """
+    Represents the state of an allocator for a tensor.
+    """
+
+    def __init__(self, ptr, size, dim=None):
+        """
+        :param ptr: The pointer to the allocated device memory.
+        :param size: The size of the allocated device memory.
+        :param dim: The dimensions of the tensor.
+        """
+        self.ptr = ptr
+        self.size = size
+        self.dim = dim
+        self.lock = threading.Lock()
+
+    def update(self, ptr=None, size=None, dims=None):
+        """
+        Updates the state of the allocator.
+
+        :param ptr: The new pointer to the allocated device memory. If None, the current pointer is not changed.
+        :param size: The new size of the allocated device memory. If None, the current size is not changed.
+        :param dims: The new dimensions of the tensor. If None, the current dimensions are not changed.
+        """
+        with self.lock:
+            if ptr is not None:
+                self.ptr = ptr
+            if size is not None:
+                self.size = size
+            if dims is not None:
+                self.dims = dims
+
+
+class MyOutputAllocator(trt.IOutputAllocator):
+    """
+    Custom output allocator class.
+    """
+
+    def __init__(self, verbose=False):
+        """
+        :param verbose: If True, enables verbose logging.
+        """
+        trt.IOutputAllocator.__init__(self)
+
+        self.lock = threading.Lock()
+        self.states = {}
+        self.verbose = verbose
+
+    def reallocate_output_async(self, tensor_name, current_memory, size, alignment, stream):
+        """
+        Reallocates output memory for the given tensor.
+
+        :param tensor_name: The name of the tensor.
+        :param current_memory: The current device memory pointer.
+        :param size: The new size of the device memory block.
+        :param alignment: The alignment of the device memory block.
+        :param stream: The CUDA stream.
+        :return: The new memory pointer.
+        """
+        size = max(size, 1)
+        ptr = current_memory
+        with self.lock:
+            if tensor_name not in self.states or size > self.states[tensor_name].size:
+                ptr = cuda_call(cudart.cudaMalloc(size))
+                if tensor_name in self.states:
+                    cuda_call(cudart.cudaFree(self.states[tensor_name].ptr))
+                    self.states[tensor_name].update(ptr=ptr, size=size)
+                else:
+                    self.states[tensor_name] = AllocatorState(ptr=ptr, size=size)
+                if self.verbose:
+                    print(f"Reallocated {size} bytes for tensor '{tensor_name}' to {ptr}")
+        return ptr
+
+    def notify_shape(self, tensor_name, dims):
+        """
+        Notifies the allocator of a change in the shape of the tensor.
+
+        :param tensor_name: The name of the tensor.
+        :param dims: The new dimensions of the tensor.
+        """
+        with self.lock:
+            assert tensor_name in self.states, f'Tensor "{tensor_name}" is not in states.'
+            self.states[tensor_name].update(dims=dims)
+            if self.verbose:
+                print(f"Updated shape for tensor '{tensor_name}': {dims}")
+
+    def __del__(self):
+        try:
+            with self.lock:
+                for tensor_name, item in self.states.items():
+                    if item.ptr is not None:
+                        cuda_call(cudart.cudaFree(item.ptr))
+                        if self.verbose:
+                            print(f"Freed memory for tensor '{tensor_name}'")
+                self.states.clear()
+        except Exception:
+            # Silently handle cleanup failures to prevent exceptions during object deletion
+            pass
+
+
+class PoolAllocator(trt.IGpuAsyncAllocator):
+    """
+    A custom GPU async allocator class that manages memory allocation and deallocation.
+
+    It utilizes the CUDA memory pool API to optimize memory allocation and minimize fragmentation.
+    """
+
+    def __init__(self):
+        """
+        Initializes the PoolAllocator instance.
+
+        Creates a CUDA memory pool with the specified properties and sets the release threshold to the maximum possible value.
+        """
+        trt.IGpuAsyncAllocator.__init__(self)
+
+        pool_props = cudart.cudaMemPoolProps()
+        pool_props.allocType = cudart.cudaMemAllocationType.cudaMemAllocationTypePinned
+        pool_props.handleTypes = cudart.cudaMemAllocationHandleType.cudaMemHandleTypeNone
+        pool_props.location.type = cudart.cudaMemLocationType.cudaMemLocationTypeDevice
+        pool_props.location.id = 0
+
+        self.pool = cuda_call(cudart.cudaMemPoolCreate(pool_props))
+
+        max_threshold = np.uint64(np.iinfo(np.uint64).max)
+        cuda_call(
+            cudart.cudaMemPoolSetAttribute(
+                self.pool, cudart.cudaMemPoolAttr.cudaMemPoolAttrReleaseThreshold, cuda.cuuint64_t(max_threshold)
+            )
+        )
+
+    def allocate_async(self, size: int, alignment: int, flags: int, stream: cudart.cudaStream_t):
+        """
+        Allocates memory asynchronously from the CUDA memory pool.
+
+        :param size: The size of the memory block to allocate.
+        :param alignment: The alignment of the memory block.
+        :param flags: The flags for the allocation.
+        :param stream: The CUDA stream for the allocation.
+        :return: The pointer to the allocated device memory.
+        """
+        ptr = cuda_call(cudart.cudaMallocFromPoolAsync(size, self.pool, stream))
+        return ptr
+
+    def deallocate_async(self, memory, stream: cudart.cudaStream_t):
+        """
+        Deallocates memory asynchronously.
+
+        :param memory: The pointer to the memory to deallocate.
+        :param stream: The CUDA stream for the deallocation.
+        :return: True if the deallocation was successful.
+        """
+        cuda_call(cudart.cudaFreeAsync(memory, stream))
+        return True
+
+    def __del__(self):
+        try:
+            if self.pool:
+                cuda_call(cudart.cudaMemPoolDestroy(self.pool))
+        except Exception:
+            # Silently handle cleanup failures to prevent exceptions during object deletion
+            pass
+
+
+class TensorRTInfer:
+    """
+    Implements inference for the FasterRCNN TensorRT engine.
+    """
+
+    def __init__(self, engine_path, use_custom_gpu_allocator=False, verbose=False):
+        """
+        Initializes the TensorRTInfer instance.
+
+        :param engine_path: The path to the serialized engine to load from disk.
+        :param use_custom_gpu_allocator: If True, uses a custom GPU allocator.
+        :param verbose: If True, enables verbose logging.
+        """
+        # Load TRT engine
+        self.logger = trt.Logger(trt.Logger.ERROR)
+        if verbose:
+            self.logger.min_severity = trt.Logger.VERBOSE
+        trt.init_libnvinfer_plugins(self.logger, namespace="")
+        with open(engine_path, "rb") as f, trt.Runtime(self.logger) as runtime:
+            assert runtime
+            if use_custom_gpu_allocator:
+                self.my_pool_allocator = PoolAllocator()
+                runtime.gpu_allocator = self.my_pool_allocator
+            self.engine = runtime.deserialize_cuda_engine(f.read())
+        assert self.engine
+        self.context = self.engine.create_execution_context()
+        assert self.context
+
+        self.my_output_allocator = MyOutputAllocator(verbose=True)
+        # Setup I/O bindings
+        self.inputs = []
+        self.outputs = []
+        self.allocations = []
+        for i in range(self.engine.num_io_tensors):
+            name = self.engine.get_tensor_name(i)
+            is_input = False
+            if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT:
+                is_input = True
+            dtype = trt.nptype(self.engine.get_tensor_dtype(name))
+
+            # trt.nptype returns a python 'type'. For here we want a numpy 'dtype' object
+            # instead to get more info about the dtype (dtype.itemsize in this case)
+            dtype = np.dtype(dtype)
+            shape = self.engine.get_tensor_shape(name)
+
+            # Use the max shape in the profile for dynamic shaped inputs
+            if is_input and any(value for value in shape if value < 0):
+                assert self.engine.num_optimization_profiles > 0
+                profile_shape = self.engine.get_tensor_profile_shape(name, 0)
+                assert len(profile_shape) == 3  # min,opt,max
+                # Set the *max* profile as binding shape
+                shape = profile_shape[2]
+
+            if is_input:
+                nbytes = np.prod(shape) * dtype.itemsize
+                allocation = cuda_call(cudart.cudaMalloc(nbytes))
+            else:
+                self.context.set_output_allocator(name, self.my_output_allocator)
+                allocation = cuda_call(
+                    cudart.cudaMalloc(128 * dtype.itemsize)
+                )  # Random number. More will be allocated using our custom allocator
+
+            binding = {
+                "index": i,
+                "name": name,
+                "dtype": dtype,
+                "shape": list(shape),
+                "allocation": allocation,
+            }
+            self.allocations.append(allocation)
+            if is_input:
+                self.inputs.append(binding)
+            else:
+                self.outputs.append(binding)
+            print(
+                f"{'Input' if is_input else 'Output'} '{binding['name']}' with shape {binding['shape']} and dtype {binding['dtype']}"
+            )
+
+        assert len(self.inputs) > 0
+        assert len(self.outputs) > 0
+        assert len(self.allocations) > 0
+
+    def input_spec(self):
+        """
+        Get the specs for the input tensor of the network. Useful to prepare memory allocations.
+        :return: Two items, the shape of the input tensor and its (numpy) datatype.
+        """
+        return self.inputs[0]["shape"], self.inputs[0]["dtype"]
+
+    def output_spec(self):
+        """
+        Get the specs for the output tensors of the network. Useful to prepare memory allocations.
+        :return: A list with two items per element, the shape and (numpy) datatype of each output tensor.
+        """
+        specs = []
+        for o in self.outputs:
+            specs.append((o["shape"], o["dtype"]))
+        return specs
+
+    def preprocess_image(self, image):
+        """
+        Preprocesses an image for inference. See also
+        https://github.com/onnx/models/tree/refs/heads/main/validated/vision/object_detection_segmentation/faster-rcnn#preprocessing-steps
+
+        :param image: The image to preprocess.
+        :return: The preprocessed image as a numpy array.
+        """
+        ratio = 800.0 / min(image.size[0], image.size[1])
+        image = image.resize((int(ratio * image.size[0]), int(ratio * image.size[1])), Image.BILINEAR)
+
+        # RGB -> BGR
+        image = np.array(image)[:, :, [2, 1, 0]].astype("float32")
+
+        # HWC -> CHW
+        image = np.transpose(image, [2, 0, 1])
+
+        # Normalize
+        mean_vec = np.array([102.9801, 115.9465, 122.7717])
+        for i in range(image.shape[0]):
+            image[i, :, :] = image[i, :, :] - mean_vec[i]
+
+        # Pad to be divisible of 32
+        padded_h = int(np.ceil(image.shape[1] / 32) * 32)
+        padded_w = int(np.ceil(image.shape[2] / 32) * 32)
+
+        padded_image = np.zeros((3, padded_h, padded_w), dtype=np.float32)
+        padded_image[:, : image.shape[1], : image.shape[2]] = image
+        image = padded_image
+
+        return image
+
+    def infer(self, arr):
+        """
+        Execute inference on an image.
+
+        :param arr: A numpy array for the input image values.
+        :return A list of outputs as numpy arrays.
+        """
+        # Copy I/O and Execute
+        common.memcpy_host_to_device(self.inputs[0]["allocation"], arr)
+        self.context.execute_v2(self.allocations)
+
+        # copy outputs to host
+        return_outputs = []
+        for output in self.outputs:
+            final_shape = self.my_output_allocator.states[output["name"]].dims
+            host_arr = np.random.random(final_shape).astype(output["dtype"])
+            device_ptr = self.my_output_allocator.states[output["name"]].ptr
+
+            nbytes = np.prod(final_shape) * output["dtype"].itemsize
+            common.memcpy_device_to_host(host_arr, device_ptr)
+
+            return_outputs.append(host_arr)
+
+        return return_outputs
+
+    def process(self, arr):
+        """
+        Execute inference on an image. The image should already be preprocessed. Memory
+        copying to and from the GPU device will be performed here.
+
+        :param arr: A numpy array holding the image values.
+        :return: A list of detected object with box, score, class included.
+        """
+        preprocess_arr = self.preprocess_image(arr.copy())
+        self.context.set_input_shape("image", preprocess_arr.shape)
+
+        # Run inference
+        outputs = self.infer(preprocess_arr)
+
+        # Post-process the results
+        scale = 800.0 / min(arr.size[0], arr.size[1])
+
+        boxes = outputs[0]
+        labels = outputs[1]
+        scores = outputs[2]
+        num = len(labels)
+
+        detections = []
+        for i in range(num):
+            if scores[i] > 0.9:
+                detections.append(
+                    {
+                        "xmin": boxes[i][0] / scale,
+                        "ymin": boxes[i][1] / scale,
+                        "xmax": boxes[i][2] / scale,
+                        "ymax": boxes[i][3] / scale,
+                        "score": scores[i],
+                        "class": labels[i] - 1,
+                    }
+                )
+        return detections
+
+
+def main(args):
+    if args.output:
+        args.output.resolve().mkdir(exist_ok=True, parents=True)
+
+    labels = []
+    if args.labels:
+        with open(args.labels) as f:
+            for label in f:
+                labels.append(label.strip())
+
+    trt_infer = TensorRTInfer(args.engine, args.use_custom_gpu_allocator, args.verbose)
+    if args.input:
+        print(f"\nInferring data in {args.input}")
+        image_paths = []
+        if args.input.is_dir():
+            for p in args.input.iterdir():
+                image_paths.append(p)
+        else:
+            image_paths.append(args.input)
+
+        for image_path in image_paths:
+            image = Image.open(image_path)
+            detections = trt_infer.process(image)
+            if args.output:
+                # Image Visualizations
+                output_path = args.output / f"{image_path.stem}.png"
+                visualize_detections(image_path, output_path, detections, labels)
+
+                # Text Results
+                output_results = ""
+                for d in detections:
+                    line = [
+                        d["xmin"],
+                        d["ymin"],
+                        d["xmax"],
+                        d["ymax"],
+                        d["score"],
+                    ]
+                    output_results += "\t".join([str(f) for f in line]) + "\n"
+                with open(args.output / f"{image_path.stem}.txt", "w") as f:
+                    f.write(output_results)
+    else:
+        print("No input provided, running in benchmark mode")
+        shape, dtype = trt_infer.input_spec()
+        batch = 255 * np.random.rand(*shape).astype(dtype)
+        trt_infer.context.set_input_shape("image", (batch.shape))
+        iterations = 200
+        times = []
+        for i in range(20):  # GPU warmup iterations
+            trt_infer.infer(batch)
+        for i in range(iterations):
+            start = time.time()
+            trt_infer.infer(batch)
+            times.append(time.time() - start)
+            print(f"Iteration {i+1} / {iterations}", end="\r")
+        print("Benchmark results include time for H2D and D2H memory copies")
+        print(f"Average Latency: {1000 * np.average(times):.3f} ms")
+
+    print("\nFinished Processing")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-e",
+        "--engine",
+        default=None,
+        required=True,
+        help="The serialized TensorRT engine",
+    )
+    parser.add_argument("-i", "--input", default=None, type=Path, help="Path to the image or directory to process")
+    parser.add_argument(
+        "-o",
+        "--output",
+        default=None,
+        type=Path,
+        help="Directory where to save the visualization results",
+    )
+    parser.add_argument(
+        "-l",
+        "--labels",
+        default="./labels_coco_80.txt",
+        help="File to use for reading the class labels from, default: ./labels_coco_80.txt",
+    )
+    parser.add_argument(
+        "-c",
+        "--use_custom_gpu_allocator",
+        action="store_true",
+        default=False,
+        help="Use a custom gpu allocator with CUDA memory pools for better performance",
+    )
+    parser.add_argument("-v", "--verbose", action="store_true", default=False, help="Set to verbose logging")
+    args = parser.parse_args()
+    main(args)

samples/python/dds_faster_rcnn/labels_coco_80.txt

@@ -0,0 +1,80 @@
+person
+bicycle
+car
+motorcycle
+airplane
+bus
+train
+truck
+boat
+traffic light
+fire hydrant
+stop sign
+parking meter
+bench
+bird
+cat
+dog
+horse
+sheep
+cow
+elephant
+bear
+zebra
+giraffe
+backpack
+umbrella
+handbag
+tie
+suitcase
+frisbee
+skis
+snowboard
+sports ball
+kite
+baseball bat
+baseball glove
+skateboard
+surfboard
+tennis racket
+bottle
+wine glass
+cup
+fork
+knife
+spoon
+bowl
+banana
+apple
+sandwich
+orange
+broccoli
+carrot
+hot dog
+pizza
+donut
+cake
+chair
+couch
+potted plant
+bed
+dining table
+toilet
+tv
+laptop
+mouse
+remote
+keyboard
+cell phone
+microwave
+oven
+toaster
+sink
+refrigerator
+book
+clock
+vase
+scissors
+teddy bear
+hair drier
+toothbrush

samples/python/dds_faster_rcnn/modify_onnx.py

@@ -0,0 +1,56 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2025 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 onnx_graphsurgeon as gs
+import onnx
+import numpy as np
+import argparse
+
+
+def modify_maskrcnn_opset12(path_to_model, output_path):
+    graph = gs.import_onnx(onnx.load(path_to_model))
+    """
+        Step 1: Remove unnecessary UINT8 cast
+            - Pattern match Cast[BOOL->UINT8] -> Cast[UINT8 -> BOOL]
+            - Fixes node 2838 - casts bool to uint8 for slice / gather. Can keep all operations in bool.
+    """
+    for node in graph.nodes:
+        if node.op == "Cast" and node.attrs["to"] == onnx.TensorProto.UINT8:
+            node.attrs["to"] = onnx.TensorProto.BOOL
+            node.outputs[0].dtype = np.bool_
+            # Need to modify output_node output to be bool as well.
+            for output_node in node.outputs[0].outputs:
+                output_node.outputs[0].dtype = np.bool_
+            print(f"Removed UINT8 casts in node {node.name}")
+
+    onnx.save(gs.export_onnx(graph.cleanup()), output_path)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-i",
+        "--input",
+        default="FasterRCNN-12.onnx",
+        help="Path to the onnx model obtained from https://github.com/onnx/models/raw/refs/heads/main/validated/vision/object_detection_segmentation/faster-rcnn/model/FasterRCNN-12.onnx",
+    )
+    parser.add_argument(
+        "-o", "--output", default="fasterrcnn12_trt.onnx", help="Desired path for the output onnx model"
+    )
+    args = parser.parse_args()
+
+    modify_maskrcnn_opset12(args.input, args.output)

samples/python/dds_faster_rcnn/requirements.txt

@@ -0,0 +1,5 @@
+Pillow==11.3.0
+cuda-python==12.9.0
+onnx==1.18.0
+onnx-graphsurgeon --index-url https://pypi.ngc.nvidia.com
+numpy==1.26.4

samples/python/dds_faster_rcnn/visualize.py

@@ -0,0 +1,199 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2025 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 PIL.Image as Image
+import PIL.ImageDraw as ImageDraw
+import PIL.ImageFont as ImageFont
+
+COLORS = [
+    "GoldenRod",
+    "MediumTurquoise",
+    "GreenYellow",
+    "SteelBlue",
+    "DarkSeaGreen",
+    "SeaShell",
+    "LightGrey",
+    "IndianRed",
+    "DarkKhaki",
+    "LawnGreen",
+    "WhiteSmoke",
+    "Peru",
+    "LightCoral",
+    "FireBrick",
+    "OldLace",
+    "LightBlue",
+    "SlateGray",
+    "OliveDrab",
+    "NavajoWhite",
+    "PaleVioletRed",
+    "SpringGreen",
+    "AliceBlue",
+    "Violet",
+    "DeepSkyBlue",
+    "Red",
+    "MediumVioletRed",
+    "PaleTurquoise",
+    "Tomato",
+    "Azure",
+    "Yellow",
+    "Cornsilk",
+    "Aquamarine",
+    "CadetBlue",
+    "CornflowerBlue",
+    "DodgerBlue",
+    "Olive",
+    "Orchid",
+    "LemonChiffon",
+    "Sienna",
+    "OrangeRed",
+    "Orange",
+    "DarkSalmon",
+    "Magenta",
+    "Wheat",
+    "Lime",
+    "GhostWhite",
+    "SlateBlue",
+    "Aqua",
+    "MediumAquaMarine",
+    "LightSlateGrey",
+    "MediumSeaGreen",
+    "SandyBrown",
+    "YellowGreen",
+    "Plum",
+    "FloralWhite",
+    "LightPink",
+    "Thistle",
+    "DarkViolet",
+    "Pink",
+    "Crimson",
+    "Chocolate",
+    "DarkGrey",
+    "Ivory",
+    "PaleGreen",
+    "DarkGoldenRod",
+    "LavenderBlush",
+    "SlateGrey",
+    "DeepPink",
+    "Gold",
+    "Cyan",
+    "LightSteelBlue",
+    "MediumPurple",
+    "ForestGreen",
+    "DarkOrange",
+    "Tan",
+    "Salmon",
+    "PaleGoldenRod",
+    "LightGreen",
+    "LightSlateGray",
+    "HoneyDew",
+    "Fuchsia",
+    "LightSeaGreen",
+    "DarkOrchid",
+    "Green",
+    "Chartreuse",
+    "LimeGreen",
+    "AntiqueWhite",
+    "Beige",
+    "Gainsboro",
+    "Bisque",
+    "SaddleBrown",
+    "Silver",
+    "Lavender",
+    "Teal",
+    "LightCyan",
+    "PapayaWhip",
+    "Purple",
+    "Coral",
+    "BurlyWood",
+    "LightGray",
+    "Snow",
+    "MistyRose",
+    "PowderBlue",
+    "DarkCyan",
+    "White",
+    "Turquoise",
+    "MediumSlateBlue",
+    "PeachPuff",
+    "Moccasin",
+    "LightSalmon",
+    "SkyBlue",
+    "Khaki",
+    "MediumSpringGreen",
+    "BlueViolet",
+    "MintCream",
+    "Linen",
+    "SeaGreen",
+    "HotPink",
+    "LightYellow",
+    "BlanchedAlmond",
+    "RoyalBlue",
+    "RosyBrown",
+    "MediumOrchid",
+    "DarkTurquoise",
+    "LightGoldenRodYellow",
+    "LightSkyBlue",
+]
+
+
+def visualize_detections(image_path, output_path, detections, labels=[]):
+    image = Image.open(image_path).convert(mode="RGB")
+    draw = ImageDraw.Draw(image)
+    line_width = 2
+    font = ImageFont.load_default()
+    for d in detections:
+        color = COLORS[d["class"] % len(COLORS)]
+        draw.line(
+            [
+                (d["xmin"], d["ymin"]),
+                (d["xmin"], d["ymax"]),
+                (d["xmax"], d["ymax"]),
+                (d["xmax"], d["ymin"]),
+                (d["xmin"], d["ymin"]),
+            ],
+            width=line_width,
+            fill=color,
+        )
+        label = f"Class {d['class']}"
+        if d["class"] < len(labels):
+            label = f"{labels[d['class']]}"
+        score = d["score"]
+        text = f"{label}: {int(100*score)}%"
+        if score < 0:
+            text = label
+        left, top, right, bottom = font.getbbox(text)
+        text_width, text_height = right - left, bottom - top
+        text_bottom = max(text_height, d["ymin"])
+        text_left = d["xmin"]
+        margin = np.ceil(0.05 * text_height)
+        draw.rectangle(
+            [
+                (text_left, text_bottom - text_height - 2 * margin),
+                (text_left + text_width, text_bottom),
+            ],
+            fill=color,
+        )
+        draw.text(
+            (text_left + margin, text_bottom - text_height - margin),
+            text,
+            fill="black",
+            font=font,
+        )
+    if output_path is None:
+        return image
+    image.save(output_path)

samples/python/detectron2/README.md

@@ -0,0 +1,191 @@
+# Detectron 2 Mask R-CNN R50-FPN 3x in TensorRT
+
+Support for Detectron 2 Mask R-CNN R50-FPN 3x model in TensorRT. This script helps with converting, running and validating this model with TensorRT.
+
+## Changelog
+
+- October 2025
+  - Migrate to strongly typed APIs.
+- Aug 2025
+  - Removed support for Python versions < 3.10.
+- Aug 2023
+  - Update ONNX version support to 1.14.0
+  - Update ONNX Runtime version support to 1.15.1 for Python>=3.8
+  - Removed support for Python versions < 3.8.
+- July 2023:
+  - Update benchmarks and include hardware used.
+- October 2022:
+  - Updated converter to support `tracing` export instead of deprecated `caffe2_tracing`
+
+## Setup
+
+In order for scripts to work we suggest an environment with TensorRT >= 8.4.1.
+
+Install TensorRT as per the [TensorRT Install Guide](https://docs.nvidia.com/deeplearning/tensorrt/install-guide/index.html). You will need to make sure the Python bindings for TensorRT are also installed correctly, these are available by installing the `python3-libnvinfer` and `python3-libnvinfer-dev` packages on your TensorRT download.
+
+Install all dependencies listed in `requirements.txt`:
+
+```
+pip install -r requirements.txt
+```
+Note: this sample cannot be run on Jetson platforms as `torch.distributed` is not available. To check whether your platform supports `torch.distributed`, open a Python shell and confirm that `torch.distributed.is_available()` returns `True`.
+
+## Model Conversion
+
+The workflow to convert Detectron 2 Mask R-CNN R50-FPN 3x model is basically Detectron 2 → ONNX → TensorRT, and so parts of this process require Detectron 2 to be installed. Official export to ONNX is documented [here](https://detectron2.readthedocs.io/en/latest/tutorials/deployment.html).
+
+### Detectron 2 Deployment
+Deployment is done through export model script located in `detectron2/tools/deploy/export_model.py` of Detectron 2 [github](https://github.com/facebookresearch/detectron2). Detectron 2 Mask R-CNN R50-FPN 3x model is dynamic with minimum testing dimension size of 800 and maximum of 1333. TensorRT plug-ins used for conversion of this model do not support dynamic shapes, as a result we have to set both height and width of the input tensor to 1344. 1344 instead of 1333 because model requires both height and width of the input tensor to be divisible by 32. In order to export this model with correct 1344x1344 resolution, we have to make a change to `export_model.py`. Currently lines 160-162:
+
+```
+aug = T.ResizeShortestEdge(
+    [cfg.INPUT.MIN_SIZE_TEST, cfg.INPUT.MIN_SIZE_TEST], cfg.INPUT.MAX_SIZE_TEST
+)
+```
+have to be changed to:
+
+```
+aug = T.ResizeShortestEdge(
+    [1344, 1344], 1344
+)
+```
+
+Export script takes `--sample-image` as one of the arguments. Such image is used to adjust input dimensions and dimensions of tensors for the rest of the network. This sample image has to be an image of 1344x1344 dimensions, which contains at least one detectable by model object. My recommendation is to upsample one of COCO dataset images to 1344x1344. Sample command:
+
+```
+python detectron2/tools/deploy/export_model.py \
+    --sample-image 1344x1344.jpg \
+    --config-file detectron2/configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml \
+    --export-method tracing \
+    --format onnx \
+    --output ./ \
+    MODEL.WEIGHTS path/to/model_final_f10217.pkl \
+    MODEL.DEVICE cuda
+
+```
+
+Where `--sample-image` is 1344x1344 image; `--config-file` path to Mask R-CNN R50-FPN 3x config, included with detectron2; `MODEL.WEIGHTS` are weights of Mask R-CNN R50-FPN 3x that can be downloaded [here](https://github.com/facebookresearch/detectron2/blob/main/MODEL_ZOO.md). Resulted `model.onnx` will be an input to conversion script.
+
+### Create ONNX Graph
+This is supported Detectron 2 model:
+
+| **Model**                                     | **Resolution** |
+| ----------------------------------------------|----------------|
+| Mask R-CNN R50-FPN 3x                         | 1344x1344      |
+
+If Detectron 2 Mask R-CNN is ready to be converted (i.e. you ran `detectron2/tools/deploy/export_model.py`), run:
+
+```
+python create_onnx.py \
+    --exported_onnx /path/to/model.onnx \
+    --onnx /path/to/converted.onnx \
+    --det2_config /detectron2/configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml \
+    --det2_weights /model_final_f10217.pkl \
+    --sample_image any_image.jpg
+```
+
+This will create the file `converted.onnx` which is ready to convert to TensorRT.
+
+It is important to mention that `--sample_image` in this case is used for anchor generation. Detectron 2 ONNX models do not have anchor data inside the graph, so anchors have to be generated "offline". If custom model is used, make sure preprocessing of your model matches what is coded in `get_anchors(self, sample_image)` function.
+
+The script has a few optional arguments, including:
+
+* `--first_nms_threshold [...]` allows overriding the default 1st NMS score threshold parameter, as the runtime latency of the NMS plugin is sensitive to this value. It's a good practice to set this value as high as possible, while still fulfilling your application requirements, to reduce inference latency. In Mask R-CNN this will be a score threshold for Region Proposal Network.
+* `--second_nms_threshold [...]` allows overriding the default 2nd NMS score threshold parameter, further improves the runtime latency of the NMS plugin. It's a good practice to set this value as high as possible, while still fulfilling your application requirements, to reduce inference latency. It will be the second and last NMS.
+* `--batch_size` allows selection of various batch sizes, default is 1.
+
+
+Optionally, you may wish to visualize the resulting ONNX graph with a tool such as [Netron](https://netron.app/).
+
+The input to the graph is a `float32` tensor with the selected input shape, containing RGB pixel data in the range of 0 to 255. All preprocessing will be performed inside the Model graph, so it is not required to further pre-process the input data.
+
+
+The outputs of the graph are the same as the outputs of the [EfficientNMS_TRT](https://github.com/NVIDIA/TensorRT/tree/master/plugin/efficientNMSPlugin) plugin and segmentation head output, name of the last node is `detection_masks`, shape is `[batch_size, max_proposals, mask_height, mask_width]`, dtype is float32.
+
+### Build TensorRT Engine
+
+TensorRT engine can be built directly with `trtexec` using the ONNX graph generated in the previous step. If it's not already in your `$PATH`, the `trtexec` binary is usually found in `/usr/bin/trtexec`, depending on your TensorRT installation method. Run:
+
+```
+trtexec --onnx=/path/to/converted.onnx --saveEngine=/path/to/engine.trt --stronglyTyped
+```
+
+However, the script `build_engine.py` is also provided in this repository for convenience, as it has been tailored to Detectron 2 2 Mask R-CNN R50-FPN 3x engine building. Run `python3 build_engine.py --help` for details on available settings.
+
+#### Benchmark Engine
+
+Optionally, you can obtain execution timing information for the built engine by using the `trtexec` utility, as:
+
+```
+trtexec \
+    --loadEngine=/path/to/engine.trt \
+    --iterations=100 --avgRuns=100
+```
+
+An inference benchmark will run, with GPU Compute latency times printed out to the console. Depending on your environment, you should see something similar to:
+
+```
+GPU Compute Time: min = 30.1864 ms, max = 37.0945 ms, mean = 34.481 ms, median = 34.4187 ms, percentile(99%) = 37.0945 ms
+```
+
+Sample results with fp32 data precision are shown below. The following results were obtained using an RTX A5000 and TensorRT 8.6.1. mAP was evaluated for the COCO val2017 dataset using the instructions in [Evaluate mAP Metric](#evaluate-map-metric).
+
+| **Precision**   | **Latency** | **bbox COCO mAP** | **segm COCO mAP** |
+| ----------------|-------------|-------------------|-------------------|
+| fp32            | 25.89 ms    | 0.402             | 0.368             |
+
+## Inference
+
+For optimal performance, inference should be done in a C++ application that takes advantage of CUDA Graphs to launch the inference request. Alternatively, the TensorRT engine built with this process can also be executed through either [Triton Inference Server](https://developer.nvidia.com/nvidia-triton-inference-server) or [DeepStream SDK](https://developer.nvidia.com/deepstream-sdk).
+
+However, for convenience, a python inference script is provided here for quick testing of the built TensorRT engine.
+
+### Inference in Python
+
+To perform object detection on a set of images with TensorRT, run:
+
+```
+python infer.py \
+    --engine /path/to/engine.trt \
+    --input /path/to/images \
+    --det2_config /detectron2/configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml \
+    --output /path/to/output \
+```
+
+Where the input path can be either a single image file, or a directory of jpg/png/bmp images.
+
+The script has a few optional arguments, including:
+* `--nms_threshold` allows overriding the default second NMS score threshold parameter.
+* `--iou_threshold` allows to set IoU threshold for the mask segmentation, default is 0.5.
+
+The detection results will be written out to the specified output directory, consisting of a visualization image, and a tab-separated results file for each input image processed.
+
+#### Sample Images
+
+![infer_1](https://drive.google.com/uc?export=view&id=1AOW9IXqjrU7eVYmaue-pqijNucXmx_s0)
+
+![infer_2](https://drive.google.com/uc?export=view&id=1m1fp2v41DOqKfj423G0-eyKVurrPNYGx)
+
+### Evaluate mAP Metric
+
+Given a validation dataset (such as [COCO val2017 data](http://images.cocodataset.org/zips/val2017.zip)), you can get the mAP metrics for the built TensorRT engine. This will use the mAP metrics tools functions from the [Detectron 2 evaluation](https://github.com/facebookresearch/detectron2/tree/main/detectron2/evaluation) repository. Make sure you follow [Use Builtin Datasets guide](https://detectron2.readthedocs.io/en/latest/tutorials/builtin_datasets.html) to correctly setup COCO or custom dataset. Additionally, run `eval_coco.py` in the same folder where `/datasets` is present, otherwise this error will appear:
+
+```
+FileNotFoundError: [Errno 2] No such file or directory: 'datasets/coco/annotations/instances_val2017.json'
+```
+
+To run evalutions, run:
+
+```
+python eval_coco.py \
+    --engine /path/to/engine.trt \
+    --input /path/to/coco/val2017 \
+    --det2_config /detectron2/configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml \
+    --det2_weights /model_final_f10217.pkl
+```
+
+The script has a few optional arguments, including:
+* `--nms_threshold` allows overriding the default second NMS score threshold parameter.
+* `--iou_threshold` allows to set IoU threshold for the mask segmentation, default is 0.5.
+
+The mAP metric is sensitive to the NMS score threshold used, as using a high threshold will reduce the model recall, resulting in a lower mAP value. It may be a good idea to build separate TensorRT engines for different purposes. That is, one engine with a default threshold (like 0.5) dedicated for mAP validation, and another engine with your application specific threshold (like 0.8) for deployment. This is why we keep the NMS threshold as a configurable parameter in the `create_onnx.py` script.

samples/python/detectron2/build_engine.py

@@ -0,0 +1,147 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2025 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 sys
+import logging
+import argparse
+
+import numpy as np
+import tensorrt as trt
+from cuda.bindings import driver as cuda, runtime as cudart
+
+from image_batcher import ImageBatcher
+
+logging.basicConfig(level=logging.INFO)
+logging.getLogger("EngineBuilder").setLevel(logging.INFO)
+log = logging.getLogger("EngineBuilder")
+
+sys.path.insert(1, os.path.join(os.path.dirname(os.path.realpath(__file__)), os.pardir))
+import common
+
+class EngineBuilder:
+    """
+    Parses an ONNX graph and builds a TensorRT engine from it.
+    """
+
+    def __init__(self, verbose=False, workspace=8):
+        """
+        :param verbose: If enabled, a higher verbosity level will be set on the TensorRT logger.
+        :param workspace: Max memory workspace to allow, in Gb.
+        """
+        self.trt_logger = trt.Logger(trt.Logger.INFO)
+        if verbose:
+            self.trt_logger.min_severity = trt.Logger.Severity.VERBOSE
+
+        trt.init_libnvinfer_plugins(self.trt_logger, namespace="")
+
+        self.builder = trt.Builder(self.trt_logger)
+        self.config = self.builder.create_builder_config()
+        self.config.set_memory_pool_limit(
+            trt.MemoryPoolType.WORKSPACE, workspace * (2**30)
+        )
+
+        self.batch_size = None
+        self.network = None
+        self.parser = None
+
+    def create_network(self, onnx_path):
+        """
+        Parse the ONNX graph and create the corresponding TensorRT network definition.
+        :param onnx_path: The path to the ONNX graph to load.
+        """
+        self.network = self.builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED))
+        self.parser = trt.OnnxParser(self.network, self.trt_logger)
+
+        onnx_path = os.path.realpath(onnx_path)
+        with open(onnx_path, "rb") as f:
+            if not self.parser.parse(f.read()):
+                log.error("Failed to load ONNX file: {}".format(onnx_path))
+                for error in range(self.parser.num_errors):
+                    log.error(self.parser.get_error(error))
+                sys.exit(1)
+
+        inputs = [self.network.get_input(i) for i in range(self.network.num_inputs)]
+        outputs = [self.network.get_output(i) for i in range(self.network.num_outputs)]
+
+        log.info("Network Description")
+        for input in inputs:
+            self.batch_size = input.shape[0]
+            log.info(
+                "Input '{}' with shape {} and dtype {}".format(
+                    input.name, input.shape, input.dtype
+                )
+            )
+        for output in outputs:
+            log.info(
+                "Output '{}' with shape {} and dtype {}".format(
+                    output.name, output.shape, output.dtype
+                )
+            )
+        assert self.batch_size > 0
+
+    def create_engine(
+        self,
+        engine_path,
+    ):
+        """
+        Build the TensorRT engine and serialize it to disk.
+        :param engine_path: The path where to serialize the engine to.
+        """
+        engine_path = os.path.realpath(engine_path)
+        engine_dir = os.path.dirname(engine_path)
+        os.makedirs(engine_dir, exist_ok=True)
+        log.info("Building fp32 Engine in {}".format(engine_path))
+
+        engine_bytes = self.builder.build_serialized_network(self.network, self.config)
+        if engine_bytes is None:
+            log.error("Failed to create engine")
+            sys.exit(1)
+
+        with open(engine_path, "wb") as f:
+            log.info("Serializing engine to file: {:}".format(engine_path))
+            f.write(engine_bytes)
+
+
+def main(args):
+    builder = EngineBuilder(args.verbose, args.workspace)
+    builder.create_network(args.onnx)
+    builder.create_engine(
+        args.engine,
+    )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-o", "--onnx", help="The input ONNX model file to load")
+    parser.add_argument("-e", "--engine", help="The output path for the TRT engine")
+    parser.add_argument(
+        "-v", "--verbose", action="store_true", help="Enable more verbose log output"
+    )
+    parser.add_argument(
+        "-w",
+        "--workspace",
+        default=1,
+        type=int,
+        help="The max memory workspace size to allow in Gb, " "default: 1",
+    )
+    args = parser.parse_args()
+    if not all([args.onnx, args.engine]):
+        parser.print_help()
+        log.error("These arguments are required: --onnx and --engine")
+        sys.exit(1)
+    main(args)

samples/python/detectron2/create_onnx.py

@@ -0,0 +1,879 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2026 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 re
+import sys
+import argparse
+import logging
+import cv2
+import onnx_graphsurgeon as gs
+import numpy as np
+import onnx
+from onnx import shape_inference
+import torch
+
+try:
+    from detectron2.engine.defaults import DefaultPredictor
+    from detectron2.modeling import build_model
+    from detectron2.config import get_cfg
+    from detectron2.structures import ImageList
+except ImportError:
+    print("Could not import Detectron 2 modules. Maybe you did not install Detectron 2")
+    print(
+        "Please install Detectron 2, check https://github.com/facebookresearch/detectron2/blob/main/INSTALL.md"
+    )
+    sys.exit(1)
+
+import onnx_utils
+
+logging.basicConfig(level=logging.INFO)
+logging.getLogger("ModelHelper").setLevel(logging.INFO)
+log = logging.getLogger("ModelHelper")
+
+
+class DET2GraphSurgeon:
+    def __init__(self, saved_model_path, config_file, weights):
+        """
+        Constructor of the Model Graph Surgeon object, to do the conversion of a Detectron 2 Mask R-CNN exported model
+        to an ONNX-TensorRT parsable model.
+        :param saved_model_path: The path pointing to the exported Detectron 2 Mask R-CNN ONNX model.
+        :param config_file: The path pointing to the Detectron 2 yaml file which describes the model.
+        :param config_file: Weights to load for the Detectron 2 model.
+        """
+
+        def det2_setup(config_file, weights):
+            """
+            Create configs and perform basic setups.
+            """
+            cfg = get_cfg()
+            cfg.merge_from_file(config_file)
+            cfg.merge_from_list(["MODEL.WEIGHTS", weights])
+            cfg.freeze()
+            return cfg
+
+        # Import exported Detectron 2 Mask R-CNN ONNX model as GraphSurgeon object.
+        self.graph = gs.import_onnx(onnx.load(saved_model_path))
+        assert self.graph
+        log.info("ONNX graph loaded successfully")
+
+        # Fold constants via ONNX-GS that exported script might've missed.
+        self.graph.fold_constants()
+
+        # Set up Detectron 2 model configuration.
+        self.det2_cfg = det2_setup(config_file, weights)
+
+        # Getting model characteristics.
+        self.fpn_out_channels = self.det2_cfg.MODEL.FPN.OUT_CHANNELS
+        self.num_classes = self.det2_cfg.MODEL.ROI_HEADS.NUM_CLASSES
+        self.first_NMS_max_proposals = self.det2_cfg.MODEL.RPN.POST_NMS_TOPK_TEST
+        self.first_NMS_iou_threshold = self.det2_cfg.MODEL.RPN.NMS_THRESH
+        self.first_NMS_score_threshold = 0.01
+        self.first_ROIAlign_pooled_size = (
+            self.det2_cfg.MODEL.ROI_BOX_HEAD.POOLER_RESOLUTION
+        )
+        self.first_ROIAlign_sampling_ratio = (
+            self.det2_cfg.MODEL.ROI_BOX_HEAD.POOLER_SAMPLING_RATIO
+        )
+        self.first_ROIAlign_type = self.det2_cfg.MODEL.ROI_BOX_HEAD.POOLER_TYPE
+        self.second_NMS_max_proposals = self.det2_cfg.TEST.DETECTIONS_PER_IMAGE
+        self.second_NMS_iou_threshold = self.det2_cfg.MODEL.ROI_HEADS.NMS_THRESH_TEST
+        self.second_NMS_score_threshold = (
+            self.det2_cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST
+        )
+        self.second_ROIAlign_pooled_size = (
+            self.det2_cfg.MODEL.ROI_MASK_HEAD.POOLER_RESOLUTION
+        )
+        self.second_ROIAlign_sampling_ratio = (
+            self.det2_cfg.MODEL.ROI_MASK_HEAD.POOLER_SAMPLING_RATIO
+        )
+        self.second_ROIAlign_type = self.det2_cfg.MODEL.ROI_MASK_HEAD.POOLER_TYPE
+        self.mask_out_res = 28
+
+        # Model characteristics.
+        log.info("Number of FPN output channels is {}".format(self.fpn_out_channels))
+        log.info("Number of classes is {}".format(self.num_classes))
+        log.info("First NMS max proposals is {}".format(self.first_NMS_max_proposals))
+        log.info("First NMS iou threshold is {}".format(self.first_NMS_iou_threshold))
+        log.info(
+            "First NMS score threshold is {}".format(self.first_NMS_score_threshold)
+        )
+        log.info("First ROIAlign type is {}".format(self.first_ROIAlign_type))
+        log.info(
+            "First ROIAlign pooled size is {}".format(self.first_ROIAlign_pooled_size)
+        )
+        log.info(
+            "First ROIAlign sampling ratio is {}".format(
+                self.first_ROIAlign_sampling_ratio
+            )
+        )
+        log.info("Second NMS max proposals is {}".format(self.second_NMS_max_proposals))
+        log.info("Second NMS iou threshold is {}".format(self.second_NMS_iou_threshold))
+        log.info(
+            "Second NMS score threshold is {}".format(self.second_NMS_score_threshold)
+        )
+        log.info("Second ROIAlign type is {}".format(self.second_ROIAlign_type))
+        log.info(
+            "Second ROIAlign pooled size is {}".format(self.second_ROIAlign_pooled_size)
+        )
+        log.info(
+            "Second ROIAlign sampling ratio is {}".format(
+                self.second_ROIAlign_sampling_ratio
+            )
+        )
+        log.info(
+            "Individual mask output resolution is {}x{}".format(
+                self.mask_out_res, self.mask_out_res
+            )
+        )
+
+        self.batch_size = None
+
+    def sanitize(self):
+        """
+        Sanitize the graph by cleaning any unconnected nodes, do a topological resort, and fold constant inputs values.
+        When possible, run shape inference on the ONNX graph to determine tensor shapes.
+        """
+
+        for i in range(3):
+            count_before = len(self.graph.nodes)
+            self.graph.cleanup().toposort()
+            try:
+                for node in self.graph.nodes:
+                    for o in node.outputs:
+                        o.shape = None
+                model = gs.export_onnx(self.graph)
+                model = shape_inference.infer_shapes(model)
+                self.graph = gs.import_onnx(model)
+            except Exception as e:
+                log.info(
+                    "Shape inference could not be performed at this time:\n{}".format(e)
+                )
+            try:
+                self.graph.fold_constants(fold_shapes=True)
+            except TypeError as e:
+                log.error(
+                    "This version of ONNX GraphSurgeon does not support folding shapes, please upgrade your "
+                    "onnx_graphsurgeon module. Error:\n{}".format(e)
+                )
+                raise
+
+            count_after = len(self.graph.nodes)
+            if count_before == count_after:
+                # No new folding occurred in this iteration, so we can stop for now.
+                break
+
+    def get_anchors(self, sample_image):
+        """
+        Detectron 2 exported ONNX does not contain anchors required for efficientNMS plug-in, so they must be generated
+        "offline" by calling actual Detectron 2 model and getting anchors from it.
+        :param sample_image: Sample image required to run through the model and obtain anchors.
+        Can be any image from a dataset. Make sure listed here Detectron 2 preprocessing steps
+        actually match your preprocessing steps. Otherwise, behavior can be unpredictable.
+        Additionally, anchors have to be generated for a fixed input dimensions,
+        meaning as soon as image leaves a preprocessor and enters predictor.model.backbone() it must have
+        a fixed dimension (1344x1344 in my case) that every single image in dataset must follow, since currently
+        TensorRT plug-ins do not support dynamic shapes.
+        """
+        # Get Detectron 2 model config and build it.
+        predictor = DefaultPredictor(self.det2_cfg)
+        model = build_model(self.det2_cfg)
+
+        # Image preprocessing.
+        input_im = cv2.imread(sample_image)
+        raw_height, raw_width = input_im.shape[:2]
+        image = predictor.aug.get_transform(input_im).apply_image(input_im)
+        image = torch.as_tensor(image.astype("float32").transpose(2, 0, 1))
+
+        # Model preprocessing.
+        inputs = [{"image": image, "height": raw_height, "width": raw_width}]
+        images = [x["image"].to(model.device) for x in inputs]
+        images = [(x - model.pixel_mean) / model.pixel_std for x in images]
+        imagelist_images = ImageList.from_tensors(images, 1344)
+
+        # Get feature maps from backbone.
+        features = predictor.model.backbone(imagelist_images.tensor)
+
+        # Get proposals from Region Proposal Network and obtain anchors from anchor generator.
+        features = [features[f] for f in predictor.model.proposal_generator.in_features]
+        det2_anchors = predictor.model.proposal_generator.anchor_generator(features)
+
+        # Extract anchors based on feature maps in ascending order (P2->P6).
+        p2_anchors = det2_anchors[0].tensor.detach().cpu().numpy()
+        p3_anchors = det2_anchors[1].tensor.detach().cpu().numpy()
+        p4_anchors = det2_anchors[2].tensor.detach().cpu().numpy()
+        p5_anchors = det2_anchors[3].tensor.detach().cpu().numpy()
+        p6_anchors = det2_anchors[4].tensor.detach().cpu().numpy()
+        final_anchors = np.concatenate(
+            (p2_anchors, p3_anchors, p4_anchors, p5_anchors, p6_anchors)
+        )
+
+        return final_anchors
+
+    def save(self, output_path):
+        """
+        Save the ONNX model to the given location.
+        :param output_path: Path pointing to the location where to write out the updated ONNX model.
+        """
+        self.graph.cleanup().toposort()
+        model = gs.export_onnx(self.graph)
+        output_path = os.path.realpath(output_path)
+        os.makedirs(os.path.dirname(output_path), exist_ok=True)
+        onnx.save(model, output_path)
+        log.info("Saved ONNX model to {}".format(output_path))
+
+    def update_preprocessor(self, batch_size):
+        """
+        Remove all the pre-processing nodes in the ONNX graph and leave only the image normalization essentials.
+        :param batch_size: The batch size to use for the ONNX graph.
+        """
+        # Set graph inputs.
+        self.batch_size = batch_size
+        self.height = self.graph.inputs[0].shape[1]
+        self.width = self.graph.inputs[0].shape[2]
+
+        input_shape = [self.batch_size, 3, self.height, self.width]
+        self.graph.inputs[0].shape = input_shape
+        self.graph.inputs[0].dtype = np.float32
+        self.graph.inputs[0].name = "input_tensor"
+
+        self.sanitize()
+        log.info(
+            "ONNX graph input shape: {} [NCHW format set]".format(
+                self.graph.inputs[0].shape
+            )
+        )
+
+        # Find the initial nodes of the graph, whatever the input is first connected to, and disconnect them.
+        for node in [
+            node for node in self.graph.nodes if self.graph.inputs[0] in node.inputs
+        ]:
+            node.inputs.clear()
+
+        # Get input tensor.
+        input_tensor = self.graph.inputs[0]
+
+        # Create preprocessing Sub node and connect input tensor to it.
+        sub_const = np.expand_dims(
+            np.asarray([255 * 0.406, 255 * 0.456, 255 * 0.485], dtype=np.float32),
+            axis=(1, 2),
+        )
+        sub_out = self.graph.op_with_const(
+            "Sub", "preprocessor/mean", input_tensor, sub_const
+        )
+
+        # Find first Div node and connect to output of Sub node.
+        div_node = self.graph.find_node_by_op("Div")
+        log.info("Found {} node".format(div_node.op))
+        div_node.inputs[0] = sub_out[0]
+
+        # Find first Conv and connect preprocessor directly to it.
+        conv_node = self.graph.find_node_by_op("Conv")
+        log.info("Found {} node".format(conv_node.op))
+        conv_node.inputs[0] = div_node.outputs[0]
+
+        # Reshape nodes tend to update the batch dimension to a fixed value of 1, they should use the batch size instead.
+        for node in [node for node in self.graph.nodes if node.op == "Reshape"]:
+            if type(node.inputs[1]) == gs.Constant and node.inputs[1].values[0] == 1:
+                node.inputs[1].values[0] = self.batch_size
+
+    def NMS(
+        self,
+        boxes,
+        scores,
+        anchors,
+        background_class,
+        score_activation,
+        max_proposals,
+        iou_threshold,
+        nms_score_threshold,
+        user_threshold,
+        nms_name=None,
+    ):
+        # Helper function to create the NMS Plugin node with the selected inputs.
+        # EfficientNMS_TRT TensorRT Plugin is suitable for our use case.
+        # :param boxes: The box predictions from the Box Net.
+        # :param scores: The class predictions from the Class Net.
+        # :param anchors: The default anchor coordinates.
+        # :param background_class: The label ID for the background class.
+        # :param max_proposals: Number of proposals made by NMS.
+        # :param score_activation: If set to True - apply sigmoid activation to the confidence scores during NMS operation,
+        # if false - no activation.
+        # :param iou_threshold: NMS intersection over union threshold, given by self.det2_cfg.
+        # :param nms_score_threshold: NMS score threshold, given by self.det2_cfg.
+        # :param user_threshold: User's given threshold to overwrite default NMS score threshold.
+        # :param nms_name: Name of NMS node in a graph, renames NMS elements accordingly in order to eliminate cycles.
+
+        if nms_name is None:
+            nms_name = ""
+        else:
+            nms_name = "_" + nms_name
+
+        # Set score threshold.
+        score_threshold = (
+            nms_score_threshold if user_threshold is None else user_threshold
+        )
+
+        # NMS Outputs.
+        nms_output_num_detections = gs.Variable(
+            name="num_detections" + nms_name, dtype=np.int32, shape=[self.batch_size, 1]
+        )
+        nms_output_boxes = gs.Variable(
+            name="detection_boxes" + nms_name,
+            dtype=np.float32,
+            shape=[self.batch_size, max_proposals, 4],
+        )
+        nms_output_scores = gs.Variable(
+            name="detection_scores" + nms_name,
+            dtype=np.float32,
+            shape=[self.batch_size, max_proposals],
+        )
+        nms_output_classes = gs.Variable(
+            name="detection_classes" + nms_name,
+            dtype=np.int32,
+            shape=[self.batch_size, max_proposals],
+        )
+
+        nms_outputs = [
+            nms_output_num_detections,
+            nms_output_boxes,
+            nms_output_scores,
+            nms_output_classes,
+        ]
+
+        # Plugin.
+        self.graph.plugin(
+            op="EfficientNMS_TRT",
+            name="nms" + nms_name,
+            inputs=[boxes, scores, anchors],
+            outputs=nms_outputs,
+            attrs={
+                "plugin_version": "1",
+                "background_class": background_class,
+                "max_output_boxes": max_proposals,
+                "score_threshold": max(0.01, score_threshold),
+                "iou_threshold": iou_threshold,
+                "score_activation": score_activation,
+                "class_agnostic": False,
+                "box_coding": 1,
+            },
+        )
+        log.info("Created nms{} with EfficientNMS_TRT plugin".format(nms_name))
+
+        return nms_outputs
+
+    def ROIAlign(
+        self,
+        rois,
+        p2,
+        p3,
+        p4,
+        p5,
+        pooled_size,
+        sampling_ratio,
+        roi_align_type,
+        num_rois,
+        ra_name,
+    ):
+        # Helper function to create the ROIAlign Plugin node with the selected inputs.
+        # PyramidROIAlign_TRT TensorRT Plugin is suitable for our use case.
+        # :param rois: Regions of interest/detection boxes outputs from preceding NMS node.
+        # :param p2: Output of p2 feature map.
+        # :param p3: Output of p3 feature map.
+        # :param p4: Output of p4 feature map.
+        # :param p5: Output of p5 feature map.
+        # :param pooled_size: Pooled output dimensions.
+        # :param sampling_ratio: Number of sampling points in the interpolation grid used to compute the output value of each pooled output bin.
+        # :param roi_align_type: Type of Detectron 2 ROIAlign op, either ROIAlign (vanilla) or ROIAlignV2 (0.5 coordinate offset).
+        # :param num_rois: Number of ROIs resulting from ROIAlign operation.
+        # :param ra_name: Name of ROIAlign node in a graph, renames ROIAlign elements accordingly in order to eliminate cycles.
+
+        # Different types of Detectron 2's ROIAlign ops require coordinate offset that is supported by PyramidROIAlign_TRT.
+        if roi_align_type == "ROIAlignV2":
+            roi_coords_transform = 2
+        elif roi_align_type == "ROIAlign":
+            roi_coords_transform = 0
+        else:
+            raise ValueError(
+                f"Unsupported ROIAlign type '{roi_align_type}'. "
+                f"Expected 'ROIAlignV2' or 'ROIAlign'."
+            )
+
+        # ROIAlign outputs.
+        roi_align_output = gs.Variable(
+            name="roi_align/output_" + ra_name,
+            dtype=np.float32,
+            shape=[
+                self.batch_size,
+                num_rois,
+                self.fpn_out_channels,
+                pooled_size,
+                pooled_size,
+            ],
+        )
+
+        # Plugin.
+        self.graph.plugin(
+            op="PyramidROIAlign_TRT",
+            name="roi_align_" + ra_name,
+            inputs=[rois, p2, p3, p4, p5],
+            outputs=[roi_align_output],
+            attrs={
+                "plugin_version": "1",
+                "fpn_scale": 224,
+                "pooled_size": pooled_size,
+                "image_size": [self.height, self.width],
+                "roi_coords_absolute": 0,
+                "roi_coords_swap": 0,
+                "roi_coords_transform": roi_coords_transform,
+                "sampling_ratio": sampling_ratio,
+            },
+        )
+        log.info("Created {} with PyramidROIAlign_TRT plugin".format(ra_name))
+
+        return roi_align_output
+
+    def process_graph(
+        self, anchors, first_nms_threshold=None, second_nms_threshold=None
+    ):
+        """
+        Processes the graph to replace the GenerateProposals and BoxWithNMSLimit operations with EfficientNMS_TRT
+        TensorRT plugin nodes and ROIAlign operations with PyramidROIAlign_TRT plugin nodes.
+        :param anchors: Anchors generated from sample image "offline" by Detectron 2, since anchors are not provided
+        inside the graph.
+        :param first_nms_threshold: Override the 1st NMS score threshold value. If set to None, use the value in the graph.
+        :param second_nms_threshold: Override the 2nd NMS score threshold value. If set to None, use the value in the graph.
+        """
+
+        def backbone():
+            """
+            Updates the graph to replace all ResizeNearest ops with ResizeNearest plugins in backbone.
+            """
+            # Get final backbone outputs.
+            p2 = self.graph.find_node_by_op_name("Conv", "/backbone/fpn_output2/Conv")
+            p3 = self.graph.find_node_by_op_name("Conv", "/backbone/fpn_output3/Conv")
+            p4 = self.graph.find_node_by_op_name("Conv", "/backbone/fpn_output4/Conv")
+            p5 = self.graph.find_node_by_op_name("Conv", "/backbone/fpn_output5/Conv")
+
+            return p2.outputs[0], p3.outputs[0], p4.outputs[0], p5.outputs[0]
+
+        def proposal_generator(anchors, first_nms_threshold):
+            """
+            Updates the graph to replace all GenerateProposals Caffe ops with one single NMS for proposals generation.
+            :param anchors: Anchors generated from sample image "offline" by Detectron 2, since anchors are not provided
+            inside the graph
+            :param first_nms_threshold: Override the 1st NMS score threshold value. If set to None, use the value in the graph.
+            """
+            # Get nodes containing final objectness logits.
+            p2_logits = self.graph.find_node_by_op_name(
+                "Flatten", "/proposal_generator/Flatten"
+            )
+            p3_logits = self.graph.find_node_by_op_name(
+                "Flatten", "/proposal_generator/Flatten_1"
+            )
+            p4_logits = self.graph.find_node_by_op_name(
+                "Flatten", "/proposal_generator/Flatten_2"
+            )
+            p5_logits = self.graph.find_node_by_op_name(
+                "Flatten", "/proposal_generator/Flatten_3"
+            )
+            p6_logits = self.graph.find_node_by_op_name(
+                "Flatten", "/proposal_generator/Flatten_4"
+            )
+
+            # Get nodes containing final anchor_deltas.
+            p2_anchors = self.graph.find_node_by_op_name(
+                "Reshape", "/proposal_generator/Reshape_1"
+            )
+            p3_anchors = self.graph.find_node_by_op_name(
+                "Reshape", "/proposal_generator/Reshape_3"
+            )
+            p4_anchors = self.graph.find_node_by_op_name(
+                "Reshape", "/proposal_generator/Reshape_5"
+            )
+            p5_anchors = self.graph.find_node_by_op_name(
+                "Reshape", "/proposal_generator/Reshape_7"
+            )
+            p6_anchors = self.graph.find_node_by_op_name(
+                "Reshape", "/proposal_generator/Reshape_9"
+            )
+
+            # Concatenate all objectness logits/scores data.
+            scores_inputs = [
+                p2_logits.outputs[0],
+                p3_logits.outputs[0],
+                p4_logits.outputs[0],
+                p5_logits.outputs[0],
+                p6_logits.outputs[0],
+            ]
+            scores_tensor = self.graph.layer(
+                name="scores",
+                op="Concat",
+                inputs=scores_inputs,
+                outputs=["scores"],
+                attrs={"axis": 1},
+            )[0]
+            # Unsqueeze to add 3rd dimension of 1 to match tensor dimensions of boxes tensor.
+            scores = self.graph.unsqueeze("scores_unsqueeze", scores_tensor, [2])[0]
+
+            # Concatenate all boxes/anchor_delta data.
+            boxes_inputs = [
+                p2_anchors.outputs[0],
+                p3_anchors.outputs[0],
+                p4_anchors.outputs[0],
+                p5_anchors.outputs[0],
+                p6_anchors.outputs[0],
+            ]
+            boxes = self.graph.layer(
+                name="boxes",
+                op="Concat",
+                inputs=boxes_inputs,
+                outputs=["anchors"],
+                attrs={"axis": 1},
+            )[0]
+
+            # Convert the anchors from Corners to CenterSize encoding.
+            anchors = np.matmul(
+                anchors,
+                [[0.5, 0, -1, 0], [0, 0.5, 0, -1], [0.5, 0, 1, 0], [0, 0.5, 0, 1]],
+            )
+            anchors = anchors / [
+                self.width,
+                self.height,
+                self.width,
+                self.height,
+            ]  # Normalize anchors to [0-1] range
+            anchors = np.expand_dims(anchors, axis=0)
+            anchors = anchors.astype(np.float32)
+            anchors = gs.Constant(name="default_anchors", values=anchors)
+
+            # Create NMS node.
+            nms_outputs = self.NMS(
+                boxes,
+                scores,
+                anchors,
+                -1,
+                False,
+                self.first_NMS_max_proposals,
+                self.first_NMS_iou_threshold,
+                self.first_NMS_score_threshold,
+                first_nms_threshold,
+                "rpn",
+            )
+
+            return nms_outputs
+
+        def roi_heads(rpn_outputs, p2, p3, p4, p5, second_nms_threshold):
+            """
+            Updates the graph to replace all ROIAlign Caffe ops with one single pyramid ROIAlign. Eliminates CollectRpnProposals
+            DistributeFpnProposals and BatchPermutation nodes that are not supported by TensorRT. Connects pyramid ROIAlign to box_head
+            and connects box_head to final box head outputs in a form of second NMS. In order to implement mask head outputs,
+            similar steps as in box_pooler are performed to replace mask_pooler. Finally, reimplemented mask_pooler is connected to
+            mask_head and mask head outputs are produced.
+            :param rpn_outputs: Outputs of the first NMS/proposal generator.
+            :param p2: Output of p2 feature map, required for ROIAlign operation.
+            :param p3: Output of p3 feature map, required for ROIAlign operation.
+            :param p4: Output of p4 feature map, required for ROIAlign operation.
+            :param p5: Output of p5 feature map, required for ROIAlign operation.
+            :param second_nms_threshold: Override the 2nd NMS score threshold value. If set to None, use the value in the graph.
+            """
+            # Create ROIAlign node.
+            box_pooler_output = self.ROIAlign(
+                rpn_outputs[1],
+                p2,
+                p3,
+                p4,
+                p5,
+                self.first_ROIAlign_pooled_size,
+                self.first_ROIAlign_sampling_ratio,
+                self.first_ROIAlign_type,
+                self.first_NMS_max_proposals,
+                "box_pooler",
+            )
+
+            # Reshape node that prepares ROIAlign/box pooler output for Gemm node that comes next.
+            box_pooler_shape = np.asarray(
+                [
+                    -1,
+                    self.fpn_out_channels
+                    * self.first_ROIAlign_pooled_size
+                    * self.first_ROIAlign_pooled_size,
+                ],
+                dtype=np.int64,
+            )
+            box_pooler_reshape = self.graph.op_with_const(
+                "Reshape", "box_pooler/reshape", box_pooler_output, box_pooler_shape
+            )
+
+            # Get first Gemm op of box head and connect box pooler to it.
+            first_box_head_gemm = self.graph.find_node_by_op_name(
+                "Gemm", "/roi_heads/box_head/fc1/Gemm"
+            )
+            first_box_head_gemm.inputs[0] = box_pooler_reshape[0]
+
+            # Get final two nodes of box predictor. Softmax op for cls_score, Gemm op for bbox_pred.
+            cls_score = self.graph.find_node_by_op_name("Softmax", "/roi_heads/Softmax")
+            bbox_pred = self.graph.find_node_by_op_name(
+                "Gemm", "/roi_heads/box_predictor/bbox_pred/Gemm"
+            )
+
+            # Linear transformation to convert box coordinates from (TopLeft, BottomRight) Corner encoding
+            # to CenterSize encoding. 1st NMS boxes are multiplied by transformation matrix in order to
+            # encode it into CenterSize format.
+            matmul_const = np.matrix(
+                "0.5 0 -1 0; 0 0.5 0 -1; 0.5 0 1 0; 0 0.5 0 1", dtype=np.float32
+            )
+            matmul_out = self.graph.matmul(
+                "RPN_NMS/detection_boxes_conversion", rpn_outputs[1], matmul_const
+            )
+
+            # Reshape node that prepares bbox_pred for scaling and second NMS.
+            bbox_pred_shape = np.asarray(
+                [self.batch_size, self.first_NMS_max_proposals, self.num_classes, 4],
+                dtype=np.int64,
+            )
+            bbox_pred_reshape = self.graph.op_with_const(
+                "Reshape", "bbox_pred/reshape", bbox_pred.outputs[0], bbox_pred_shape
+            )
+
+            # 0.1, 0.1, 0.2, 0.2 are localization head variance numbers, they scale bbox_pred_reshape, in order to get accurate coordinates.
+            scale_adj = np.expand_dims(
+                np.asarray([0.1, 0.1, 0.2, 0.2], dtype=np.float32), axis=(0, 1)
+            )
+            final_bbox_pred = self.graph.op_with_const(
+                "Mul", "bbox_pred/scale", bbox_pred_reshape[0], scale_adj
+            )
+
+            # Reshape node that prepares cls_score for slicing and second NMS.
+            cls_score_shape = np.array(
+                [self.batch_size, self.first_NMS_max_proposals, self.num_classes + 1],
+                dtype=np.int64,
+            )
+            cls_score_reshape = self.graph.op_with_const(
+                "Reshape", "cls_score/reshape", cls_score.outputs[0], cls_score_shape
+            )
+
+            # Slice operation to adjust third dimension of cls_score tensor, deletion of background class (81 in Detectron 2).
+            final_cls_score = self.graph.slice(
+                "cls_score/slicer", cls_score_reshape[0], 0, self.num_classes, 2
+            )
+
+            # Create NMS node.
+            nms_outputs = self.NMS(
+                final_bbox_pred[0],
+                final_cls_score[0],
+                matmul_out[0],
+                -1,
+                False,
+                self.second_NMS_max_proposals,
+                self.second_NMS_iou_threshold,
+                self.second_NMS_score_threshold,
+                second_nms_threshold,
+                "box_outputs",
+            )
+
+            # Create ROIAlign node.
+            mask_pooler_output = self.ROIAlign(
+                nms_outputs[1],
+                p2,
+                p3,
+                p4,
+                p5,
+                self.second_ROIAlign_pooled_size,
+                self.second_ROIAlign_sampling_ratio,
+                self.second_ROIAlign_type,
+                self.second_NMS_max_proposals,
+                "mask_pooler",
+            )
+
+            # Reshape mask pooler output.
+            mask_pooler_shape = np.asarray(
+                [
+                    self.second_NMS_max_proposals * self.batch_size,
+                    self.fpn_out_channels,
+                    self.second_ROIAlign_pooled_size,
+                    self.second_ROIAlign_pooled_size,
+                ],
+                dtype=np.int64,
+            )
+            mask_pooler_reshape_node = self.graph.op_with_const(
+                "Reshape", "mask_pooler/reshape", mask_pooler_output, mask_pooler_shape
+            )
+
+            # Get first Conv op in mask head and connect ROIAlign's squeezed output to it.
+            mask_head_conv = self.graph.find_node_by_op_name(
+                "Conv", "/roi_heads/mask_head/mask_fcn1/Conv"
+            )
+            mask_head_conv.inputs[0] = mask_pooler_reshape_node[0]
+
+            # Reshape node that is preparing 2nd NMS class outputs for Add node that comes next.
+            classes_reshape_shape = np.asarray(
+                [self.second_NMS_max_proposals * self.batch_size], dtype=np.int64
+            )
+            classes_reshape_node = self.graph.op_with_const(
+                "Reshape",
+                "box_outputs/reshape_classes",
+                nms_outputs[3],
+                classes_reshape_shape,
+            )
+
+            # This loop will generate an array used in Add node, which eventually will help Gather node to pick the single
+            # class of interest per bounding box, instead of creating 80 masks for every single bounding box.
+            add_array = []
+            for i in range(self.second_NMS_max_proposals * self.batch_size):
+                if i == 0:
+                    start_pos = 0
+                else:
+                    start_pos = i * self.num_classes
+                add_array.append(start_pos)
+
+            # This Add node is one of the Gather node inputs, Gather node performs gather on 0th axis of data tensor
+            # and requires indices that set tensors to be withing bounds, this Add node provides the bounds for Gather.
+            add_array = np.asarray(add_array, dtype=np.int32)
+            classes_add_node = self.graph.op_with_const(
+                "Add", "box_outputs/add", classes_reshape_node[0], add_array
+            )
+
+            # Get the last Conv op in mask head and reshape it to correctly gather class of interest's masks.
+            last_conv = self.graph.find_node_by_op_name(
+                "Conv", "/roi_heads/mask_head/predictor/Conv"
+            )
+            last_conv_reshape_shape = np.asarray(
+                [
+                    self.second_NMS_max_proposals * self.num_classes * self.batch_size,
+                    self.mask_out_res,
+                    self.mask_out_res,
+                ],
+                dtype=np.int64,
+            )
+            last_conv_reshape_node = self.graph.op_with_const(
+                "Reshape",
+                "mask_head/reshape_all_masks",
+                last_conv.outputs[0],
+                last_conv_reshape_shape,
+            )
+
+            # Gather node that selects only masks belonging to detected class, 79 other masks are discarded.
+            final_gather = self.graph.gather(
+                "mask_head/final_gather",
+                last_conv_reshape_node[0],
+                classes_add_node[0],
+                0,
+            )
+
+            # Get last Sigmoid node and connect Gather node to it.
+            mask_head_sigmoid = self.graph.find_node_by_op_name(
+                "Sigmoid", "/roi_heads/mask_head/Sigmoid"
+            )
+            mask_head_sigmoid.inputs[0] = final_gather[0]
+
+            # Final Reshape node, reshapes output of Sigmoid, important for various batch_size support (not tested yet).
+            final_graph_reshape_shape = np.asarray(
+                [
+                    self.batch_size,
+                    self.second_NMS_max_proposals,
+                    self.mask_out_res,
+                    self.mask_out_res,
+                ],
+                dtype=np.int64,
+            )
+            final_graph_reshape_node = self.graph.op_with_const(
+                "Reshape",
+                "mask_head/final_reshape",
+                mask_head_sigmoid.outputs[0],
+                final_graph_reshape_shape,
+            )
+            final_graph_reshape_node[0].dtype = np.float32
+            final_graph_reshape_node[0].name = "detection_masks"
+
+            return nms_outputs, final_graph_reshape_node[0]
+
+        # Only Detectron 2's Mask-RCNN R50-FPN 3x is supported currently.
+        p2, p3, p4, p5 = backbone()
+        rpn_outputs = proposal_generator(anchors, first_nms_threshold)
+        box_head_outputs, mask_head_output = roi_heads(
+            rpn_outputs, p2, p3, p4, p5, second_nms_threshold
+        )
+        # Append segmentation head output.
+        box_head_outputs.append(mask_head_output)
+        # Set graph outputs, both bbox and segmentation heads.
+        self.graph.outputs = box_head_outputs
+        self.sanitize()
+
+
+def main(args):
+    det2_gs = DET2GraphSurgeon(args.exported_onnx, args.det2_config, args.det2_weights)
+    det2_gs.update_preprocessor(args.batch_size)
+    anchors = det2_gs.get_anchors(args.sample_image)
+    det2_gs.process_graph(anchors, args.first_nms_threshold, args.second_nms_threshold)
+    det2_gs.save(args.onnx)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-i",
+        "--exported_onnx",
+        help="The exported to ONNX Detectron 2 Mask R-CNN",
+        type=str,
+    )
+    parser.add_argument(
+        "-o", "--onnx", help="The output ONNX model file to write", type=str
+    )
+    parser.add_argument(
+        "-c",
+        "--det2_config",
+        help="The Detectron 2 config file (.yaml) for the model",
+        type=str,
+    )
+    parser.add_argument(
+        "-w", "--det2_weights", help="The Detectron 2 model weights (.pkl)", type=str
+    )
+    parser.add_argument(
+        "-s", "--sample_image", help="Sample image for anchors generation", type=str
+    )
+    parser.add_argument(
+        "-b", "--batch_size", help="Batch size for the model", type=int, default=1
+    )
+    parser.add_argument(
+        "-t1",
+        "--first_nms_threshold",
+        help="Override the score threshold for the 1st NMS operation",
+        type=float,
+    )
+    parser.add_argument(
+        "-t2",
+        "--second_nms_threshold",
+        help="Override the score threshold for the 2nd NMS operation",
+        type=float,
+    )
+    args = parser.parse_args()
+    if not all(
+        [
+            args.exported_onnx,
+            args.onnx,
+            args.det2_config,
+            args.det2_weights,
+            args.sample_image,
+        ]
+    ):
+        parser.print_help()
+        print(
+            "\nThese arguments are required: --exported_onnx --onnx --det2_config --det2_weights and --sample_image"
+        )
+        sys.exit(1)
+    main(args)

samples/python/detectron2/eval_coco.py

@@ -0,0 +1,161 @@
+#
+# 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 sys
+import argparse
+import numpy as np
+import torch
+from PIL import Image
+from infer import TensorRTInfer
+from image_batcher import ImageBatcher
+
+try:
+    from detectron2.config import get_cfg
+    from detectron2.data import MetadataCatalog
+    from detectron2.evaluation import COCOEvaluator
+    from detectron2.structures import Instances, Boxes, ROIMasks
+except ImportError:
+    print("Could not import Detectron 2 modules. Maybe you did not install Detectron 2")
+    print(
+        "Please install Detectron 2, check https://github.com/facebookresearch/detectron2/blob/main/INSTALL.md"
+    )
+    sys.exit(1)
+
+
+def build_evaluator(dataset_name):
+    """
+    Create evaluator for a COCO dataset.
+    Currently only Mask R-CNN is supported, dataset of interest is COCO, so only COCOEvaluator is implemented.
+    """
+    evaluator_type = MetadataCatalog.get(dataset_name).evaluator_type
+    if evaluator_type in ["coco"]:
+        return COCOEvaluator(dataset_name)
+    else:
+        raise NotImplementedError("Evaluator type is not supported")
+
+
+def setup(config_file, weights):
+    """
+    Create config and perform basic setup.
+    """
+    cfg = get_cfg()
+    cfg.merge_from_file(config_file)
+    cfg.merge_from_list(["MODEL.WEIGHTS", weights])
+    cfg.freeze()
+    return cfg
+
+
+def main(args):
+    # Set up Detectron 2 config and build evaluator.
+    cfg = setup(args.det2_config, args.det2_weights)
+    dataset_name = cfg.DATASETS.TEST[0]
+    evaluator = build_evaluator(dataset_name)
+    evaluator.reset()
+
+    trt_infer = TensorRTInfer(args.engine)
+    batcher = ImageBatcher(
+        args.input, *trt_infer.input_spec(), config_file=args.det2_config
+    )
+
+    for batch, images, scales in batcher.get_batch():
+        print(
+            "Processing Image {} / {}".format(batcher.image_index, batcher.num_images),
+            end="\r",
+        )
+        detections = trt_infer.infer(batch, scales, args.nms_threshold)
+        for i in range(len(images)):
+            # Get inference image resolution.
+            infer_im = Image.open(images[i])
+            im_width, im_height = infer_im.size
+            pred_boxes = []
+            scores = []
+            pred_classes = []
+            # Number of detections.
+            num_instances = len(detections[i])
+            # Reserve numpy array to hold all mask predictions per image.
+            pred_masks = np.empty((num_instances, 28, 28), dtype=np.float32)
+            # Image ID, required for Detectron 2 evaluations.
+            source_id = int(os.path.splitext(os.path.basename(images[i]))[0])
+            # Loop over every single detection.
+            for n in range(num_instances):
+                det = detections[i][n]
+                # Append box coordinates data.
+                pred_boxes.append([det["ymin"], det["xmin"], det["ymax"], det["xmax"]])
+                # Append score.
+                scores.append(det["score"])
+                # Append class.
+                pred_classes.append(det["class"])
+                # Append mask.
+                pred_masks[n] = det["mask"]
+            # Create new Instances object required for Detectron 2 evalutions and add:
+            # boxes, scores, pred_classes, pred_masks.
+            image_shape = (im_height, im_width)
+            instances = Instances(image_shape)
+            instances.pred_boxes = Boxes(pred_boxes)
+            instances.scores = torch.tensor(scores)
+            instances.pred_classes = torch.tensor(pred_classes)
+            roi_masks = ROIMasks(torch.tensor(pred_masks))
+            instances.pred_masks = roi_masks.to_bitmasks(
+                instances.pred_boxes, im_height, im_width, args.iou_threshold
+            ).tensor
+            # Process evaluations per image.
+            image_dict = [{"instances": instances}]
+            input_dict = [{"image_id": source_id}]
+            evaluator.process(input_dict, image_dict)
+
+    # Final evaluations, generation of mAP accuracy performance.
+    evaluator.evaluate()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-e", "--engine", help="The TensorRT engine to infer with.")
+    parser.add_argument(
+        "-i",
+        "--input",
+        help="The input to infer, either a single image path, or a directory of images.",
+    )
+    parser.add_argument(
+        "-c",
+        "--det2_config",
+        help="The Detectron 2 config file (.yaml) for the model",
+        type=str,
+    )
+    parser.add_argument(
+        "-w", "--det2_weights", help="The Detectron 2 model weights (.pkl)", type=str
+    )
+    parser.add_argument(
+        "-t",
+        "--nms_threshold",
+        type=float,
+        help="Override the score threshold for the NMS operation, if higher than the threshold in the engine.",
+    )
+    parser.add_argument(
+        "--iou_threshold",
+        default=0.5,
+        type=float,
+        help="Select the IoU threshold for the mask segmentation. Range is 0 to 1. Pixel values more than threshold will become 1, less 0.",
+    )
+    args = parser.parse_args()
+    if not all([args.engine, args.input, args.det2_config, args.det2_weights]):
+        parser.print_help()
+        print(
+            "\nThese arguments are required: --engine --input --det2_config and --det2_weights"
+        )
+        sys.exit(1)
+    main(args)

samples/python/detectron2/image_batcher.py

@@ -0,0 +1,222 @@
+#
+# 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 sys
+import numpy as np
+from PIL import Image
+
+try:
+    from detectron2.config import get_cfg
+except ImportError:
+    print("Could not import Detectron 2 modules. Maybe you did not install Detectron 2")
+    print(
+        "Please install Detectron 2, check https://github.com/facebookresearch/detectron2/blob/main/INSTALL.md"
+    )
+    sys.exit(1)
+
+
+class ImageBatcher:
+    """
+    Creates batches of pre-processed images.
+    """
+
+    def __init__(
+        self,
+        input,
+        shape,
+        dtype,
+        max_num_images=None,
+        exact_batches=False,
+        config_file=None,
+    ):
+        """
+        :param input: The input directory to read images from.
+        :param shape: The tensor shape of the batch to prepare, either in NCHW or NHWC format.
+        :param dtype: The (numpy) datatype to cast the batched data to.
+        :param max_num_images: The maximum number of images to read from the directory.
+        :param exact_batches: This defines how to handle a number of images that is not an exact multiple of the batch
+        size. If false, it will pad the final batch with zeros to reach the batch size. If true, it will *remove* the
+        last few images in excess of a batch size multiple, to guarantee batches are exact (useful for calibration).
+        :param config_file: The path pointing to the Detectron 2 yaml file which describes the model.
+        """
+
+        def det2_setup(config_file):
+            """
+            Create configs and perform basic setups.
+            """
+            cfg = get_cfg()
+            if config_file is not None:
+                cfg.merge_from_file(config_file)
+            cfg.freeze()
+            return cfg
+
+        # Set up Detectron 2 model configuration.
+        self.det2_cfg = det2_setup(config_file)
+
+        # Extract min and max dimensions for testing.
+        self.min_size_test = self.det2_cfg.INPUT.MIN_SIZE_TEST
+        self.max_size_test = self.det2_cfg.INPUT.MAX_SIZE_TEST
+
+        # Find images in the given input path.
+        input = os.path.realpath(input)
+        self.images = []
+
+        extensions = [".jpg", ".jpeg", ".png", ".bmp", ".ppm"]
+
+        def is_image(path):
+            return (
+                os.path.isfile(path) and os.path.splitext(path)[1].lower() in extensions
+            )
+
+        if os.path.isdir(input):
+            self.images = [
+                os.path.join(input, f)
+                for f in os.listdir(input)
+                if is_image(os.path.join(input, f))
+            ]
+            self.images.sort()
+        elif os.path.isfile(input):
+            if is_image(input):
+                self.images.append(input)
+        self.num_images = len(self.images)
+        if self.num_images < 1:
+            print("No valid {} images found in {}".format("/".join(extensions), input))
+            sys.exit(1)
+
+        # Handle Tensor Shape.
+        self.dtype = dtype
+        self.shape = shape
+        assert len(self.shape) == 4
+        self.batch_size = shape[0]
+        assert self.batch_size > 0
+        self.format = None
+        self.width = -1
+        self.height = -1
+        if self.shape[1] == 3:
+            self.format = "NCHW"
+            self.height = self.shape[2]
+            self.width = self.shape[3]
+        elif self.shape[3] == 3:
+            self.format = "NHWC"
+            self.height = self.shape[1]
+            self.width = self.shape[2]
+        assert all([self.format, self.width > 0, self.height > 0])
+
+        # Adapt the number of images as needed.
+        if max_num_images and 0 < max_num_images < len(self.images):
+            self.num_images = max_num_images
+        if exact_batches:
+            self.num_images = self.batch_size * (self.num_images // self.batch_size)
+        if self.num_images < 1:
+            print("Not enough images to create batches")
+            sys.exit(1)
+        self.images = self.images[0 : self.num_images]
+
+        # Subdivide the list of images into batches.
+        self.num_batches = 1 + int((self.num_images - 1) / self.batch_size)
+        self.batches = []
+        for i in range(self.num_batches):
+            start = i * self.batch_size
+            end = min(start + self.batch_size, self.num_images)
+            self.batches.append(self.images[start:end])
+
+        # Indices.
+        self.image_index = 0
+        self.batch_index = 0
+
+    def preprocess_image(self, image_path):
+        """
+        The image preprocessor loads an image from disk and prepares it as needed for batching. This includes padding,
+        resizing, normalization, data type casting, and transposing.
+        This Image Batcher implements one algorithm for now:
+        * Resizes and pads the image to fit the input size.
+        :param image_path: The path to the image on disk to load.
+        :return: Two values: A numpy array holding the image sample, ready to be contacatenated into the rest of the
+        batch, and the resize scale used, if any.
+        """
+
+        def resize_pad(image, pad_color=(0, 0, 0)):
+            """
+            A subroutine to implement padding and resizing. This will resize the image to fit fully within the input
+            size, and pads the remaining bottom-right portions with the value provided.
+            :param image: The PIL image object
+            :pad_color: The RGB values to use for the padded area. Default: Black/Zeros.
+            :return: Two values: The PIL image object already padded and cropped, and the resize scale used.
+            """
+
+            # Get characteristics.
+            width, height = image.size
+
+            # Replicates behavior of ResizeShortestEdge augmentation.
+            size = self.min_size_test * 1.0
+            pre_scale = size / min(height, width)
+            if height < width:
+                newh, neww = size, pre_scale * width
+            else:
+                newh, neww = pre_scale * height, size
+
+            # If delta between min and max dimensions is so that max sized dimension reaches self.max_size_test
+            # before min dimension reaches self.min_size_test, keeping the same aspect ratio. We still need to
+            # maintain the same aspect ratio and keep max dimension at self.max_size_test.
+            if max(newh, neww) > self.max_size_test:
+                pre_scale = self.max_size_test * 1.0 / max(newh, neww)
+                newh = newh * pre_scale
+                neww = neww * pre_scale
+            neww = int(neww + 0.5)
+            newh = int(newh + 0.5)
+
+            # Scaling factor for normalized box coordinates scaling in post-processing.
+            scaling = max(newh / height, neww / width)
+
+            # Padding.
+            image = image.resize((neww, newh), resample=Image.BILINEAR)
+            pad = Image.new("RGB", (self.width, self.height))
+            pad.paste(pad_color, [0, 0, self.width, self.height])
+            pad.paste(image)
+            return pad, scaling
+
+        scale = None
+        image = Image.open(image_path)
+        image = image.convert(mode="RGB")
+        # Pad with mean values of COCO dataset, since padding is applied before actual model's
+        # preprocessor steps (Sub, Div ops), we need to pad with mean values in order to reverse
+        # the effects of Sub and Div, so that padding after model's preprocessor will be with actual 0s.
+        image, scale = resize_pad(image, (124, 116, 104))
+        image = np.asarray(image, dtype=np.float32)
+        # Change HWC -> CHW.
+        image = np.transpose(image, (2, 0, 1))
+        # Change RGB -> BGR.
+        return image[[2, 1, 0]], scale
+
+    def get_batch(self):
+        """
+        Retrieve the batches. This is a generator object, so you can use it within a loop as:
+        for batch, images in batcher.get_batch():
+           ...
+        Or outside of a batch with the next() function.
+        :return: A generator yielding three items per iteration: a numpy array holding a batch of images, the list of
+        paths to the images loaded within this batch, and the list of resize scales for each image in the batch.
+        """
+        for i, batch_images in enumerate(self.batches):
+            batch_data = np.zeros(self.shape, dtype=self.dtype)
+            batch_scales = [None] * len(batch_images)
+            for i, image in enumerate(batch_images):
+                self.image_index += 1
+                batch_data[i], batch_scales[i] = self.preprocess_image(image)
+            self.batch_index += 1
+            yield batch_data, batch_images, batch_scales

samples/python/detectron2/infer.py

@@ -0,0 +1,325 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2025 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 sys
+import argparse
+import numpy as np
+import tensorrt as trt
+from cuda.bindings import runtime as cudart
+from image_batcher import ImageBatcher
+from visualize import visualize_detections
+
+sys.path.insert(1, os.path.join(os.path.dirname(os.path.realpath(__file__)), os.pardir))
+import common
+
+
+class TensorRTInfer:
+    """
+    Implements inference for the Model TensorRT engine.
+    """
+
+    def __init__(self, engine_path):
+        """
+        :param engine_path: The path to the serialized engine to load from disk.
+        """
+
+        # Load TRT engine
+        self.logger = trt.Logger(trt.Logger.ERROR)
+        trt.init_libnvinfer_plugins(self.logger, namespace="")
+        with open(engine_path, "rb") as f, trt.Runtime(self.logger) as runtime:
+            assert runtime
+            self.engine = runtime.deserialize_cuda_engine(f.read())
+        assert self.engine
+        self.context = self.engine.create_execution_context()
+        assert self.context
+
+        # Setup I/O bindings
+        self.inputs = []
+        self.outputs = []
+        self.device_memories = []
+        for i in range(self.engine.num_io_tensors):
+            name = self.engine.get_tensor_name(i)
+            is_input = False
+            if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT:
+                is_input = True
+            dtype = self.engine.get_tensor_dtype(name)
+            shape = self.engine.get_tensor_shape(name)
+            if is_input:
+                self.batch_size = shape[0]
+            size = np.dtype(trt.nptype(dtype)).itemsize
+            for s in shape:
+                size *= s
+            device_mem = common.DeviceMem(size)
+            binding = {
+                "index": i,
+                "name": name,
+                "dtype": np.dtype(trt.nptype(dtype)),
+                "shape": list(shape),
+                "allocation": device_mem.device_ptr,
+                "size": size,
+            }
+            self.device_memories.append(device_mem)
+            if is_input:
+                self.inputs.append(binding)
+            else:
+                self.outputs.append(binding)
+
+        assert self.batch_size > 0
+        assert len(self.inputs) > 0
+        assert len(self.outputs) > 0
+        assert len(self.device_memories) > 0
+
+    def input_spec(self):
+        """
+        Get the specs for the input tensor of the network. Useful to prepare memory allocations.
+        :return: Two items, the shape of the input tensor and its (numpy) datatype.
+        """
+        return self.inputs[0]["shape"], self.inputs[0]["dtype"]
+
+    def output_spec(self):
+        """
+        Get the specs for the output tensors of the network. Useful to prepare memory allocations.
+        :return: A list with two items per element, the shape and (numpy) datatype of each output tensor.
+        """
+        specs = []
+        for o in self.outputs:
+            specs.append((o["shape"], o["dtype"]))
+        return specs
+
+    def infer(self, batch, scales=None, nms_threshold=None):
+        """
+        Execute inference on a batch of images. The images should already be batched and preprocessed, as prepared by
+        the ImageBatcher class. Memory copying to and from the GPU device will be performed here.
+        :param batch: A numpy array holding the image batch.
+        :param scales: The image resize scales for each image in this batch. Default: No scale postprocessing applied.
+        :return: A nested list for each image in the batch and each detection in the list.
+        """
+
+        # Prepare the output data.
+        outputs = []
+        for shape, dtype in self.output_spec():
+            outputs.append(np.zeros(shape, dtype))
+
+        # Process I/O and execute the network.
+        common.memcpy_host_to_device(
+            self.inputs[0]["allocation"], np.ascontiguousarray(batch)
+        )
+
+        self.context.execute_v2(self.allocations)
+        for o in range(len(outputs)):
+            common.memcpy_device_to_host(outputs[o], self.outputs[o]["allocation"])
+
+        # Process the results.
+        nums = outputs[0]
+        boxes = outputs[1]
+        scores = outputs[2]
+        pred_classes = outputs[3]
+        masks = outputs[4]
+
+        detections = []
+        for i in range(self.batch_size):
+            detections.append([])
+            for n in range(int(nums[i])):
+                # Select a mask.
+                mask = masks[i][n]
+
+                # Calculate scaling values for bboxes.
+                scale = self.inputs[0]["shape"][2]
+                scale /= scales[i]
+                scale_y = scale
+                scale_x = scale
+
+                if nms_threshold and scores[i][n] < nms_threshold:
+                    continue
+                # Append to detections
+                detections[i].append(
+                    {
+                        "ymin": boxes[i][n][0] * scale_y,
+                        "xmin": boxes[i][n][1] * scale_x,
+                        "ymax": boxes[i][n][2] * scale_y,
+                        "xmax": boxes[i][n][3] * scale_x,
+                        "score": scores[i][n],
+                        "class": int(pred_classes[i][n]),
+                        "mask": mask,
+                    }
+                )
+        return detections
+
+
+def main(args):
+    output_dir = os.path.realpath(args.output)
+    os.makedirs(output_dir, exist_ok=True)
+
+    labels = [
+        "person",
+        "bicycle",
+        "car",
+        "motorcycle",
+        "airplane",
+        "bus",
+        "train",
+        "truck",
+        "boat",
+        "traffic light",
+        "fire hydrant",
+        "stop sign",
+        "parking meter",
+        "bench",
+        "bird",
+        "cat",
+        "dog",
+        "horse",
+        "sheep",
+        "cow",
+        "elephant",
+        "bear",
+        "zebra",
+        "giraffe",
+        "backpack",
+        "umbrella",
+        "handbag",
+        "tie",
+        "suitcase",
+        "frisbee",
+        "skis",
+        "snowboard",
+        "sports ball",
+        "kite",
+        "baseball bat",
+        "baseball glove",
+        "skateboard",
+        "surfboard",
+        "tennis racket",
+        "bottle",
+        "wine glass",
+        "cup",
+        "fork",
+        "knife",
+        "spoon",
+        "bowl",
+        "banana",
+        "apple",
+        "sandwich",
+        "orange",
+        "broccoli",
+        "carrot",
+        "hot dog",
+        "pizza",
+        "donut",
+        "cake",
+        "chair",
+        "couch",
+        "potted plant",
+        "bed",
+        "dining table",
+        "toilet",
+        "tv",
+        "laptop",
+        "mouse",
+        "remote",
+        "keyboard",
+        "cell phone",
+        "microwave",
+        "oven",
+        "toaster",
+        "sink",
+        "refrigerator",
+        "book",
+        "clock",
+        "vase",
+        "scissors",
+        "teddy bear",
+        "hair drier",
+        "toothbrush",
+    ]
+
+    trt_infer = TensorRTInfer(args.engine)
+    batcher = ImageBatcher(
+        args.input, *trt_infer.input_spec(), config_file=args.det2_config
+    )
+    for batch, images, scales in batcher.get_batch():
+        print(
+            "Processing Image {} / {}".format(batcher.image_index, batcher.num_images),
+            end="\r",
+        )
+        detections = trt_infer.infer(batch, scales, args.nms_threshold)
+        for i in range(len(images)):
+            basename = os.path.splitext(os.path.basename(images[i]))[0]
+            # Image Visualizations
+            output_path = os.path.join(output_dir, "{}.png".format(basename))
+            visualize_detections(
+                images[i], output_path, detections[i], labels, args.iou_threshold
+            )
+            # Text Results
+            output_results = ""
+            for d in detections[i]:
+                line = [
+                    d["xmin"],
+                    d["ymin"],
+                    d["xmax"],
+                    d["ymax"],
+                    d["score"],
+                    d["class"],
+                ]
+                output_results += "\t".join([str(f) for f in line]) + "\n"
+            with open(os.path.join(args.output, "{}.txt".format(basename)), "w") as f:
+                f.write(output_results)
+    print()
+    print("Finished Processing")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-e", "--engine", default=None, help="The serialized TensorRT engine"
+    )
+    parser.add_argument(
+        "-i", "--input", default=None, help="Path to the image or directory to process"
+    )
+    parser.add_argument(
+        "-c",
+        "--det2_config",
+        help="The Detectron 2 config file (.yaml) for the model",
+        type=str,
+    )
+    parser.add_argument(
+        "-o",
+        "--output",
+        default=None,
+        help="Directory where to save the visualization results",
+    )
+    parser.add_argument(
+        "-t",
+        "--nms_threshold",
+        type=float,
+        help="Override the score threshold for the NMS operation, if higher than the threshold in the engine.",
+    )
+    parser.add_argument(
+        "--iou_threshold",
+        default=0.5,
+        type=float,
+        help="Select the IoU threshold for the mask segmentation. Range is 0 to 1. Pixel values more than threshold will become 1, less 0",
+    )
+    args = parser.parse_args()
+    if not all([args.engine, args.input, args.output, args.det2_config]):
+        parser.print_help()
+        print(
+            "\nThese arguments are required: --engine --input --output and --det2_config"
+        )
+        sys.exit(1)
+    main(args)

samples/python/detectron2/onnx_utils.py

@@ -0,0 +1,332 @@
+#
+# 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 logging
+import numpy as np
+import onnx_graphsurgeon as gs
+
+logging.basicConfig(level=logging.INFO)
+logging.getLogger("ModelHelper").setLevel(logging.INFO)
+log = logging.getLogger("ModelHelper")
+
+
+@gs.Graph.register()
+def op_with_const(self, op, name, input, value):
+    """
+    Add an operation with constant to the graph which will operate on the input tensor with the value(s) given.
+    :param op: The ONNX operation to perform, i.e. "Add" or "Mul".
+    :param input: The tensor to operate on.
+    :param value: The value array to operate with.
+    :param name: The name to use for the node.
+    """
+    input_tensor = input if type(input) is gs.Variable else input[0]
+    log.debug("Created {} node '{}': {}".format(op, name, value.squeeze()))
+    const = gs.Constant(name="{}_value:0".format(name), values=value)
+    return self.layer(
+        name=name, op=op, inputs=[input_tensor, const], outputs=[name + ":0"]
+    )
+
+
+@gs.Graph.register()
+def matmul(self, name, input, value):
+    """
+    Add MatMul operation to the graph which will operate on the input tensor with the value(s) given.
+    :param input: The tensor to operate on.
+    :param value: The linear transformation matrix to operate with.
+    :param name: The name to use for the node.
+    """
+    input_tensor = input if type(input) is gs.Variable else input[0]
+    log.debug("Created {} node '{}': {}".format("MatMul", name, value.squeeze()))
+    const = gs.Constant(name="{}_value:0".format(name), values=value)
+    return self.layer(
+        name=name, op="MatMul", inputs=[input_tensor, const], outputs=[name + ":0"]
+    )
+
+
+@gs.Graph.register()
+def clip(self, name, input, clip_min, clip_max):
+    """
+    Add Clip operation to the graph which will operate on the input tensor with the value(s) given.
+    :param input: The tensor to operate on.
+    :param name: The name to use for the node.
+    :param clip_min: Minimum value to include, less is clipped.
+    :param clip_max: Maximum value to include, more is clipped.
+    """
+    input_tensor = input if type(input) is gs.Variable else input[0]
+    log.debug("Created {} node '{}".format("Clip", name))
+    const_min = gs.Constant(
+        name="{}_value:0".format(name), values=np.asarray([clip_min], dtype=np.float32)
+    )
+    const_max = gs.Constant(
+        name="{}_value:1".format(name), values=np.asarray([clip_max], dtype=np.float32)
+    )
+    return self.layer(
+        name=name,
+        op="Clip",
+        inputs=[input_tensor, const_min, const_max],
+        outputs=[name + ":0"],
+    )
+
+
+@gs.Graph.register()
+def slice(self, name, input, starts, ends, axes):
+    """
+    Add Slice operation to the graph which will operate on the input tensor with the value(s) given.
+    :param op: The ONNX operation to perform, i.e. "Add" or "Mul".
+    :param input: The tensor to operate on.
+    :param name: The name to use for the node.
+    :param starts: Value at which Slice starts.
+    :param ends: Value at which Slice ends.
+    :param axes: Axes on which Slice operation should be performed.
+    """
+
+    input_tensor = input if type(input) is gs.Variable else input[0]
+    log.debug("Created {} node '{}".format("Slice", name))
+    const_start = gs.Constant(
+        name="{}_value:0".format(name), values=np.asarray([starts], dtype=np.int64)
+    )
+    const_end = gs.Constant(
+        name="{}_value:1".format(name), values=np.asarray([ends], dtype=np.int64)
+    )
+    const_axes = gs.Constant(
+        name="{}_value:2".format(name), values=np.asarray([axes], dtype=np.int64)
+    )
+    return self.layer(
+        name=name,
+        op="Slice",
+        inputs=[input_tensor, const_start, const_end, const_axes],
+        outputs=[name + ":0"],
+    )
+
+
+@gs.Graph.register()
+def unsqueeze(self, name, input, axes=[3]):
+    """
+    Adds to the graph an Unsqueeze node for the given axes and to the given input.
+    :param self: The gs.Graph object being extended.
+    :param name: The name to use for the node.
+    :param input: The tensor to be "unsqueezed".
+    :param axes: A list of axes on which to add the new dimension(s).
+    :return: The first output tensor, to allow chained graph construction.
+    """
+    input_tensor = input if type(input) is gs.Variable else input[0]
+    log.debug("Created Unsqueeze node '{}': {}".format(name, axes))
+    return self.layer(
+        name=name,
+        op="Unsqueeze",
+        inputs=[input_tensor],
+        outputs=[name + ":0"],
+        attrs={"axes": axes},
+    )
+
+
+@gs.Graph.register()
+def squeeze(self, name, input, axes=[2]):
+    """
+    Adds to the graph an Squeeze node for the given axes and to the given input.
+    :param self: The gs.Graph object being extended.
+    :param name: The name to use for the node.
+    :param input: The tensor to be "squeezed".
+    :param axes: A list of axes on which to remove a dimension(s).
+    :return: The first output tensor, to allow chained graph construction.
+    """
+    input_tensor = input if type(input) is gs.Variable else input[0]
+    log.debug("Created Squeeze node '{}': {}".format(name, axes))
+    return self.layer(
+        name=name,
+        op="Squeeze",
+        inputs=[input_tensor],
+        outputs=[name + ":0"],
+        attrs={"axes": axes},
+    )
+
+
+@gs.Graph.register()
+def gather(self, name, data, indices, axes=0):
+    """
+    Adds to the graph a Gather node for the given axes and to the given input.
+    :param self: The gs.Graph object being extended.
+    :param name: The name to use for the node.
+    :param data: Data from which to gather specific tensors.
+    :param indices: Indices by which to gather data tensors.
+    :param axes: A list of axes on which to perform gather operation
+    """
+    data_tensor = data if type(data) is gs.Variable else data[0]
+    indices_tensor = indices if type(indices) is gs.Variable else indices[0]
+    log.debug("Created Gather node '{}': {}".format(name, axes))
+    return self.layer(
+        name=name,
+        op="Gather",
+        inputs=[data_tensor, indices_tensor],
+        outputs=[name + ":0"],
+        attrs={"axes": axes},
+    )
+
+
+@gs.Graph.register()
+def transpose(self, name, input, perm):
+    """
+    Adds to the graph a Transpose node for the given axes permutation and to the given input.
+    :param self: The gs.Graph object being extended.
+    :param name: The name to use for the node.
+    :param input: The tensor to be transposed.
+    :param perm: A list of axes defining their order after transposing occurs.
+    :return: The first output tensor, to allow chained graph construction.
+    """
+    input_tensor = input if type(input) is gs.Variable else input[0]
+    log.debug("Created Transpose node '{}': {}".format(name, perm))
+    return self.layer(
+        name=name,
+        op="Transpose",
+        inputs=[input_tensor],
+        outputs=[name + ":0"],
+        attrs={"perm": perm},
+    )
+
+
+@gs.Graph.register()
+def sigmoid(self, name, input):
+    """
+    Adds to the graph a Sigmoid node for the given input.
+    :param self: The gs.Graph object being extended.
+    :param name: The name to use for the node.
+    :param input: The tensor to be applied to.
+    :return: The first output tensor, to allow chained graph construction.
+    """
+    input_tensor = input if type(input) is gs.Variable else input[0]
+    log.debug("Created Sigmoid node '{}'".format(name))
+    return self.layer(
+        name=name, op="Sigmoid", inputs=[input_tensor], outputs=[name + ":0"]
+    )
+
+
+@gs.Graph.register()
+def plugin(self, op, name, inputs: list, outputs: list, attrs):
+    """
+    Adds to the graph a TensorRT plugin node with the given name, inputs and outputs. The attrs dictionary holds
+    attributes to be added to the plugin node.
+    :param self: The gs.Graph object being extended.
+    :param op: The registered name for the TensorRT plugin.
+    :param name: The name to use for the node.
+    :param inputs: The list of tensors to use an inputs.
+    :param outputs: The list of tensors to use as outputs.
+    :param attrs: The dictionary to use as attributes.
+    :return: The first output tensor, to allow chained graph construction.
+    """
+    log.debug("Created TRT Plugin node '{}': {}".format(name, attrs))
+    return self.layer(op=op, name=name, inputs=inputs, outputs=outputs, attrs=attrs)
+
+
+@gs.Graph.register()
+def find_node_by_op(self, op):
+    """
+    Finds the first node in the graph with the given operation name.
+    :param self: The gs.Graph object being extended.
+    :param op: The operation name to search for.
+    :return: The first node matching that performs that op.
+    """
+    for node in self.nodes:
+        if node.op == op:
+            return node
+    return None
+
+
+@gs.Graph.register()
+def find_node_by_op_name(self, op, name):
+    """
+    Finds the first node in the graph with the given operation name.
+    :param self: The gs.Graph object being extended.
+    :param op: The operation name to search for.
+    :param name: Selected node name.
+    :return: The first node matching that performs that op.
+    """
+    for node in self.nodes:
+        if node.op == op and node.name == name:
+            return node
+    return None
+
+
+@gs.Graph.register()
+def find_node_by_op_input_output_name(
+    self, op, input_name, output_name, input_pos=0, output_pos=0
+):
+    """
+    Finds the first node in the graph with the given operation name.
+    :param self: The gs.Graph object being extended.
+    :param op: The operation name to search for.
+    :param input_pos: Which input to consider, default is 0.
+    :param output_pos: Which output to consider, default is 0.
+    :param input_name: Selected input's name.
+    :param output_name: Selected output's name.
+    :return: The first node matching that performs that op.
+    """
+    for node in self.nodes:
+        if (
+            node.op == op
+            and node.inputs[input_pos].name == input_name
+            and node.outputs[output_pos].name == output_name
+        ):
+            return node
+    return None
+
+
+@gs.Graph.register()
+def find_descendant_by_op(self, node, op, depth=10):
+    """
+    Starting from the given node, finds a node lower in the graph matching the given operation name.
+    This is not an exhaustive graph search.
+    In order to graph search bfs is used, so runtime complexity is O(V+E).
+    :param self: The gs.Graph object being extended.
+    :param node: The node to start searching from.
+    :param op: The operation name to search for.
+    :param depth: Stop searching after traversing these many nodes.
+    :return: The first descendant node matching that performs that op.
+    """
+    queue = []
+    for i in range(depth):
+        queue.append(node.o())
+        while queue:
+            node = queue.pop(0)
+            if node.op == op:
+                return node
+            for child in node.outputs[0].outputs:
+                queue.append(child)
+    return None
+
+
+@gs.Graph.register()
+def find_ancestor_by_op(self, node, op, depth=10):
+    """
+    Starting from the given node, finds a node higher in the graph matching the given operation name.
+    This is not an exhaustive graph search.
+    In order to graph search bfs is used, so runtime complexity is O(V+E).
+    :param self: The gs.Graph object being extended.
+    :param node: The node to start searching from.
+    :param op: The operation name to search for.
+    :param depth: Stop searching after traversing these many nodes.
+    :return: The first ancestor node matching that performs that op.
+    """
+    queue = []
+    for i in range(depth):
+        queue.append(node.i())
+        while queue:
+            node = queue.pop(0)
+            if node.op == op:
+                return node
+            for child in node.inputs[-1].inputs:
+                queue.append(child)
+    return None

samples/python/detectron2/requirements.txt

@@ -0,0 +1,12 @@
+onnx==1.18.0
+onnxruntime==1.18.1
+onnxscript>=0.6.0
+Pillow==11.3.0
+git+https://github.com/facebookresearch/detectron2.git
+git+https://github.com/NVIDIA/TensorRT#subdirectory=tools/onnx-graphsurgeon
+cuda-python==12.9.0
+pywin32; platform_system == "Windows"
+pyyaml==6.0.3
+requests==2.32.4
+tqdm==4.66.4
+numpy==1.26.4

samples/python/detectron2/visualize.py

@@ -0,0 +1,255 @@
+#
+# 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 PIL.Image as Image
+import PIL.ImageDraw as ImageDraw
+import PIL.ImageFont as ImageFont
+import PIL.ImageFilter as ImageFilter
+
+
+COLORS = [
+    "GoldenRod",
+    "MediumTurquoise",
+    "GreenYellow",
+    "SteelBlue",
+    "DarkSeaGreen",
+    "SeaShell",
+    "LightGrey",
+    "IndianRed",
+    "DarkKhaki",
+    "LawnGreen",
+    "WhiteSmoke",
+    "Peru",
+    "LightCoral",
+    "FireBrick",
+    "OldLace",
+    "LightBlue",
+    "SlateGray",
+    "OliveDrab",
+    "NavajoWhite",
+    "PaleVioletRed",
+    "SpringGreen",
+    "AliceBlue",
+    "Violet",
+    "DeepSkyBlue",
+    "Red",
+    "MediumVioletRed",
+    "PaleTurquoise",
+    "Tomato",
+    "Azure",
+    "Yellow",
+    "Cornsilk",
+    "Aquamarine",
+    "CadetBlue",
+    "CornflowerBlue",
+    "DodgerBlue",
+    "Olive",
+    "Orchid",
+    "LemonChiffon",
+    "Sienna",
+    "OrangeRed",
+    "Orange",
+    "DarkSalmon",
+    "Magenta",
+    "Wheat",
+    "Lime",
+    "GhostWhite",
+    "SlateBlue",
+    "Aqua",
+    "MediumAquaMarine",
+    "LightSlateGrey",
+    "MediumSeaGreen",
+    "SandyBrown",
+    "YellowGreen",
+    "Plum",
+    "FloralWhite",
+    "LightPink",
+    "Thistle",
+    "DarkViolet",
+    "Pink",
+    "Crimson",
+    "Chocolate",
+    "DarkGrey",
+    "Ivory",
+    "PaleGreen",
+    "DarkGoldenRod",
+    "LavenderBlush",
+    "SlateGrey",
+    "DeepPink",
+    "Gold",
+    "Cyan",
+    "LightSteelBlue",
+    "MediumPurple",
+    "ForestGreen",
+    "DarkOrange",
+    "Tan",
+    "Salmon",
+    "PaleGoldenRod",
+    "LightGreen",
+    "LightSlateGray",
+    "HoneyDew",
+    "Fuchsia",
+    "LightSeaGreen",
+    "DarkOrchid",
+    "Green",
+    "Chartreuse",
+    "LimeGreen",
+    "AntiqueWhite",
+    "Beige",
+    "Gainsboro",
+    "Bisque",
+    "SaddleBrown",
+    "Silver",
+    "Lavender",
+    "Teal",
+    "LightCyan",
+    "PapayaWhip",
+    "Purple",
+    "Coral",
+    "BurlyWood",
+    "LightGray",
+    "Snow",
+    "MistyRose",
+    "PowderBlue",
+    "DarkCyan",
+    "White",
+    "Turquoise",
+    "MediumSlateBlue",
+    "PeachPuff",
+    "Moccasin",
+    "LightSalmon",
+    "SkyBlue",
+    "Khaki",
+    "MediumSpringGreen",
+    "BlueViolet",
+    "MintCream",
+    "Linen",
+    "SeaGreen",
+    "HotPink",
+    "LightYellow",
+    "BlanchedAlmond",
+    "RoyalBlue",
+    "RosyBrown",
+    "MediumOrchid",
+    "DarkTurquoise",
+    "LightGoldenRodYellow",
+    "LightSkyBlue",
+]
+
+
+# Overlay mask with transparency on top of the image.
+def overlay(image, mask, color, alpha_transparency=0.5):
+    for channel in range(3):
+        image[:, :, channel] = np.where(
+            mask == 1,
+            image[:, :, channel] * (1 - alpha_transparency)
+            + alpha_transparency * color[channel] * 255,
+            image[:, :, channel],
+        )
+    return image
+
+
+def visualize_detections(
+    image_path, output_path, detections, labels=[], iou_threshold=0.5
+):
+    image = Image.open(image_path).convert(mode="RGB")
+    # Get image dimensions.
+    im_width, im_height = image.size
+    line_width = 2
+    font = ImageFont.load_default()
+    for d in detections:
+        color = COLORS[d["class"] % len(COLORS)]
+        # Dynamically convert PIL color into RGB numpy array.
+        pixel_color = Image.new("RGB", (1, 1), color)
+        # Normalize.
+        np_color = (np.asarray(pixel_color)[0][0]) / 255
+        # TRT instance segmentation masks.
+        if isinstance(d["mask"], np.ndarray) and d["mask"].shape == (28, 28):
+            # PyTorch uses [x1,y1,x2,y2] format instead of regular [y1,x1,y2,x2].
+            d["ymin"], d["xmin"], d["ymax"], d["xmax"] = (
+                d["xmin"],
+                d["ymin"],
+                d["xmax"],
+                d["ymax"],
+            )
+            # Get detection bbox resolution.
+            det_width = round(d["xmax"] - d["xmin"])
+            det_height = round(d["ymax"] - d["ymin"])
+            # Slight scaling, to get binary masks after float32 -> uint8
+            # conversion, if not scaled all pixels are zero.
+            mask = d["mask"] > iou_threshold
+            # Convert float32 -> uint8.
+            mask = mask.astype(np.uint8)
+            # Create an image out of predicted mask array.
+            small_mask = Image.fromarray(mask)
+            # Upsample mask to detection bbox's size.
+            mask = small_mask.resize((det_width, det_height), resample=Image.BILINEAR)
+            # Create an original image sized template for correct mask placement.
+            pad = Image.new("L", (im_width, im_height))
+            # Place your mask according to detection bbox placement.
+            pad.paste(mask, (round(d["xmin"]), (round(d["ymin"]))))
+            # Reconvert mask into numpy array for evaluation.
+            padded_mask = np.array(pad)
+            # Creat np.array from original image, copy in order to modify.
+            image_copy = np.asarray(image).copy()
+            # Image with overlaid mask.
+            masked_image = overlay(image_copy, padded_mask, np_color)
+            # Reconvert back to PIL.
+            image = Image.fromarray(masked_image)
+
+        # Bbox lines.
+        draw = ImageDraw.Draw(image)
+        draw.line(
+            [
+                (d["xmin"], d["ymin"]),
+                (d["xmin"], d["ymax"]),
+                (d["xmax"], d["ymax"]),
+                (d["xmax"], d["ymin"]),
+                (d["xmin"], d["ymin"]),
+            ],
+            width=line_width,
+            fill=color,
+        )
+        label = "Class {}".format(d["class"])
+        if d["class"] < len(labels):
+            label = "{}".format(labels[d["class"]])
+        score = d["score"]
+        text = "{}: {}%".format(label, int(100 * score))
+        if score < 0:
+            text = label
+        left, top, right, bottom = font.getbbox(text)
+        text_width, text_height = right - left, bottom - top
+        text_bottom = max(text_height, d["ymin"])
+        text_left = d["xmin"]
+        margin = np.ceil(0.05 * text_height)
+        draw.rectangle(
+            [
+                (text_left, text_bottom - text_height - 2 * margin),
+                (text_left + text_width, text_bottom),
+            ],
+            fill=color,
+        )
+        draw.text(
+            (text_left + margin, text_bottom - text_height - margin),
+            text,
+            fill="black",
+            font=font,
+        )
+    if output_path is None:
+        return image
+    image.save(output_path)

samples/python/downloader.py

@@ -0,0 +1,259 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2025 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 argparse
+import errno
+import hashlib
+import logging
+import os
+import sys
+
+
+logger = logging.getLogger("downloader")
+
+
+class DataFile:
+    """Holder of a data file."""
+
+    def __init__(self, attr):
+        self.attr = attr
+        self.path = attr["path"]
+        self.url = attr["url"]
+        if "checksum" not in attr:
+            logger.warning("Checksum of %s not provided!", self.path)
+        self.checksum = attr.get("checksum", None)
+
+    def __str__(self):
+        return str(self.attr)
+
+
+class SampleData:
+    """Holder of data files of an sample."""
+
+    def __init__(self, attr):
+        self.attr = attr
+        self.sample = attr["sample"]
+        files = attr.get("files", [])
+        self.files = [DataFile(f) for f in files]
+
+    def __str__(self):
+        return str(self.attr)
+
+
+def _loadYAML(yaml_path):
+    with open(yaml_path, "rb") as f:
+        import yaml
+
+        y = yaml.load(f, yaml.FullLoader)
+        return SampleData(y)
+
+
+def _checkMD5(path, refMD5):
+    md5 = hashlib.md5(open(path, "rb").read()).hexdigest()
+    return md5 == refMD5
+
+
+def _createDirIfNeeded(path):
+    the_dir = os.path.dirname(path)
+    try:
+        os.makedirs(the_dir)
+    except OSError as e:
+        if e.errno != errno.EEXIST:
+            raise
+
+
+def download(data_dir, yaml_path, retries, overwrite=False):
+    """Download the data files specified in YAML file to a directory.
+
+    Return false if the downloaded file or the local copy (if not overwrite) has a different checksum.
+    """
+    sample_data = _loadYAML(yaml_path)
+    logger.info("Downloading data for %s", sample_data.sample)
+
+    def _downloadFile(path, url, retries):
+        logger.info("Downloading %s from %s", path, url)
+        import requests
+        from requests.adapters import HTTPAdapter, Retry
+
+        session = requests.Session()
+        retries = Retry(total=retries, backoff_factor=0.5)
+        session.mount("http://", HTTPAdapter(max_retries=retries))
+        session.mount("https://", HTTPAdapter(max_retries=retries))
+        try:
+            r = session.get(url, stream=True, timeout=60)
+            
+            if r.status_code == 200:
+                logger.info("Connecting to %s is successful.", url)
+
+                size = int(r.headers.get("content-length", 0))
+                from tqdm import tqdm
+
+                progress_bar = tqdm(total=size, unit="iB", unit_scale=True)
+                with open(path, "wb") as fd:
+                    for chunk in r.iter_content(chunk_size=1024):
+                        progress_bar.update(len(chunk))
+                        fd.write(chunk)
+                progress_bar.close()
+                return True
+            else:
+                logger.info("Failed to connect to %s with status code: %s.", url, r.status_code)
+                return False
+        
+        except requests.exceptions.ConnectionError as e:
+            logger.debug("Connection failed after retries:", e)
+        except requests.exceptions.Timeout as e:
+            logger.debug("A timeout occurred:", e)
+        except requests.exceptions.RequestException as e:
+            logger.debug("Error occurred while requesting connection to %s: %s.", url, e)
+        return False
+
+    allGood = True
+    for f in sample_data.files:
+        fpath = os.path.join(data_dir, f.path)
+        if os.path.exists(fpath):
+            if _checkMD5(fpath, f.checksum):
+                logger.info("Found local copy %s, skip downloading.", fpath)
+                continue
+            else:
+                logger.warning("Local copy %s has a different checksum!", fpath)
+                if overwrite:
+                    logging.warning("Removing local copy %s", fpath)
+                    os.remove(fpath)
+                else:
+                    allGood = False
+                    continue
+        _createDirIfNeeded(fpath)
+        assert _downloadFile(fpath, f.url, retries=retries)
+        if not _checkMD5(fpath, f.checksum):
+            logger.error("The downloaded file %s has a different checksum!", fpath)
+            allGood = False
+
+    return allGood
+
+
+def _parseArgs():
+    parser = argparse.ArgumentParser(description="Downloader of TensorRT sample data files.")
+    parser.add_argument(
+        "-d",
+        "--data",
+        help="Specify the data directory, data will be downloaded to there. $TRT_DATA_DIR will be overwritten by this argument.",
+    )
+    parser.add_argument(
+        "-f",
+        "--file",
+        help="Specify the path to the download.yml, default to `download.yml` in the working directory",
+        default="download.yml",
+    )
+    parser.add_argument(
+        "-o",
+        "--overwrite",
+        help="Force to overwrite if MD5 check failed",
+        action="store_true",
+        default=False,
+    )
+    parser.add_argument(
+        "-v",
+        "--verify",
+        help="Verify if the data has been downloaded. Will not download if specified.",
+        action="store_true",
+        default=False,
+    )
+    parser.add_argument(
+        "-r",
+        "--retries",
+        help="Number of retries for download",
+        type=int,
+        default=10,
+    )
+    args, _ = parser.parse_known_args()
+    data = os.environ.get("TRT_DATA_DIR", None) if args.data is None else args.data
+    if data is None:
+        raise ValueError("Data directory must be specified by either `-d $DATA` or environment variable $TRT_DATA_DIR.")
+
+    return data, args
+
+
+def verifyChecksum(data_dir, yaml_path):
+    """Verify the checksum of the files described by the YAML.
+
+    Return false of any of the file doesn't existed or checksum is different with the YAML.
+    """
+    sample_data = _loadYAML(yaml_path)
+    logger.info("Verifying data files and their MD5 for %s", sample_data.sample)
+
+    allGood = True
+    for f in sample_data.files:
+        fpath = os.path.join(data_dir, f.path)
+        if os.path.exists(fpath):
+            if _checkMD5(fpath, f.checksum):
+                logger.info("MD5 match for local copy %s", fpath)
+            else:
+                logger.error("Local file %s has a different checksum!", fpath)
+                allGood = False
+        else:
+            allGood = False
+            logger.error("Data file %s doesn't have a local copy", f.path)
+
+    return allGood
+
+
+def main():
+    data, args = _parseArgs()
+    logging.basicConfig()
+    logger.setLevel(logging.INFO)
+
+    ret = True
+    if args.verify:
+        ret = verifyChecksum(data, args.file)
+    else:
+        ret = download(data, args.file, args.retries, args.overwrite)
+
+    if not ret:
+        # Error of downloading or checksum
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()
+
+
+TRT_DATA_DIR = None
+
+
+def getFilePath(path):
+    """Util to get the full path to the downloaded data files.
+
+    It only works when the sample doesn't have any other command line argument.
+    """
+    global TRT_DATA_DIR
+    if not TRT_DATA_DIR:
+        parser = argparse.ArgumentParser(description="Helper of data file download tool")
+        parser.add_argument(
+            "-d",
+            "--data",
+            help="Specify the data directory where it is saved in. $TRT_DATA_DIR will be overwritten by this argument.",
+        )
+        args, _ = parser.parse_known_args()
+        TRT_DATA_DIR = os.environ.get("TRT_DATA_DIR", None) if args.data is None else args.data
+    if TRT_DATA_DIR is None:
+        raise ValueError("Data directory must be specified by either `-d $DATA` or environment variable $TRT_DATA_DIR.")
+
+    fullpath = os.path.join(TRT_DATA_DIR, path)
+    if not os.path.exists(fullpath):
+        raise ValueError("Data file %s doesn't exist!" % fullpath)
+
+    return fullpath

samples/python/engine_refit_onnx_bidaf/README.md

@@ -0,0 +1,163 @@
+# TensorRT Engine Refitting of ONNX models.
+
+**Table Of Contents**
+- [Description](#description)
+- [How does this sample work?](#how-does-this-sample-work)
+- [Prerequisites](#prerequisites)
+- [Running the sample](#running-the-sample)
+- [Additional resources](#additional-resources)
+- [License](#license)
+- [Changelog](#changelog)
+- [Known issues](#known-issues)
+
+## Description
+
+This sample shows how to refit an engine built from an ONNX model via parsers. A modified version of the [ONNX BiDAF model](https://github.com/onnx/models/tree/main/validated/text/machine_comprehension/bidirectional_attention_flow) is used as the sample model, which implements the Bi-Directional Attention Flow (BiDAF) network described in the paper [Bidirectional Attention Flow for Machine Comprehension](https://arxiv.org/abs/1611.01603).
+
+## How does this sample work?
+
+This sample replaces unsupported nodes (HardMax / Compress) in the original ONNX model via ONNX-graphsurgeon (in `prepare_model.py`) and build a refittable TensorRT engine.
+The engine is then refitted with fake weights and correct weights, each followed by inference on sample context and query sentences in `build_and_refit_engine.py`.
+
+## Prerequisites
+
+Dependencies required for this sample
+
+1. Install the dependencies for Python:
+```bash
+pip3 install -r requirements.txt
+```
+
+2. TensorRT
+
+3. [ONNX-GraphSurgeon](https://github.com/NVIDIA/TensorRT/tree/main/tools/onnx-graphsurgeon)
+
+4. Download sample data. See the "Download Sample Data" section of [the general setup guide](../README.md).
+
+## Running the sample
+
+The data directory needs to be specified (either via `-d /path/to/data` or environment varaiable `TRT_DATA_DIR`)
+when running these scripts. An error will be thrown if not. Taking `TRT_DATA_DIR` approach in following example.
+
+* Prepare the ONNX model. (The data directory needs to be specified.)
+  ```bash
+  python3 prepare_model.py
+  ```
+
+The output should look similar to the following:
+```
+Modifying the ONNX model ...
+Modified ONNX model saved as bidaf-modified.onnx
+Done.
+```
+
+The script will modify the original model from [onnx/models](https://github.com/onnx/models/raw/c02f8c8699fc12273649e658b8d2a1a8e32a35d0/text/machine_comprehension/bidirectional_attention_flow/model/bidaf-9.onnx) and save an ONNX model that can be parsed and run by TensorRT.
+
+The original ONNX model contains four CategoryMapper nodes to map the four input string arrays to int arrays.
+Since TensorRT does not support string data type and CategoryMapper nodes, we dump out the four maps for the four nodes as json files (`model/CategoryMapper_{4-6}.json`) and use them to preprocess input data.
+Now the four inputs become four outputs of the original CategoryMapper nodes.
+
+And unsupported HardMax nodes and Compress nodes are replaced by ArgMax nodes and Gather nodes, respectively.
+
+
+* Build a TensorRT engine, refit the engine and run inference.
+`python3 build_and_refit_engine.py --weights-location GPU`
+
+The script will build a TensorRT engine from the modified ONNX model, and then refit the engine from GPU weights and run inference on sample context and query sentences.
+
+When running the above command for the first time, the output should look similar to the following:
+```
+Loading ONNX file from path bidaf-modified.onnx...
+Beginning ONNX file parsing
+[09/25/2023-08:48:16] [TRT] [W] ModelImporter.cpp:407: Make sure input CategoryMapper_4 has Int64 binding.
+[09/25/2023-08:48:16] [TRT] [W] ModelImporter.cpp:407: Make sure input CategoryMapper_5 has Int64 binding.
+[09/25/2023-08:48:16] [TRT] [W] ModelImporter.cpp:407: Make sure input CategoryMapper_6 has Int64 binding.
+[09/25/2023-08:48:16] [TRT] [W] ModelImporter.cpp:407: Make sure input CategoryMapper_7 has Int64 binding.
+Completed parsing of ONNX file
+Network inputs:
+CategoryMapper_4 <class 'numpy.int64'> (-1, 1)
+CategoryMapper_5 <class 'numpy.int64'> (-1, 1, 1, 16)
+CategoryMapper_6 <class 'numpy.int64'> (-1, 1)
+CategoryMapper_7 <class 'numpy.int64'> (-1, 1, 1, 16)
+Building an engine from file bidaf-modified.onnx; this may take a while...
+Completed creating Engine
+Refitting engine from GPU weights...
+Engine refitted in 39.88 ms.
+Doing inference...
+Doing inference...
+Refitting engine from GPU weights...
+Engine refitted in 0.27 ms.
+Doing inference...
+Doing inference...
+Passed
+```
+
+Note that refitting for second time will be much faster than the first time.
+When running the above command again, engine will be deserialized from the plan file, the output should look similar to the following:
+```
+Reading engine from file bidaf.trt...
+Refitting engine from GPU weights...
+Engine refitted in 32.64 ms.
+Doing inference...
+Doing inference...
+Refitting engine from GPU weights...
+Engine refitted in 0.41 ms.
+Doing inference...
+Doing inference...
+Passed
+```
+
+To refit the engine from CPU weights, change the command to be `python3 build_and_refit_engine.py --weights-location CPU`. And the output should look similar to the following
+```
+Reading engine from file bidaf.trt...
+Refitting engine from CPU weights...
+Engine refitted in 45.18 ms.
+Doing inference...
+Doing inference...
+Refitting engine from CPU weights...
+Engine refitted in 1.20 ms.
+Doing inference...
+Doing inference...
+Passed
+```
+
+There is also an option `--version-compatible` to enable engine version compatibility. If installed, `tensorrt_dispatch` package will used for refitting and running version compatible engines instead of `tensorrt` package.
+To build and refit a version compatible engine, run the command `python3 build_and_refit_engine.py --version-compatible` and the output should look similar to the above cases.
+
+# Additional resources
+
+The following resources provide a deeper understanding about the model used in this sample:
+
+**Model**
+- [Bidirectional Attention Flow for Machine Comprehension](https://arxiv.org/abs/1611.01603)
+
+**Documentation**
+- [Introduction To NVIDIA’s TensorRT Samples](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sample-support-guide/index.html#samples)
+- [Working With TensorRT Using The Python API](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#python_topics)
+- [Importing A Model Using A Parser In Python](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#import_model_python)
+- [NVIDIA’s TensorRT Documentation Library](https://docs.nvidia.com/deeplearning/sdk/tensorrt-archived/index.html)
+
+# License
+
+For terms and conditions for use, reproduction, and distribution, see the [TensorRT Software License Agreement](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sla/index.html) documentation.
+
+# Changelog
+
+October 2025
+  - Migrate to strongly typed APIs.
+
+August 2025:
+  - Removed support for Python versions < 3.10.
+
+January 2024:
+  - Add support for refitting version compatible engines.
+
+August 2023:
+  - Add support for refitting engines from GPU weights.
+  - Removed support for Python versions < 3.8.
+
+October 2020: This sample was recreated, updated and reviewed.
+
+# Known issues
+
+There are no known issues in this sample.

samples/python/engine_refit_onnx_bidaf/build_and_refit_engine.py

@@ -0,0 +1,298 @@
+#!/usr/bin/env python3
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2025 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 sys
+
+import numpy as np
+import argparse
+import tensorrt as trt
+
+sys.path.insert(1, os.path.join(sys.path[0], os.path.pardir))
+from cuda.bindings import runtime as cudart
+
+TRT_LOGGER = trt.Logger()
+
+
+def get_plan(onnx_file_path, engine_file_path, version_compatible):
+    """Attempts to load a serialized engine if available, otherwise builds a new TensorRT engine and saves it."""
+
+    def build_plan():
+        """Takes an ONNX file and creates a TensorRT engine to run inference with"""
+        import tensorrt as trt
+
+        builder = trt.Builder(TRT_LOGGER)
+        network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED))
+        parser = trt.OnnxParser(network, TRT_LOGGER)
+
+        # Parse model file
+        print("Loading ONNX file from path {}...".format(onnx_file_path))
+        with open(onnx_file_path, "rb") as model:
+            print("Beginning ONNX file parsing")
+            if not parser.parse(model.read()):
+                print("ERROR: Failed to parse the ONNX file.")
+                for error in range(parser.num_errors):
+                    print(parser.get_error(error))
+                return None
+        print("Completed parsing of ONNX file")
+
+        # Print input info
+        print("Network inputs:")
+        for i in range(network.num_inputs):
+            tensor = network.get_input(i)
+            print(tensor.name, trt.nptype(tensor.dtype), tensor.shape)
+
+        config = builder.create_builder_config()
+        config.set_flag(trt.BuilderFlag.REFIT)
+        if version_compatible:
+            config.set_flag(trt.BuilderFlag.VERSION_COMPATIBLE)
+
+        for opt in [6, 10]:
+            profile = builder.create_optimization_profile()
+
+            input0_min = (1, 1)
+            input0_opt = (opt, 1)
+            input0_max = (15, 1)
+            profile.set_shape(
+                network.get_input(0).name,
+                min=input0_min,
+                opt=input0_opt,
+                max=input0_max,
+            )
+
+            input1_min = (1, 1, 1, 16)
+            input1_opt = (opt, 1, 1, 16)
+            input1_max = (15, 1, 1, 16)
+            profile.set_shape(
+                network.get_input(1).name,
+                min=input1_min,
+                opt=input1_opt,
+                max=input1_max,
+            )
+
+            input2_min = (1, 1)
+            input2_opt = (opt, 1)
+            input2_max = (15, 1)
+            profile.set_shape(
+                network.get_input(2).name,
+                min=input2_min,
+                opt=input2_opt,
+                max=input2_max,
+            )
+
+            input3_min = (1, 1, 1, 16)
+            input3_opt = (opt, 1, 1, 16)
+            input3_max = (15, 1, 1, 16)
+            profile.set_shape(
+                network.get_input(3).name,
+                min=input3_min,
+                opt=input3_opt,
+                max=input3_max,
+            )
+
+            config.add_optimization_profile(profile)
+
+        print(
+            "Building an engine from file {}; this may take a while...".format(
+                onnx_file_path
+            )
+        )
+        plan = builder.build_serialized_network(network, config)
+        print("Completed creating Engine")
+
+        with open(engine_file_path, "wb") as f:
+            f.write(plan)
+        return plan
+
+    if os.path.exists(engine_file_path):
+        # If a serialized engine exists, use it instead of building an engine.
+        print("Reading engine from file {}...".format(engine_file_path))
+        f = open(engine_file_path, "rb")
+        return f.read()
+    return build_plan()
+
+
+def main():
+    global trt
+    global TRT_LOGGER
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-l",
+        "--weights-location",
+        dest="weights_location",
+        default="GPU",
+        choices=["GPU", "CPU"],
+        help="The location for weights passed to refitter, either GPU/CPU, default: GPU",
+    )
+    parser.add_argument(
+        "--version-compatible",
+        dest="version_compatible",
+        action="store_true",
+        help="Build a version compatible engine for refitting",
+    )
+    args = parser.parse_args()
+
+    onnx_file_path = "bidaf-modified.onnx"
+    engine_file_path = "bidaf{}.trt".format("-vc" if args.version_compatible else "")
+
+    plan = get_plan(onnx_file_path, engine_file_path, args.version_compatible)
+
+    if args.version_compatible:
+        # Try using dispatch runtime for refitting and inference. If failed, fallback to full runtime.
+        try:
+            del sys.modules["tensorrt"]
+            sys.modules["tensorrt"] = __import__("tensorrt_dispatch")
+            sys.modules["trt"] = sys.modules["tensorrt"]
+            import tensorrt_dispatch as trt
+
+            print(
+                "Importing tensorrt_dispatch instead of full tensorrt for refitting and running vc engines."
+            )
+        except:
+            print(
+                "Failed to import tensorrt_dispatch for refitting and running vc engines. Please install the package first!"
+            )
+            sys.modules["tensorrt"] = __import__("tensorrt")
+        TRT_LOGGER = trt.Logger()
+
+    engine = None
+    with open(engine_file_path, "rb") as f:
+        runtime = trt.Runtime(TRT_LOGGER)
+        if args.version_compatible:
+            runtime.engine_host_code_allowed = True
+        engine = runtime.deserialize_cuda_engine(plan)
+
+    # should be after get_engine
+    from data_processing import get_inputs, preprocess
+    import common_runtime as common
+
+    # input
+    context = "A quick brown fox jumps over the lazy dog."
+    query = "What color is the fox?"
+    cw_str, _ = preprocess(context)
+    # get ravelled data
+    cw, cc, qw, qc = get_inputs(context, query)
+
+    # Do inference with TensorRT
+    weights_names = ["Parameter576_B_0", "W_0"]
+    refit_weights_dict = {
+        name: np.load("{}.npy".format(name)) for name in weights_names
+    }
+    fake_weights_dict = {
+        name: np.ones_like(weights) for name, weights in refit_weights_dict.items()
+    }
+    device_mem_dict = {}
+    if args.weights_location == "GPU":
+        for name, weights in refit_weights_dict.items():
+            nbytes = weights.size * weights.itemsize
+            device_mem_dict[name] = common.DeviceMem(nbytes)
+
+    execution_context = engine.create_execution_context()
+    refitter = trt.Refitter(engine, TRT_LOGGER)
+    # Skip weights validation since we are confident that the new weights are similar to the weights used to build engine.
+    refitter.weights_validation = False
+    # To get a list of all refittable weights' names
+    # in the network, use refitter.get_all_weights().
+
+    if args.weights_location == "GPU":
+        for name, device_mem in device_mem_dict.items():
+            device_weights = trt.Weights(
+                trt.DataType.FLOAT, device_mem.device_ptr, refit_weights_dict[name].size
+            )
+            weights_prototype = refitter.get_weights_prototype(name)
+            assert device_weights.dtype == weights_prototype.dtype
+            assert device_weights.size == weights_prototype.size
+            refitter.set_named_weights(name, device_weights, trt.TensorLocation.DEVICE)
+
+    for weights_dict, answer_correct in [
+        (fake_weights_dict, False),
+        (refit_weights_dict, True),
+    ]:
+        import time
+
+        T1 = time.perf_counter()
+        device_mem_list = []
+        # Refit named weights via set_named_weights
+        for name in weights_names:
+            host_weights = weights_dict[name]
+            if args.weights_location == "CPU":
+                weights = host_weights
+                location = trt.TensorLocation.HOST
+                refitter.set_named_weights(name, weights, location)
+            else:
+                common.memcpy_host_to_device(device_mem_dict[name].device_ptr, host_weights)
+
+        # Get missing weights names. This should return empty lists in this case.
+        missing_weights = refitter.get_missing_weights()
+        assert (
+            len(missing_weights) == 0
+        ), "Refitter found missing weights. Call set_named_weights() or set_weights() for all missing weights"
+
+        print(f"Refitting engine from {args.weights_location} weights...")
+        # Refit the engine with the new weights. This will return True if the refit operation succeeded.
+        assert refitter.refit_cuda_engine()
+
+        T2 = time.perf_counter()
+        print("Engine refitted in {:.2f} ms.".format((T2 - T1) * 1000))
+
+        for profile_idx in range(engine.num_optimization_profiles):
+            print("Doing inference...")
+            # Do inference
+            inputs, outputs, bindings = common.allocate_buffers(
+                engine, profile_idx
+            )
+            padding_bindings = [0] * (len(bindings) * profile_idx)
+            new_bindings = padding_bindings + bindings
+
+            # Use context manager for proper stream lifecycle management
+            with common.CudaStreamContext() as stream:
+                # Set host input. The common.do_inference function will copy the input to the GPU before executing.
+                inputs[0].host = cw
+                inputs[1].host = cc
+                inputs[2].host = qw
+                inputs[3].host = qc
+                execution_context.set_optimization_profile_async(profile_idx, stream.stream)
+                execution_context.set_input_shape("CategoryMapper_4", (10, 1))
+                execution_context.set_input_shape("CategoryMapper_5", (10, 1, 1, 16))
+                execution_context.set_input_shape("CategoryMapper_6", (6, 1))
+                execution_context.set_input_shape("CategoryMapper_7", (6, 1, 1, 16))
+
+                trt_outputs = common.do_inference(
+                    execution_context,
+                    engine=engine,
+                    bindings=bindings,
+                    inputs=inputs,
+                    outputs=outputs,
+                    stream=stream,
+                )
+
+            start = trt_outputs[0].item()
+            end = trt_outputs[1].item()
+            answer = [w.encode() for w in cw_str[start : end + 1].reshape(-1)]
+            assert answer_correct == (answer == [b"brown"]), answer
+            common.free_buffers(inputs, outputs)
+
+    for _, device_mem in device_mem_dict.items():
+        device_mem.free()
+
+    print("Passed")
+
+
+if __name__ == "__main__":
+    main()

samples/python/engine_refit_onnx_bidaf/data_processing.py

@@ -0,0 +1,64 @@
+#
+# 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 nltk
+from nltk import word_tokenize
+import json
+import tensorrt as trt
+
+
+def preprocess(text):
+    try:
+        nltk.data.find("tokenizers/punkt_tab")
+    except LookupError:
+        nltk.download("punkt_tab")
+    tokens = word_tokenize(text)
+    # split into lower-case word tokens, in numpy array with shape of (seq, 1)
+    words = np.asarray([w.lower() for w in tokens]).reshape(-1, 1)
+    # split words into chars, in numpy array with shape of (seq, 1, 1, 16)
+    chars = [[c for c in t][:16] for t in tokens]
+    chars = [cs + [""] * (16 - len(cs)) for cs in chars]
+    chars = np.asarray(chars).reshape(-1, 1, 1, 16)
+    return words, chars
+
+
+def get_map_func(filepath):
+    file = open(filepath)
+    category_map = json.load(file)
+    category_mapper = dict(
+        zip(category_map["cats_strings"], category_map["cats_int64s"])
+    )
+    default_int64 = category_map["default_int64"]
+    func = lambda s: category_mapper.get(s, default_int64)
+    return np.vectorize(func)
+
+
+def get_inputs(context, query):
+    cw, cc = preprocess(context)
+    qw, qc = preprocess(query)
+
+    context_word_func = get_map_func("CategoryMapper_4.json")
+    context_char_func = get_map_func("CategoryMapper_5.json")
+    query_word_func = get_map_func("CategoryMapper_6.json")
+    query_char_func = get_map_func("CategoryMapper_7.json")
+
+    cw_input = context_word_func(cw).astype(trt.nptype(trt.int32)).ravel()
+    cc_input = context_char_func(cc).astype(trt.nptype(trt.int32)).ravel()
+    qw_input = query_word_func(qw).astype(trt.nptype(trt.int32)).ravel()
+    qc_input = query_char_func(qc).astype(trt.nptype(trt.int32)).ravel()
+    return cw_input, cc_input, qw_input, qc_input

samples/python/engine_refit_onnx_bidaf/download.yml

@@ -0,0 +1,21 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 2020-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.
+#
+sample: engine_refit_onnx_bidaf
+files:
+  - path: samples/python/engine_refit_onnx_bidaf/bidaf-original.onnx
+    url: https://github.com/onnx/models/raw/c02f8c8699fc12273649e658b8d2a1a8e32a35d0/text/machine_comprehension/bidirectional_attention_flow/model/bidaf-9.onnx
+    checksum: cf11f1eceb4731f8dd39345467fe94a1

samples/python/engine_refit_onnx_bidaf/prepare_model.py

@@ -0,0 +1,107 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2025 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 onnx_graphsurgeon as gs
+import onnx
+import numpy as np
+import json
+
+import sys, os
+
+sys.path.insert(1, os.path.join(sys.path[0], os.path.pardir))
+from downloader import getFilePath
+
+
+def drop_category_mapper_nodes(graph):
+    new_inputs = []
+    for org_input in graph.inputs:
+        # head node, simply disconnect it with others
+        assert len(org_input.outputs) == 1
+        category_mapper_node = org_input.outputs[0]
+        assert category_mapper_node.op == "CategoryMapper"
+        assert len(category_mapper_node.outputs) == 1
+        new_inputs.append(category_mapper_node.outputs[0])
+        category_mapper_node.inputs.clear()
+        category_mapper_node.outputs.clear()
+
+        # Save mapping info to preprocess inputs.
+        with open(category_mapper_node.name + ".json", "w") as fp:
+            json.dump(category_mapper_node.attrs, fp)
+
+    graph.inputs = new_inputs
+
+
+def replace_unsupported_ops(graph):
+    # replace hardmax with ArgMax
+    hardmaxes = [node for node in graph.nodes if node.op == "Hardmax"]
+    assert len(hardmaxes) == 1
+    hardmax = hardmaxes[0]
+    hardmax.op = "ArgMax"
+    hardmax.name = "ArgMax(org:" + hardmax.name + ")"
+    hardmax.attrs["axis"] = 1
+    hardmax.attrs["keepdims"] = 0
+
+    cast = hardmax.o()
+    reshape = cast.o()
+
+    hardmax.outputs = reshape.outputs
+    assert len(hardmax.outputs) == 1
+    hardmax.outputs[0].dtype = np.int64
+    hardmax.outputs[0].shape = [1]
+
+    compress = reshape.o()
+    compress.op = "Gather"
+    compress.name = "Gather(org:" + compress.name + ")"
+    compress.attrs["axis"] = 1
+
+    cast.outputs.clear()
+    reshape.outputs.clear()
+    # Remove the node from the graph completely
+    graph.cleanup().toposort()
+
+
+def save_weights_for_refitting(graph):
+    # Save weights for refitting
+    tmap = graph.tensors()
+    np.save("Parameter576_B_0.npy", tmap["Parameter576_B_0"].values)
+    np.save("W_0.npy", tmap["W_0"].values)
+
+
+def main():
+    org_model_file_path = getFilePath(
+        "samples/python/engine_refit_onnx_bidaf/bidaf-original.onnx"
+    )
+
+    print("Modifying the ONNX model ...")
+    original_model = onnx.load(org_model_file_path)
+    graph = gs.import_onnx(original_model)
+
+    drop_category_mapper_nodes(graph)
+    replace_unsupported_ops(graph)
+    save_weights_for_refitting(graph)
+
+    new_model = gs.export_onnx(graph)
+
+    modified_model_name = "bidaf-modified.onnx"
+    onnx.checker.check_model(new_model)
+    onnx.save(new_model, modified_model_name)
+    print("Modified ONNX model saved as {}".format(modified_model_name))
+    print("Done.")
+
+
+if __name__ == "__main__":
+    main()

samples/python/engine_refit_onnx_bidaf/requirements.txt

@@ -0,0 +1,9 @@
+onnx==1.18.0
+nltk==3.9.1
+wget==3.2
+cuda-python==12.9.0
+pywin32; platform_system == "Windows"
+pyyaml==6.0.3
+requests==2.32.4
+tqdm==4.66.4
+numpy==1.26.4

samples/python/introductory_parser_samples/README.md

@@ -0,0 +1,101 @@
+# Introduction To Importing ONNX Models Into TensorRT Using Python
+
+
+**Table Of Contents**
+- [Description](#description)
+- [How does this sample work?](#how-does-this-sample-work)
+	* [onnx_resnet50](#onnx_resnet50)
+- [Prerequisites](#prerequisites)
+- [Running the sample](#running-the-sample)
+	* [Sample `--help` options](#sample-help-options)
+- [Additional resources](#additional-resources)
+- [License](#license)
+- [Changelog](#changelog)
+- [Known issues](#known-issues)
+
+## Description
+
+This sample, introductory_parser_samples, is a Python sample which uses TensorRT and its included ONNX parser, to perform inference with ResNet-50 models saved in ONNX format.
+
+## How does this sample work?
+
+### onnx_resnet50
+
+This sample demonstrates how to build an engine from an ONNX model file using the open-source ONNX parser and then run inference. The ONNX parser can be used with any framework that supports the ONNX format (typically `.onnx` files).
+
+## Prerequisites
+
+1. Install the dependencies for Python.
+
+```bash
+pip3 install -r requirements.txt
+```
+
+2. Preparing sample data
+
+See [Preparing sample data](../../README.md#preparing-sample-data) in the main samples README.
+
+## Running the sample
+
+1.  Run the sample to create a TensorRT inference engine and run inference:
+	`python3 onnx_resnet50.py`
+
+	**Note:** If the TensorRT sample data is not installed in the default location, the `data` directory must be specified. For example: `python3 onnx_resnet50.py -d $TRT_DATADIR`
+
+2.  Verify that the sample ran successfully. If the sample runs successfully you should see output similar to the following:
+	`Correctly recognized data/samples/resnet50/reflex_camera.jpeg as reflex camera`
+
+### Sample --help options
+
+To see the full list of available options and their descriptions, use the `-h` or `--help` command line option. For example:
+```
+usage: onnx_resnet50.py [-h] [-d DATADIR]
+
+Runs a ResNet50 network with a TensorRT inference engine.
+
+optional arguments:
+ -h, --help            show this help message and exit
+ -d DATADIR, --datadir DATADIR
+                       Location of the TensorRT sample data directory.
+                       (default: /usr/src/tensorrt/data)
+```
+
+# Additional resources
+
+The following resources provide a deeper understanding about importing a model into TensorRT using Python:
+
+**ResNet-50**
+- [Deep Residual Learning for Image Recognition](https://arxiv.org/pdf/1512.03385.pdf)
+
+**Parsers**
+- [ONNX Parser](https://docs.nvidia.com/deeplearning/sdk/tensorrt-api/python_api/parsers/Onnx/pyOnnx.html)
+
+**Documentation**
+- [Introduction To NVIDIA’s TensorRT Samples](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sample-support-guide/index.html#samples)
+- [Working With TensorRT Using The Python API](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#python_topics)
+- [Importing A Model Using A Parser In Python](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#import_model_python)
+- [NVIDIA’s TensorRT Documentation Library](https://docs.nvidia.com/deeplearning/sdk/tensorrt-archived/index.html)
+
+# License
+
+For terms and conditions for use, reproduction, and distribution, see the [TensorRT Software License Agreement](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sla/index.html) documentation.
+
+# Changelog
+October 2025
+Migrate to strongly typed APIs.
+
+August 2025
+Removed support for Python versions < 3.10.
+
+August 2023
+Removed support for Python versions < 3.8.
+
+August 2022
+Removed options for Caffe and UFF parsers.
+
+February 2019
+This `README.md` file was recreated, updated and reviewed.
+
+# Known issues
+
+There are no known issues in this sample

samples/python/introductory_parser_samples/onnx_resnet50.py

@@ -0,0 +1,135 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2025 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
+
+# This sample uses an ONNX ResNet50 Model to create a TensorRT Inference Engine
+import random
+import sys
+
+import numpy as np
+
+import tensorrt as trt
+from PIL import Image
+
+sys.path.insert(1, os.path.join(sys.path[0], os.path.pardir))
+import common
+
+
+class ModelData(object):
+    MODEL_PATH = "ResNet50.onnx"
+    INPUT_SHAPE = (3, 224, 224)
+    # We can convert TensorRT data types to numpy types with trt.nptype()
+    DTYPE = trt.float32
+
+
+# You can set the logger severity higher to suppress messages (or lower to display more messages).
+TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
+
+
+# The Onnx path is used for Onnx models.
+def build_engine_onnx(model_file):
+    builder = trt.Builder(TRT_LOGGER)
+    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED))
+    config = builder.create_builder_config()
+    parser = trt.OnnxParser(network, TRT_LOGGER)
+
+    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, common.GiB(1))
+    # Load the Onnx model and parse it in order to populate the TensorRT network.
+    with open(model_file, "rb") as model:
+        if not parser.parse(model.read()):
+            print("ERROR: Failed to parse the ONNX file.")
+            for error in range(parser.num_errors):
+                print(parser.get_error(error))
+            return None
+
+    engine_bytes = builder.build_serialized_network(network, config)
+    runtime = trt.Runtime(TRT_LOGGER)
+    return runtime.deserialize_cuda_engine(engine_bytes)
+
+
+def load_normalized_test_case(test_image, pagelocked_buffer):
+    # Converts the input image to a CHW Numpy array
+    def normalize_image(image):
+        # Resize, antialias and transpose the image to CHW.
+        c, h, w = ModelData.INPUT_SHAPE
+        image_arr = (
+            np.asarray(image.resize((w, h), Image.LANCZOS))
+            .transpose([2, 0, 1])
+            .astype(trt.nptype(ModelData.DTYPE))
+            .ravel()
+        )
+        # This particular ResNet50 model requires some preprocessing, specifically, mean normalization.
+        return (image_arr / 255.0 - 0.45) / 0.225
+
+    # Normalize the image and copy to pagelocked memory.
+    np.copyto(pagelocked_buffer, normalize_image(Image.open(test_image)))
+    return test_image
+
+
+def main():
+    # Set the data path to the directory that contains the trained models and test images for inference.
+    _, data_files = common.find_sample_data(
+        description="Runs a ResNet50 network with a TensorRT inference engine.",
+        subfolder="resnet50",
+        find_files=[
+            "binoculars.jpeg",
+            "reflex_camera.jpeg",
+            "tabby_tiger_cat.jpg",
+            ModelData.MODEL_PATH,
+            "class_labels.txt",
+        ],
+    )
+    # Get test images, models and labels.
+    test_images = data_files[0:3]
+    onnx_model_file, labels_file = data_files[3:]
+    labels = open(labels_file, "r").read().split("\n")
+
+    # Build a TensorRT engine.
+    engine = build_engine_onnx(onnx_model_file)
+    # Inference is the same regardless of which parser is used to build the engine, since the model architecture is the same.
+    # Allocate buffers
+    inputs, outputs, bindings = common.allocate_buffers(engine)
+    # Contexts are used to perform inference.
+    context = engine.create_execution_context()
+
+    # Use context manager for proper stream lifecycle management
+    with common.CudaStreamContext() as stream:
+        # Load a normalized test case into the host input page-locked buffer.
+        test_image = random.choice(test_images)
+        test_case = load_normalized_test_case(test_image, inputs[0].host)
+        # Run the engine. The output will be a 1D tensor of length 1000, where each value represents the
+        # probability that the image corresponds to that label
+        trt_outputs = common.do_inference(
+            context,
+            engine=engine,
+            bindings=bindings,
+            inputs=inputs,
+            outputs=outputs,
+            stream=stream,
+        )
+        # We use the highest probability as our prediction. Its index corresponds to the predicted label.
+        pred = labels[np.argmax(trt_outputs[0])]
+        common.free_buffers(inputs, outputs)
+    if "_".join(pred.split()) in os.path.splitext(os.path.basename(test_case))[0]:
+        print("Correctly recognized " + test_case + " as " + pred)
+    else:
+        print("Incorrectly recognized " + test_case + " as " + pred)
+
+
+if __name__ == "__main__":
+    main()

samples/python/introductory_parser_samples/requirements.txt

@@ -0,0 +1,7 @@
+Pillow==11.3.0
+cuda-python==12.9.0
+pywin32; platform_system == "Windows"
+pyyaml==6.0.3
+requests==2.32.4
+tqdm==4.66.4
+numpy==1.26.4

samples/python/network_api_pytorch_mnist/README.md

@@ -0,0 +1,118 @@
+# “Hello World” For TensorRT Using PyTorch And Python
+
+**Table Of Contents**
+
+- [Description](#description)
+- [How does this sample work?](#how-does-this-sample-work)
+    * [TensorRT API layers and ops](#tensorrt-api-layers-and-ops)
+- [Prerequisites](#prerequisites)
+- [Running the sample](#running-the-sample)
+    * [Sample `--help` options](#sample-help-options)
+- [Additional resources](#additional-resources)
+- [License](#license)
+- [Changelog](#changelog)
+- [Known issues](#known-issues)
+
+## Description
+
+This sample, `network_api_pytorch_mnist`, trains a convolutional model on the [MNIST](https://ossci-datasets.s3.amazonaws.com/mnist/) dataset and runs inference with a TensorRT engine.
+
+## How does this sample work?
+
+This sample is an end-to-end sample that trains a model in PyTorch, recreates the network in TensorRT, imports weights from the trained model, and finally runs inference with a TensorRT engine. For more information, see [Creating A Network Definition In Python](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#network_python).
+
+The `sample.py` script imports the functions from the `mnist.py` script for training the PyTorch model, as well as retrieving test cases from the PyTorch Data Loader.
+
+### TensorRT API layers and ops
+
+In this sample, the following layers are used. For more information about these layers, see the [TensorRT Developer Guide: Layers](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#layers) documentation.
+
+[Activation layer](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#activation-layer)
+The Activation layer implements element-wise activation functions. Specifically, this sample uses the Activation layer with the type `RELU`.
+
+[Convolution layer](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#convolution-layer)
+The Convolution layer computes a 2D (channel, height, and width) convolution, with or without bias.
+
+[MatrixMultiplyLayer](https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#matrixmultiply-layer)
+The MatrixMultiply layer implements a matrix multiplication.
+(The [FullyConnected layer](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#fullyconnected-layer) is deprecated since 8.4.
+The bias of FullyConnected semantic can be added with an
+[ElementwiseLayer](https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#elementwise-layer) of `SUM` operation.)
+
+[Pooling layer](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#pooling-layer)
+The Pooling layer implements pooling within a channel. Supported pooling types are `maximum`, `average` and `maximum-average blend`.
+
+## Prerequisites
+
+1. Upgrade pip version and install the sample dependencies.
+    ```bash
+    pip3 install --upgrade pip
+    pip3 install -r requirements.txt
+    ```
+
+To run this sample you must be using Python 3.6 or newer.
+
+On PowerPC systems, you will need to manually install PyTorch using IBM's [PowerAI](https://www.ibm.com/support/knowledgecenter/SS5SF7_1.6.0/navigation/pai_install.htm).
+
+2. Preparing sample data
+
+See [Preparing sample data](../../README.md#preparing-sample-data) in the main samples README.
+
+The MNIST dataset can be found under `$TRT_DATADIR/mnist`.
+
+## Running the sample
+
+1.  Run the sample to create a TensorRT inference engine and run inference:
+    `python3 sample.py`
+
+2.  Verify that the sample ran successfully. If the sample runs successfully you should see a match between the test case and the prediction.
+     ```
+    Test Case: 0
+    Prediction: 0
+    ```
+
+### Sample --help options
+
+To see the full list of available options and their descriptions, use the `-h` or `--help` command line option.
+
+# Additional resources
+
+The following resources provide a deeper understanding about getting started with TensorRT using Python:
+
+**Model**
+- [MNIST model](https://github.com/pytorch/examples/tree/master/mnist)
+
+**Dataset**
+- [MNIST database](https://ossci-datasets.s3.amazonaws.com/mnist/)
+
+**Documentation**
+- [Introduction To NVIDIA’s TensorRT Samples](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sample-support-guide/index.html#samples)
+- [Working With TensorRT Using The Python API](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#python_topics)
+- [NVIDIA’s TensorRT Documentation Library](https://docs.nvidia.com/deeplearning/sdk/tensorrt-archived/index.html)
+
+# License
+
+For terms and conditions for use, reproduction, and distribution, see the [TensorRT Software License Agreement](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sla/index.html) documentation.
+
+# Changelog
+October 2025
+Migrate to strongly typed APIs.
+
+August 2025
+Removed support for Python versions < 3.10.
+
+August 2023
+Removed support for Python versions < 3.8.
+
+September 2021
+Updated the sample to use explicit batch network definition.
+
+March 2021
+Documented the Python version limitations.
+
+February 2019
+This `README.md` file was recreated, updated and reviewed.
+
+# Known issues
+
+This sample only supports Python 3.6+ due to `torch` and `torchvision` version requirements.

samples/python/network_api_pytorch_mnist/model.py

@@ -0,0 +1,156 @@
+#
+# 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.
+#
+
+# This file contains functions for training a PyTorch MNIST Model
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torchvision import datasets, transforms
+from torch.autograd import Variable
+
+import numpy as np
+
+from random import randint
+
+
+# Network
+class Net(nn.Module):
+    def __init__(self):
+        super(Net, self).__init__()
+        self.conv1 = nn.Conv2d(1, 20, kernel_size=5)
+        self.conv2 = nn.Conv2d(20, 50, kernel_size=5)
+        self.fc1 = nn.Linear(800, 500)
+        self.fc2 = nn.Linear(500, 10)
+
+    def forward(self, x):
+        x = F.max_pool2d(self.conv1(x), kernel_size=2, stride=2)
+        x = F.max_pool2d(self.conv2(x), kernel_size=2, stride=2)
+        x = x.view(-1, 800)
+        x = F.relu(self.fc1(x))
+        x = self.fc2(x)
+        return F.log_softmax(x, dim=1)
+
+
+class MnistModel(object):
+    def __init__(self):
+        self.batch_size = 64
+        self.test_batch_size = 100
+        self.learning_rate = 0.0025
+        self.sgd_momentum = 0.9
+        self.log_interval = 100
+        # Fetch MNIST data set.
+        self.train_loader = torch.utils.data.DataLoader(
+            datasets.MNIST(
+                "/tmp/mnist/data",
+                train=True,
+                download=True,
+                transform=transforms.Compose(
+                    [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
+                ),
+            ),
+            batch_size=self.batch_size,
+            shuffle=True,
+            num_workers=1,
+            timeout=600,
+        )
+        self.test_loader = torch.utils.data.DataLoader(
+            datasets.MNIST(
+                "/tmp/mnist/data",
+                train=False,
+                transform=transforms.Compose(
+                    [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
+                ),
+            ),
+            batch_size=self.test_batch_size,
+            shuffle=True,
+            num_workers=1,
+            timeout=600,
+        )
+        self.network = Net()
+        if torch.cuda.is_available():
+            self.network = self.network.to("cuda")
+
+    # Train the network for one or more epochs, validating after each epoch.
+    def learn(self, num_epochs=2):
+        # Train the network for a single epoch
+        def train(epoch):
+            self.network.train()
+            optimizer = optim.SGD(
+                self.network.parameters(),
+                lr=self.learning_rate,
+                momentum=self.sgd_momentum,
+            )
+            for batch, (data, target) in enumerate(self.train_loader):
+                if torch.cuda.is_available():
+                    data = data.to("cuda")
+                    target = target.to("cuda")
+                data, target = Variable(data), Variable(target)
+                optimizer.zero_grad()
+                output = self.network(data)
+                loss = F.nll_loss(output, target)
+                loss.backward()
+                optimizer.step()
+                if batch % self.log_interval == 0:
+                    print(
+                        "Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}".format(
+                            epoch,
+                            batch * len(data),
+                            len(self.train_loader.dataset),
+                            100.0 * batch / len(self.train_loader),
+                            loss.data.item(),
+                        )
+                    )
+
+        # Test the network
+        def test(epoch):
+            self.network.eval()
+            test_loss = 0
+            correct = 0
+            for data, target in self.test_loader:
+                with torch.no_grad():
+                    if torch.cuda.is_available():
+                        data = data.to("cuda")
+                        target = target.to("cuda")
+                    data, target = Variable(data), Variable(target)
+                output = self.network(data)
+                test_loss += F.nll_loss(output, target).data.item()
+                pred = output.data.max(1)[1]
+                correct += pred.eq(target.data).cpu().sum()
+            test_loss /= len(self.test_loader)
+            print(
+                "\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n".format(
+                    test_loss,
+                    correct,
+                    len(self.test_loader.dataset),
+                    100.0 * correct / len(self.test_loader.dataset),
+                )
+            )
+
+        for e in range(num_epochs):
+            train(e + 1)
+            test(e + 1)
+
+    def get_weights(self):
+        return self.network.state_dict()
+
+    def get_random_testcase(self):
+        data, target = next(iter(self.test_loader))
+        case_num = randint(0, len(data) - 1)
+        test_case = data.cpu().numpy()[case_num].ravel().astype(np.float32)
+        test_name = target.cpu().numpy()[case_num]
+        return test_case, test_name

samples/python/network_api_pytorch_mnist/requirements.txt

@@ -0,0 +1,9 @@
+Pillow==11.3.0
+torch
+torchvision
+cuda-python==12.9.0
+pywin32; platform_system == "Windows"
+pyyaml==6.0.3
+requests==2.32.4
+tqdm==4.66.4
+numpy==1.26.4

samples/python/network_api_pytorch_mnist/sample.py

@@ -0,0 +1,179 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2025 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 sys
+
+# This sample uses an MNIST PyTorch model to create a TensorRT Inference Engine
+import model
+import numpy as np
+
+import tensorrt as trt
+
+sys.path.insert(1, os.path.join(sys.path[0], os.path.pardir))
+import common
+
+# You can set the logger severity higher to suppress messages (or lower to display more messages).
+TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
+
+
+class ModelData(object):
+    INPUT_NAME = "data"
+    INPUT_SHAPE = (1, 1, 28, 28)
+    OUTPUT_NAME = "prob"
+    OUTPUT_SIZE = 10
+    DTYPE = trt.float32
+
+
+def populate_network(network, weights):
+    # Configure the network layers based on the weights provided.
+    input_tensor = network.add_input(
+        name=ModelData.INPUT_NAME, dtype=ModelData.DTYPE, shape=ModelData.INPUT_SHAPE
+    )
+
+    def add_matmul_as_fc(net, input, outputs, w, b):
+        assert len(input.shape) >= 3
+        m = 1 if len(input.shape) == 3 else input.shape[0]
+        k = int(np.prod(input.shape) / m)
+        assert np.prod(input.shape) == m * k
+        n = int(w.size / k)
+        assert w.size == n * k
+        assert b.size == n
+
+        input_reshape = net.add_shuffle(input)
+        input_reshape.reshape_dims = trt.Dims2(m, k)
+
+        filter_const = net.add_constant(trt.Dims2(n, k), w)
+        mm = net.add_matrix_multiply(
+            input_reshape.get_output(0),
+            trt.MatrixOperation.NONE,
+            filter_const.get_output(0),
+            trt.MatrixOperation.TRANSPOSE,
+        )
+
+        bias_const = net.add_constant(trt.Dims2(1, n), b)
+        bias_add = net.add_elementwise(
+            mm.get_output(0), bias_const.get_output(0), trt.ElementWiseOperation.SUM
+        )
+
+        output_reshape = net.add_shuffle(bias_add.get_output(0))
+        output_reshape.reshape_dims = trt.Dims4(m, n, 1, 1)
+        return output_reshape
+
+    conv1_w = weights["conv1.weight"].cpu().numpy()
+    conv1_b = weights["conv1.bias"].cpu().numpy()
+    conv1 = network.add_convolution_nd(
+        input=input_tensor,
+        num_output_maps=20,
+        kernel_shape=(5, 5),
+        kernel=conv1_w,
+        bias=conv1_b,
+    )
+    conv1.stride_nd = (1, 1)
+
+    pool1 = network.add_pooling_nd(
+        input=conv1.get_output(0), type=trt.PoolingType.MAX, window_size=(2, 2)
+    )
+    pool1.stride_nd = trt.Dims2(2, 2)
+
+    conv2_w = weights["conv2.weight"].cpu().numpy()
+    conv2_b = weights["conv2.bias"].cpu().numpy()
+    conv2 = network.add_convolution_nd(
+        pool1.get_output(0), 50, (5, 5), conv2_w, conv2_b
+    )
+    conv2.stride_nd = (1, 1)
+
+    pool2 = network.add_pooling_nd(conv2.get_output(0), trt.PoolingType.MAX, (2, 2))
+    pool2.stride_nd = trt.Dims2(2, 2)
+
+    fc1_w = weights["fc1.weight"].cpu().numpy()
+    fc1_b = weights["fc1.bias"].cpu().numpy()
+    fc1 = add_matmul_as_fc(network, pool2.get_output(0), 500, fc1_w, fc1_b)
+
+    relu1 = network.add_activation(
+        input=fc1.get_output(0), type=trt.ActivationType.RELU
+    )
+
+    fc2_w = weights["fc2.weight"].cpu().numpy()
+    fc2_b = weights["fc2.bias"].cpu().numpy()
+    fc2 = add_matmul_as_fc(
+        network, relu1.get_output(0), ModelData.OUTPUT_SIZE, fc2_w, fc2_b
+    )
+
+    fc2.get_output(0).name = ModelData.OUTPUT_NAME
+    network.mark_output(tensor=fc2.get_output(0))
+
+
+def build_engine(weights):
+    # For more information on TRT basics, refer to the introductory samples.
+    builder = trt.Builder(TRT_LOGGER)
+    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED))
+    config = builder.create_builder_config()
+    runtime = trt.Runtime(TRT_LOGGER)
+
+    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, common.GiB(1))
+    # Populate the network using weights from the PyTorch model.
+    populate_network(network, weights)
+    # Build and return an engine.
+    plan = builder.build_serialized_network(network, config)
+    return runtime.deserialize_cuda_engine(plan)
+
+
+# Loads a random test case from pytorch's DataLoader
+def load_random_test_case(model, pagelocked_buffer):
+    # Select an image at random to be the test case.
+    img, expected_output = model.get_random_testcase()
+    # Copy to the pagelocked input buffer
+    np.copyto(pagelocked_buffer, img)
+    return expected_output
+
+
+def main():
+    common.add_help(description="Runs an MNIST network using a PyTorch model")
+    # Train the PyTorch model
+    mnist_model = model.MnistModel()
+    mnist_model.learn()
+    weights = mnist_model.get_weights()
+    # Do inference with TensorRT.
+    engine = build_engine(weights)
+
+    # Build an engine, allocate buffers and create a stream.
+    # For more information on buffer allocation, refer to the introductory samples.
+    inputs, outputs, bindings = common.allocate_buffers(engine)
+    context = engine.create_execution_context()
+
+    # Use context manager for proper stream lifecycle management
+    with common.CudaStreamContext() as stream:
+        case_num = load_random_test_case(mnist_model, pagelocked_buffer=inputs[0].host)
+        # For more information on performing inference, refer to the introductory samples.
+        # The common.do_inference function will return a list of outputs - we only have one in this case.
+        [output] = common.do_inference(
+            context,
+            engine=engine,
+            bindings=bindings,
+            inputs=inputs,
+            outputs=outputs,
+            stream=stream,
+        )
+        pred = np.argmax(output)
+        common.free_buffers(inputs, outputs)
+    print("Test Case: " + str(case_num))
+    print("Prediction: " + str(pred))
+
+
+if __name__ == "__main__":
+    main()

samples/python/non_zero_plugin/README.md

@@ -0,0 +1,110 @@
+# Python-based NonZero Plugin for TensorRT using IPluginV3
+
+## Description
+
+This sample, `non_zero_plugin`, implements a Python-based plugin for the NonZero operation, configurable to use a `CUDA Python` or `PyTorch` backend.
+
+NonZero is an operation where the non-zero indices of the input tensor is found.
+
+## How does this sample work?
+
+This sample creates and runs a TensorRT engine built from a network containing a single NonZeroPlugin node. It demonstrates how
+custom layers with data-dependent output shapes can be implemented and added to a TensorRT network using Python.
+
+### Implementing a NonZero plugin using IPluginV3 interface
+
+Until `IPluginV3` (and associated interfaces), TensorRT plugins could not have outputs whose shapes depended on the input values (they could only depend
+on input shapes). `IPluginV3OneBuild` which exposes a build capability for `IPluginV3`, provides support for such data-dependent output shapes.
+
+`NonZeroPlugin` in this sample is written to handle 2-D input tensors of shape $R \times C$. Assume that the tensor contains $K$ non-zero elements and that the
+non-zero indices are required in a row ordering (each set of indices in its own row). Then the output shape would be $K \times 2$.
+
+The output shapes are expressed to the TensorRT builder through the `IPluginV3OneBuild.get_output_shapes()` API. Expressing the second dimension of the output is
+straightforward:
+```
+# output_dims[0] = trt.DimsExprs(2)
+output_dims[0][1] = exprBuilder.constant(2)
+```
+
+The extent of each data-dependent dimension in the plugin must be expressed in terms of a *_size tensor_*. A size tensor is a scalar output of type
+`trt.int32` or `trt.int64` that must be added as one of the plugin outputs. In this case, it is sufficient to declare one size tensor to denote the extent of the
+first dimension of the non-zero indices output. To declare a size tensor, one must provide an upper-bound and optimum value for its extent as `IDimensionExpr`s. These can be formed through the `IExprBuilder` argument passed to the `IPluginV3OneBuild.get_output_shapes()` method.
+ - For unknown inputs, the upper-bound is the total number of elements in the input
+	```
+	upper_bound = exprBuilder.operation(trt.DimensionOperation.PROD, inputs[0][0], inputs[0][1])
+	```
+ - A good estimate for the optimum is that half of the elements are non-zero
+	```
+	opt_value = exprBuilder.operation(trt.DimensionOperation.FLOOR_DIV, upper_bound, exprBuilder.constant(2))
+	```
+
+Now we can declare the size tensor using the `IExprBuilder.declare_size_tensor()` method, which also requires the specification of the output index at which the size tensor would reside. Let us place it after the non-zero indices output:
+```
+num_non_zero_size_tensor = exprBuilder.declare_size_tensor(1, opt_value, upper_bound)
+```
+
+Now we are ready to specify the extent of the first dimension of the non-zero indices output:
+```
+# output_dims[0] = trt.DimsExprs(0)
+output_dims[0][0] = num_non_zero_size_tensor
+```
+Note that the size tensor is declared to be a scalar (0-D):
+
+### Creating network and building the engine
+
+To add the plugin to the network, the `INetworkDefinition::add_plugin_v3()` method must be used.
+
+Similar to `IPluginCreator` used for V2 plugins, V3 plugins must be accompanied by the registration of a plugin creator implementing the `IPluginCreatorV3One` interface.
+
+## Running the sample
+
+1.  Run the sample to create a TensorRT inference engine and run inference:
+    `python3 non_zero_plugin.py [-h] [--precision {fp32,fp16}] [--backend {cuda_python,torch}] [--net_type {onnx,inetdef}]`
+
+2.  Verify that the sample ran successfully. If the sample runs successfully you should see the following message:
+     ```
+    Inference result correct!
+    ```
+
+### Sample `--help` options
+
+To see the full list of available options and their descriptions, use the `-h` or `--help` command line option.
+
+
+# Additional resources
+
+The following resources provide a deeper understanding about the V3 TensorRT plugins and the NonZero operation:
+
+**NonZero**
+- [ONNX: NonZero](https://onnx.ai/onnx/operators/onnx__NonZero.html)
+
+**C++-based NonZero Plugin sample**
+- [NonZero C++ Plugin](../../sampleNonZeroPlugin/)
+
+**TensorRT plugins**
+- [Extending TensorRT with Custom Layers](https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#extending)
+- [TensorRT Python-based Plugins](https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/#add_custom_layer_python)
+
+**Other documentation**
+- [Introduction To NVIDIA’s TensorRT Samples](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sample-support-guide/index.html#samples)
+- [Working With TensorRT Using The Python API](https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/#python_topics)
+- [NVIDIA’s TensorRT Documentation Library](https://docs.nvidia.com/deeplearning/sdk/tensorrt-archived/index.html)
+
+# License
+
+For terms and conditions for use, reproduction, and distribution, see the [TensorRT Software License Agreement](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sla/index.html) documentation.
+
+# Changelog
+
+October 2025
+Migrate to strongly typed APIs.
+
+August 2025
+Removed support for Python versions < 3.10.
+
+April 2024
+This is the first version of this `README.md` file.
+
+# Known issues
+
+There are no known issues in this sample.

samples/python/non_zero_plugin/non_zero_plugin.py

@@ -0,0 +1,352 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025 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 onnx_graphsurgeon as gs
+import numpy as np
+import onnx
+import os
+import sys
+
+import tensorrt as trt
+from polygraphy.backend.trt import (
+    CreateConfig,
+    EngineFromNetwork,
+    NetworkFromOnnxPath,
+    TrtRunner,
+    create_network,
+    engine_from_network,
+)
+
+import argparse
+
+from polygraphy import mod
+
+sys.path.insert(1, os.path.join(os.path.dirname(os.path.realpath(__file__)), os.pardir))
+from plugin_utils import cuda_call, KernelHelper, UnownedMemory, volume
+
+
+cuda = mod.lazy_import("cuda.bindings.driver")
+cudart = mod.lazy_import("cuda.bindings.runtime")
+nvrtc = mod.lazy_import("cuda.bindings.nvrtc")
+
+torch = mod.lazy_import("torch")
+cp = mod.lazy_import("cupy")
+
+non_zero_half_kernel = r'''
+#include <cuda_fp16.h>
+extern "C" __global__
+void find_non_zero_indices_half(
+    half const* X, int* indices, unsigned long long* count, int R, int C)
+{
+    static_assert(sizeof(unsigned long long) == 8U, "unsigned long long must be 8 bytes in NVCC");
+    int row = blockIdx.x * blockDim.x + threadIdx.x;
+
+    // Check if the row index is within bounds
+    if (row < R)
+    {
+        for (int col = 0; col < C; ++col)
+        {
+            half const z = static_cast<half>(0.F);
+            if (X[col + C * row] != z)
+            {
+                // Increment count atomically and get the previous value
+                unsigned long long index = atomicAdd(count, 1ULL);
+                indices[2 * index] = row;
+                indices[2 * index + 1] = col;
+            }
+        }
+    }
+}
+'''
+
+non_zero_float_kernel = r'''
+extern "C" __global__
+void find_non_zero_indices_float(
+    float const* X, int* indices, unsigned long long* count, int R, int C)
+{
+    static_assert(sizeof(unsigned long long) == 8U, "unsigned long long must be 8 bytes in NVCC");
+    int row = blockIdx.x * blockDim.x + threadIdx.x;
+
+    // Check if the row index is within bounds
+    if (row < R)
+    {
+        for (int col = 0; col < C; ++col)
+        {
+            if (X[col + C * row] != 0.F)
+            {
+                // Increment count atomically and get the previous value
+                unsigned long long index = atomicAdd(count, 1ULL);
+                indices[2 * index] = row;
+                indices[2 * index + 1] = col;
+            }
+        }
+    }
+}
+'''
+
+class NonZeroPlugin(trt.IPluginV3, trt.IPluginV3OneCore, trt.IPluginV3OneBuild, trt.IPluginV3OneRuntime):
+    def __init__(self, backend = None):
+        trt.IPluginV3.__init__(self)
+        trt.IPluginV3OneCore.__init__(self)
+        trt.IPluginV3OneBuild.__init__(self)
+        trt.IPluginV3OneRuntime.__init__(self)
+
+        self.num_outputs = 2
+        self.plugin_namespace = ""
+        self.plugin_name = "NonZeroPlugin"
+        self.plugin_version = "1"
+
+        if backend is not None:
+            self.backend = backend.tobytes().decode("utf-8")
+        else:
+            self.backend = "cuda_python"
+
+        self.cuDevice = None
+
+    def get_capability_interface(self, type):
+        return self
+
+    def get_output_data_types(self, input_types):
+        return [trt.DataType.INT32, trt.DataType.INT64]
+
+    def get_output_shapes(self, inputs, shape_inputs, exprBuilder):
+        # First output is 2-D
+        # Second output is a size tensor, which must be declared a scalar (0-D)
+        output_dims = [trt.DimsExprs(2), trt.DimsExprs(0)]
+
+        upper_bound = exprBuilder.operation(trt.DimensionOperation.PROD, inputs[0][0], inputs[0][1])
+        opt_value = exprBuilder.operation(trt.DimensionOperation.FLOOR_DIV, upper_bound, exprBuilder.constant(2))
+        num_non_zero_size_tensor = exprBuilder.declare_size_tensor(1, opt_value, upper_bound)
+
+        output_dims[0][0] = num_non_zero_size_tensor
+        output_dims[0][1] = exprBuilder.constant(2)
+
+        return output_dims
+
+    def get_fields_to_serialize(self):
+        return trt.PluginFieldCollection(
+            [
+                trt.PluginField(
+                    "backend", self.backend.encode(), trt.PluginFieldType.CHAR
+                )
+            ]
+        )
+
+    def configure_plugin(self, inp, out):
+        if self.backend == "cuda_python":
+            self.cuDevice = cuda_call(cuda.cuDeviceGet(0))
+
+    def on_shape_change(self, inp, out):
+        if self.backend == "cuda_python":
+            self.cuDevice = cuda_call(cuda.cuDeviceGet(0))
+
+    def supports_format_combination(self, pos, in_out, num_inputs):
+        assert num_inputs == 1
+        assert pos < len(in_out)
+
+        type_ok = False
+
+        # first input should be float16 or float32
+        if pos == 0:
+            type_ok = in_out[0].desc.type == trt.DataType.FLOAT or in_out[0].desc.type == trt.DataType.HALF
+        elif pos == 1:
+            type_ok = in_out[1].desc.type == trt.DataType.INT32
+        else: # pos == 2
+            # size tensor outputs must be NCHW INT64
+            type_ok = in_out[2].desc.type == trt.DataType.INT64
+
+        return in_out[pos].desc.format == trt.TensorFormat.LINEAR and type_ok
+
+    def enqueue(self, input_desc, output_desc, inputs, outputs, workspace, stream):
+        inp_dtype = trt.nptype(input_desc[0].type)
+
+        if self.backend == "cuda_python":
+            R = input_desc[0].dims[0]
+            C = input_desc[0].dims[1]
+
+            blockSize = 256
+            numBlocks = int((C + blockSize - 1) // blockSize)
+
+            d_in = np.array([inputs[0]], dtype=np.uint64)
+            d_out_0 = np.array([outputs[0]], dtype=np.uint64)
+            d_out_1 = np.array([outputs[1]], dtype=np.uint64)
+
+            args = [d_in, d_out_0, d_out_1, np.array(R, dtype=np.uint32), np.array(C, dtype=np.uint32)]
+            kernelArgs = np.array([arg.ctypes.data for arg in args], dtype=np.uint64)
+
+            stream_ptr = np.array([stream], dtype=np.uint64)
+
+            if inp_dtype == np.float32:
+                kernelHelper = KernelHelper(non_zero_float_kernel, int(self.cuDevice))
+                _non_zero_float_kernel = kernelHelper.getFunction(b'find_non_zero_indices_float')
+                cuda_call(cuda.cuLaunchKernel(_non_zero_float_kernel,
+                                            numBlocks, 1, 1,
+                                            blockSize, 1, 1,
+                                            0,
+                                            stream_ptr,
+                                            kernelArgs, 0))
+            elif inp_dtype == np.float16:
+                kernelHelper = KernelHelper(non_zero_half_kernel, int(self.cuDevice))
+                _non_zero_half_kernel = kernelHelper.getFunction(b'find_non_zero_indices_half')
+                cuda_call(cuda.cuLaunchKernel(_non_zero_half_kernel,
+                                            numBlocks, 1, 1,
+                                            blockSize, 1, 1,
+                                            0,
+                                            stream_ptr,
+                                            kernelArgs, 0))
+            else:
+                raise ValueError("inp_dtype not valid")
+
+        elif self.backend == "torch":
+            inp_mem = UnownedMemory(inputs[0], input_desc[0].dims, inp_dtype)
+
+            out_mem = UnownedMemory(
+                outputs[0], 2 * volume(input_desc[0].dims), np.int32
+            )
+
+            out_1_mem = UnownedMemory(outputs[1], 1, np.int64)
+
+            a_t = torch.as_tensor(inp_mem.d, device="cuda")
+            out = torch.nonzero(a_t)
+
+            out_mem.d[: volume(out.shape)] = cp.reshape(cp.asarray(out), (-1,))
+            cp.copyto(out_1_mem.d, cp.reshape(cp.asarray([out.shape[0]]), (-1,)))
+
+        else:
+            raise ValueError(f"backend not valid: {self.backend}")
+
+    def attach_to_context(self, context):
+        return self.clone()
+
+    def set_tactic(self, tactic):
+        pass
+
+    def clone(self):
+        cloned_plugin = NonZeroPlugin()
+        cloned_plugin.__dict__.update(self.__dict__)
+        return cloned_plugin
+
+    #
+    # The following defaults take effect since the respective methods are not overriden
+    #
+
+    # def get_valid_tactics(self):
+    #     return []
+
+    # def get_workspace_size(self, input_desc, output_desc):
+    #     return 0
+
+    # def destroy(self):
+    #     pass
+
+
+class NonZeroPluginCreator(trt.IPluginCreatorV3One):
+    def __init__(self):
+        trt.IPluginCreatorV3One.__init__(self)
+        self.name = "NonZeroPlugin"
+        self.plugin_namespace = ""
+        self.plugin_version = "1"
+        self.field_names = trt.PluginFieldCollection(
+            [trt.PluginField("backend", np.array([]), trt.PluginFieldType.CHAR)]
+        )
+
+    def create_plugin(self, name, fc, phase):
+        backend = None
+        for f in fc:
+            if f.name == "backend":
+                backend = f.data[:-1] if f.data[-1] == 0 else f.data
+        return NonZeroPlugin(backend)
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--precision', type=str, default="fp32", choices=["fp32", "fp16"])
+    parser.add_argument("--backend", type=str, default="torch", choices=["cuda_python", "torch"])
+    parser.add_argument('--net_type', type=str, default="onnx", choices=["onnx", "inetdef"])
+
+    args = parser.parse_args()
+
+    if args.backend == "cuda_python":
+        # Initialize CUDA and create default context
+        cuda_call(cudart.cudaFree(0))
+
+    elif args.backend == "torch":
+        # Initialize CUDA and create default context
+        torch.cuda.init()
+
+    precision = np.float32 if args.precision == "fp32" else np.float16
+
+    inp_shape = (128, 128)
+    X = np.random.normal(size=inp_shape).astype(precision)
+    # Zero out a random set of indices
+    indices = np.random.choice(np.prod(inp_shape), replace=False, size=np.random.randint(0, np.prod(inp_shape) + 1))
+    X[np.unravel_index(indices, inp_shape)] = 0
+
+    # Register plugin creator
+    plg_registry = trt.get_plugin_registry()
+    my_plugin_creator = NonZeroPluginCreator()
+    plg_registry.register_creator(my_plugin_creator, "")
+
+    if args.net_type == "onnx":
+        # create ONNX model
+        onnx_path = "test_NonZeroPlugin.onnx"
+        inputX = gs.Variable(name="X", shape=inp_shape, dtype=precision)
+        Y = gs.Variable(name="Y", dtype=np.int32)
+        Y_num = gs.Variable(name="Y_num", dtype=np.int64)
+        nonZeroPluginNode = gs.Node(
+            name="NonZeroPlugin",
+            op="NonZeroPlugin",
+            inputs=[inputX],
+            outputs=[Y, Y_num],
+            attrs={"backend": args.backend.encode()},
+        )
+        graph = gs.Graph(nodes=[nonZeroPluginNode], inputs=[inputX], outputs=[Y], opset=16)
+        onnx.save(gs.export_onnx(graph), onnx_path)
+
+        # build engine
+        build_engine = EngineFromNetwork(
+            NetworkFromOnnxPath(onnx_path, strongly_typed=True), CreateConfig()
+        )
+    else:
+        # Create plugin object
+        builder, network = create_network(strongly_typed=True)
+        plg_creator = plg_registry.get_creator("NonZeroPlugin", "1", "")
+        plugin_fields_list = [
+            trt.PluginField("backend", args.backend.encode(), trt.PluginFieldType.CHAR)
+        ]
+        pfc = trt.PluginFieldCollection(plugin_fields_list)
+        plugin = plg_creator.create_plugin("NonZeroPlugin", pfc, trt.TensorRTPhase.BUILD)
+
+        # Populate network
+        inputX = network.add_input(name="X", dtype=trt.float32 if precision==np.float32 else trt.float16, shape=inp_shape)
+        out = network.add_plugin_v3([inputX], [], plugin)
+        out.get_output(0).name = "Y"
+        network.mark_output(tensor=out.get_output(0))
+        build_engine = engine_from_network((builder, network), CreateConfig())
+
+    # Compare against Numpy's nonzero
+    Y_ref = np.transpose(np.nonzero(X))
+
+    # Run
+    with TrtRunner(build_engine, "trt_runner") as runner:
+        outputs = runner.infer({"X": X})
+        Y = outputs["Y"]
+        Y = Y[np.lexsort(np.fliplr(Y).T)]
+
+        if np.allclose(Y, Y_ref):
+            print("Inference result correct!")
+        else:
+            print("Inference result incorrect!")

samples/python/non_zero_plugin/requirements.txt

@@ -0,0 +1,13 @@
+cuda-python==12.9.0
+cupy-cuda12x
+torch
+--extra-index-url https://pypi.ngc.nvidia.com
+polygraphy>=0.50.1
+colored
+numpy==1.26.4
+--extra-index-url https://pypi.ngc.nvidia.com
+onnx-graphsurgeon
+pywin32; platform_system == "Windows"
+pyyaml==6.0.3
+requests==2.32.4
+tqdm==4.66.4

samples/python/onnx_custom_plugin/CMakeLists.txt

@@ -0,0 +1,81 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2026 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.
+#
+
+cmake_minimum_required(VERSION 3.31 FATAL_ERROR)
+project(CustomHardMax LANGUAGES CXX CUDA)
+
+find_package(CUDAToolkit REQUIRED)
+
+if(NOT MSVC)
+    # Enable all compile warnings
+    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall -Wno-long-long -pedantic -Wno-deprecated-declarations")
+    set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -Xcompiler -Wno-deprecated-declarations")
+endif()
+
+# Sets variable to a value if variable is unset.
+macro(set_ifndef var val)
+    if(NOT ${var})
+        set(${var} ${val})
+    endif()
+    message(STATUS "Configurable variable ${var} set to ${${var}}")
+endmacro()
+
+# -------- CONFIGURATION --------
+if(NOT MSVC)
+    set_ifndef(TRT_LIB /usr/lib/x86_64-linux-gnu)
+    set_ifndef(TRT_INCLUDE /usr/include/x86_64-linux-gnu)
+
+endif()
+
+# Find dependencies:
+message("\nThe following variables are derived from the values of the previous variables unless provided explicitly:\n")
+
+# TensorRT’s nvinfer lib
+find_library(
+    _NVINFER_LIB nvinfer
+    HINTS ${TRT_LIB}
+    PATH_SUFFIXES lib lib64)
+set_ifndef(NVINFER_LIB ${_NVINFER_LIB})
+
+# -------- BUILDING --------
+
+add_definitions(-DTENSORRT_BUILD_LIB)
+
+# Add include directories
+file(REAL_PATH ${CMAKE_SOURCE_DIR} CMAKE_SOURCE_DIR_REALPATH)
+
+get_filename_component(SAMPLES_COMMON_DIR ${CMAKE_SOURCE_DIR_REALPATH}/../../common/ ABSOLUTE)
+get_filename_component(SHARED_DIR ${CMAKE_SOURCE_DIR_REALPATH}/../../../shared ABSOLUTE)
+get_filename_component(SAMPLES_DIR ${CMAKE_SOURCE_DIR_REALPATH}/../../ ABSOLUTE)
+# Define Hardmax plugin library target
+add_library(
+    customHardmaxPlugin MODULE
+    ${SAMPLES_COMMON_DIR}/logger.cpp ${SHARED_DIR}/utils/fileLock.cpp
+    ${CMAKE_SOURCE_DIR_REALPATH}/plugin/customHardmaxPlugin.cpp ${CMAKE_SOURCE_DIR_REALPATH}/plugin/customHardmaxPlugin.h)
+
+target_include_directories(customHardmaxPlugin PRIVATE
+    ${CUDAToolkit_INCLUDE_DIRS} ${TRT_INCLUDE} ${CMAKE_SOURCE_DIR_REALPATH}/plugin/
+    ${SAMPLES_COMMON_DIR} ${SAMPLES_DIR} ${SHARED_DIR})
+
+# Use C++11
+target_compile_features(customHardmaxPlugin PUBLIC cxx_std_17)
+
+# Link TensorRT’s nvinfer lib
+target_link_libraries(customHardmaxPlugin PRIVATE ${NVINFER_LIB})
+target_link_libraries(customHardmaxPlugin PRIVATE CUDA::cudart_static)
+target_link_libraries(customHardmaxPlugin PRIVATE CUDA::cublas)
+target_link_libraries(customHardmaxPlugin PRIVATE CUDA::cuda_driver)

samples/python/onnx_custom_plugin/README.md

@@ -0,0 +1,177 @@
+# Adding A Custom Layer Implementation to Your ONNX Network
+
+**Table Of Contents**
+- [Description](#description)
+- [How does this sample work?](#how-does-this-sample-work)
+- [Prerequisites](#prerequisites)
+- [Download and preprocess the ONNX model](#download-the-onnx-model)
+- [Running the sample](#running-the-sample)
+- [Additional resources](#additional-resources)
+- [License](#license)
+- [Changelog](#changelog)
+- [Known issues](#known-issues)
+
+## Description
+
+This sample, `onnx_custom_plugin`, demonstrates how to use plugins written in C++ with the TensorRT Python bindings and ONNX Parser. This sample uses the [BiDAF Model](https://github.com/onnx/models/tree/main/text/machine_comprehension/bidirectional_attention_flow) from ONNX Model Zoo.
+
+## How does this sample work?
+
+This sample implements a Hardmax layer using cuBLAS, wraps the implementation in a TensorRT plugin (with a corresponding plugin creator) and then generates a shared library module containing its code. The user then dynamically loads this library in Python, which causes the plugin to be registered in TensorRT's PluginRegistry and makes it available to the ONNX parser.
+
+This sample includes:
+
+`plugin/`
+This directory contains files for the Hardmax layer plugin.
+
+`customHardmaxPlugin.cpp`
+A custom TensorRT plugin implementation.
+
+`customHardmaxPlugin.h`
+The Hardmax Plugin headers.
+
+`model.py`
+This script downloads the BiDAF onnx model and uses Onnx Graphsurgeon to replace layers unsupported by TensorRT.
+
+`sample.py`
+This script loads the ONNX model and performs inference using TensorRT.
+
+`load_plugin_lib.py`
+This script contains a helper function to load the customHardmaxPlugin library in Python.
+
+`test_custom_hardmax_plugin.py`
+This script tests the Hardmax Plugin against a reference numpy implementation.
+
+`requirements.txt`
+This file specifies all the Python packages required to run this Python sample.
+
+## Prerequisites
+
+For specific software versions, see the [TensorRT Installation Guide](https://docs.nvidia.com/deeplearning/sdk/tensorrt-archived/index.html).
+
+1. Install the dependencies for Python.
+
+```bash
+pip3 install -r requirements.txt
+```
+
+2. [Install CMake](https://cmake.org/download/).
+
+3. [Install Cublas](https://developer.nvidia.com/cublas).
+
+4. (For Windows builds) [Visual Studio](https://visualstudio.microsoft.com/vs/older-downloads/) 2017 Community or Enterprise edition
+
+## Download and preprocess the ONNX model
+
+Run the model script to download the BiDAF model from the Onnx Model Zoo. The script will replace the `Hardmax` layer with an op called `CustomHardmax` to match the custom Plugin name. It will also replace the unsupported `Compress` node with an equivalent operation, and remove the `CategoryMapper` nodes which do a String-to-Int conversion of the model inputs.
+
+```bash
+python3 model.py
+```
+
+## Running the sample
+
+1.  Build the plugin and its corresponding Python bindings.
+
+   - On Linux, run:
+      ```bash
+      mkdir build && pushd build
+      cmake .. && make -j
+      popd
+      ```
+
+      **NOTE:** If any of the dependencies are not installed in their default locations, you can manually specify them. For example:
+      ```bash
+      cmake .. -DCMAKE_CUDA_COMPILER=/usr/local/cuda-x.x/bin/nvcc # (Or adding /path/to/nvcc into $PATH)
+               -DCUDA_INC_DIR=/usr/local/cuda-x.x/include/  # (Or adding /path/to/cuda/include into $CPLUS_INCLUDE_PATH)
+               -DTRT_LIB=/path/to/tensorrt/lib/
+               -DTRT_INCLUDE=/path/to/tensorrt/include/
+      ```
+
+   - On Windows, run the following in Powershell, replacing paths appropriately:
+      ```ps1
+      mkdir build; pushd build
+      cmake .. -G "Visual Studio 15 Win64" /
+         -DTRT_LIB=C:\path\to\tensorrt\lib /
+         -DTRT_INCLUDE=C:\path\to\tensorrt\lib /
+         -DCUDA_INC_DIR="C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v<CUDA_VERSION>\include" /
+         -DCUDA_LIB_DIR="C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v<CUDA_VERSION>\lib\x64"
+      # NOTE: msbuild is usually located under C:\Program Files (x86)\Microsoft Visual Studio\2017\<EDITION>\MSBuild\<VERSION>\Bin
+      #   You should add this path to your PATH environment variable.
+      msbuild ALL_BUILD.vcxproj
+      popd
+      ```
+
+   The command `cmake ..` displays a complete list of configurable variables. If a variable is set to `VARIABLE_NAME-NOTFOUND`, then you’ll need to specify it manually or set the variable it is derived from correctly.
+
+2.  Run inference using TensorRT with the custom Hardmax plugin implementation:
+   ```bash
+   python3 sample.py
+   ```
+
+3.  Verify that the sample ran successfully.
+   ```
+   === Testing ===
+
+   Input context: Garry the lion is 5 years old. He lives in the savanna.
+   Input query: Where does the lion live?
+   Model prediction:  savanna
+
+   Input context: A quick brown fox jumps over the lazy dog.
+   Input query: What color is the fox?
+   Model prediction:  brown
+   ```
+
+   The model can also be run interactively:
+   ```bash
+   python3 sample.py --interactive
+   ```
+
+   The context and query can then be entered from the command line:
+
+   ```
+   === Testing ===
+   Enter context: Waldo wears a striped shirt. He also wears glasses.
+   Enter query: Who wears glasses?
+   Model prediction:  waldo
+   ```
+
+# Additional resources
+
+The following resources provide a deeper understanding about getting started with TensorRT using Python:
+
+**Model**
+- [BiDAF model](https://allenai.github.io/bi-att-flow/)
+
+**Documentation**
+- [Introduction To NVIDIA’s TensorRT Samples](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sample-support-guide/index.html#samples)
+- [Working With TensorRT Using The Python API](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#python_topics)
+- [NVIDIA’s TensorRT Documentation Library](https://docs.nvidia.com/deeplearning/sdk/tensorrt-archived/index.html)
+
+# License
+
+For terms and conditions for use, reproduction, and distribution, see the [TensorRT Software License Agreement](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sla/index.html) documentation.
+
+# Changelog
+March 2026
+  - Migrate HardmaxPlugin from IPluginV2DynamicExt to IPluginV3.
+
+October 2025
+  - Migrate to strongly typed APIs.
+
+August 2025:
+  - Removed support for Python versions < 3.10.
+
+January 2024:
+  - Create cublas handle with cublasCreate instead of using the cublasContext argument from attachToContext.
+  - Added the Cublas library as a prerequisite.
+
+August 2023:
+  - Update ONNX version support to 1.14.0
+  - Removed support for Python versions < 3.8.
+
+September 2022: This `README.md` file was created and reviewed.
+
+# Known issues
+
+There are no known issues in this sample.

samples/python/onnx_custom_plugin/load_plugin_lib.py

@@ -0,0 +1,58 @@
+#
+# 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 ctypes
+
+WORKING_DIR = os.environ.get("TRT_WORKING_DIR") or os.path.dirname(
+    os.path.realpath(__file__)
+)
+IS_WINDOWS = os.name == "nt"
+if IS_WINDOWS:
+    HARDMAX_PLUGIN_LIBRARY_NAME = "customHardmaxPlugin.dll"
+    HARDMAX_PLUGIN_LIBRARY = [
+        os.path.join(WORKING_DIR, "build", "Debug", HARDMAX_PLUGIN_LIBRARY_NAME),
+        os.path.join(WORKING_DIR, "build", "Release", HARDMAX_PLUGIN_LIBRARY_NAME),
+    ]
+else:
+    HARDMAX_PLUGIN_LIBRARY_NAME = "libcustomHardmaxPlugin.so"
+    HARDMAX_PLUGIN_LIBRARY = [
+        os.path.join(WORKING_DIR, "build", HARDMAX_PLUGIN_LIBRARY_NAME)
+    ]
+
+
+def load_plugin_lib():
+    for plugin_lib in HARDMAX_PLUGIN_LIBRARY:
+        if os.path.isfile(plugin_lib):
+            try:
+                # Python specifies that winmode is 0 by default, but some implementations
+                # incorrectly default to None instead. See:
+                # https://docs.python.org/3.8/library/ctypes.html
+                # https://github.com/python/cpython/blob/3.10/Lib/ctypes/__init__.py#L343
+                ctypes.CDLL(plugin_lib, winmode=0)
+            except TypeError:
+                # winmode only introduced in python 3.8
+                ctypes.CDLL(plugin_lib)
+            return
+
+    raise IOError(
+        "\n{}\n{}\n{}\n".format(
+            "Failed to load library ({}).".format(HARDMAX_PLUGIN_LIBRARY_NAME),
+            "Please build the Hardmax sample plugin.",
+            "For more information, see the included README.md",
+        )
+    )

samples/python/onnx_custom_plugin/model.py

@@ -0,0 +1,129 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2026 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 json
+
+import wget
+import onnx
+import onnx_graphsurgeon as gs
+
+MODEL_URL = "https://github.com/onnx/models/raw/e77240a62df68ed13e3138a5812553a552b857bb/text/machine_comprehension/bidirectional_attention_flow/model/bidaf-9.onnx"
+
+WORKING_DIR = os.environ.get("TRT_WORKING_DIR") or os.path.dirname(
+    os.path.realpath(__file__)
+)
+MODEL_DIR = os.path.join(WORKING_DIR, "models")
+RAW_MODEL_PATH = os.path.join(MODEL_DIR, "bidaf-9.onnx")
+TRT_MODEL_PATH = os.path.join(MODEL_DIR, "bidaf-9-trt.onnx")
+
+
+def _do_graph_surgery(raw_model_path, trt_model_path):
+    graph = gs.import_onnx(onnx.load(raw_model_path))
+
+    # Replace unsupported Hardmax with our CustomHardmax op
+    hardmax_node = None
+    for node in graph.nodes:
+        if node.op == "Hardmax":
+            node.op = "CustomHardmax"
+            hardmax_node = node
+    assert hardmax_node is not None, "Model does not contain a Hardmax node"
+
+    # The original onnx model also uses another unsupported op called "Compress".
+    # "Compress" returns values from the first tensor for all indices which evaluate to
+    # True in the second tensor. In our case the second Tensor is the output of Hardmax,
+    # so exactly one index will evaluate to true because the value at it will be 1, and
+    # all other values will be 0. We can achieve the same result as "Compress" by taking the
+    # dot product of our value tensor and the Hardmax output.
+    #
+    # So, we will replace the subgraph Compress(Transpose_29, Cast(Reshape(Hardmax)))
+    # with the subgraph Einsum(Transpose_29, Hardmax) where the equation in Einsum takes the dot product.
+    node_by_name = {node.name: node for node in graph.nodes}
+    transpose_node = node_by_name["Transpose_29"]
+    compress_node = node_by_name["Compress_31"]
+
+    einsum_node = gs.Node(
+        "Einsum",
+        "Dot_of_Hardmax_and_Transpose",
+        attrs={"equation": "ij,ij->i"},  # "Dot product" of 2d tensors
+        inputs=[hardmax_node.outputs[0], transpose_node.outputs[0]],
+        outputs=[compress_node.outputs[0]],
+    )
+    graph.nodes.append(einsum_node)
+
+    # Separate the old subgraph which will be deleted with graph.cleanup()
+    hardmax_node.o().inputs.clear()
+    transpose_node.o().inputs.clear()
+    compress_node.outputs.clear()
+
+    # Also remove the CategoryMapper nodes which convert strings to integers as the first step in the model.
+    # We need to convert the following structure:
+    #
+    #      Input as                        Converted to
+    #   String tokens                     Integer tokens
+    #  ---------------->[CategoryMapper]------------------>[Rest of Model]
+    #
+    # into the following:
+    #
+    #      Input as
+    #   Integer tokens
+    #  ------------------>[Rest of Model]
+    #
+    # Later we will feed the model the integer tokens directly.
+    # Note: list conversion is necessary because we modify graph.nodes in the for loop.
+    category_mapper_nodes = [
+        node for node in graph.nodes if node.op == "CategoryMapper"
+    ]
+    for node in category_mapper_nodes:
+        # Remove CategoryMapper node from onnx graph
+        graph.nodes.remove(node)
+
+        # Also remove references its inputs in the graph's inputs
+        for input_tensor in node.inputs:
+            graph.inputs.remove(input_tensor)
+
+        # The graph's new inputs are the Integer tokens output by CategoryMapper
+        graph.inputs += node.outputs
+
+        # Save String->Int map
+        with open(node.name + ".json", "w") as fp:
+            json.dump(node.attrs, fp)
+
+    graph.cleanup().toposort()
+    onnx.save(gs.export_onnx(graph), trt_model_path)
+
+
+def make_trt_compatible_onnx_model():
+    os.makedirs(MODEL_DIR, exist_ok=True)
+    if not os.path.exists(RAW_MODEL_PATH):
+        wget.download(MODEL_URL, out=RAW_MODEL_PATH)
+        print("\nDownloaded BiDAF model from Onnx Model Zoo")
+    print("Performing graph surgery on Onnx Model Zoo BiDAF model")
+    _do_graph_surgery(RAW_MODEL_PATH, TRT_MODEL_PATH)
+    print("Graph Surgery complete!")
+
+
+def main():
+    if os.path.exists(TRT_MODEL_PATH):
+        print("TRT-compatible onnx model already exists!")
+    else:
+        print("TRT-compatible onnx model not found, generating...")
+        make_trt_compatible_onnx_model()
+
+
+if __name__ == "__main__":
+    main()

samples/python/onnx_custom_plugin/plugin/customHardmaxPlugin.cpp

@@ -0,0 +1,394 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2026 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.
+ */
+
+#include "customHardmaxPlugin.h"
+#include "NvInferPlugin.h"
+#include "common.h" // volume(), ASSERT
+#include "logger.h" // sample::gLogError
+#include <cuda.h>
+#include <memory>
+#include <string_view>
+
+using namespace nvinfer1;
+
+#define CUDRIVER_CALL(call)                                                                                            \
+    {                                                                                                                  \
+        cudaError_enum s_ = call;                                                                                      \
+        if (s_ != CUDA_SUCCESS)                                                                                        \
+        {                                                                                                              \
+            char const *errName_, *errDesc_;                                                                           \
+            cuGetErrorName(s_, &errName_);                                                                             \
+            cuGetErrorString(s_, &errDesc_);                                                                           \
+            sample::gLogError << "CUDA Error: " << errName_ << " " << errDesc_ << std::endl;                           \
+            return s_;                                                                                                 \
+        }                                                                                                              \
+    }
+
+#define CUDA_CALL(call)                                                                                                \
+    {                                                                                                                  \
+        cudaError_t s_ = call;                                                                                         \
+        if (s_ != cudaSuccess)                                                                                         \
+        {                                                                                                              \
+            sample::gLogError << "CUDA Error: " << cudaGetErrorName(s_) << " " << cudaGetErrorString(s_) << std::endl; \
+            return s_;                                                                                                 \
+        }                                                                                                              \
+    }
+
+#define CUBLAS_CALL(call)                                                                                              \
+    {                                                                                                                  \
+        cublasStatus_t s_ = call;                                                                                      \
+        if (s_ != CUBLAS_STATUS_SUCCESS)                                                                               \
+        {                                                                                                              \
+            sample::gLogError << "cuBLAS Error: " << s_ << std::endl;                                                  \
+            return s_;                                                                                                 \
+        }                                                                                                              \
+    }
+
+REGISTER_TENSORRT_PLUGIN(HardmaxPluginCreator);
+
+namespace
+{
+constexpr char const* kHARDMAX_NAME{"CustomHardmax"};
+constexpr char const* kHARDMAX_VERSION{"1"};
+} // namespace
+
+HardmaxPlugin::HardmaxPlugin(int32_t axis)
+    : mAxis(axis)
+{
+}
+
+HardmaxPlugin::HardmaxPlugin(HardmaxPlugin const& other)
+    : mNamespace(other.mNamespace)
+    , mAxisSize(other.mAxisSize)
+    , mDimProductOuter(other.mDimProductOuter)
+    , mDimProductInner(other.mDimProductInner)
+    , mCublas(nullptr)
+    , mAxis(other.mAxis)
+{
+}
+
+HardmaxPlugin::~HardmaxPlugin()
+{
+    if (mCublas)
+    {
+        cublasDestroy(mCublas);
+    }
+}
+
+// IPluginV3 methods
+
+IPluginCapability* HardmaxPlugin::getCapabilityInterface(PluginCapabilityType type) noexcept
+{
+    if (type == PluginCapabilityType::kBUILD)
+    {
+        return static_cast<IPluginV3OneBuild*>(this);
+    }
+    if (type == PluginCapabilityType::kRUNTIME)
+    {
+        return static_cast<IPluginV3OneRuntime*>(this);
+    }
+    ASSERT(type == PluginCapabilityType::kCORE);
+    return static_cast<IPluginV3OneCore*>(this);
+}
+
+IPluginV3* HardmaxPlugin::clone() noexcept
+{
+    auto plugin = std::make_unique<HardmaxPlugin>(*this);
+    return plugin.release();
+}
+
+// IPluginV3OneCore methods
+
+char const* HardmaxPlugin::getPluginName() const noexcept
+{
+    return kHARDMAX_NAME;
+}
+
+char const* HardmaxPlugin::getPluginVersion() const noexcept
+{
+    return kHARDMAX_VERSION;
+}
+
+char const* HardmaxPlugin::getPluginNamespace() const noexcept
+{
+    return mNamespace.c_str();
+}
+
+// IPluginV3OneBuild methods
+
+int32_t HardmaxPlugin::getNbOutputs() const noexcept
+{
+    return 1;
+}
+
+int32_t HardmaxPlugin::getOutputDataTypes(
+    DataType* outputTypes, int32_t nbOutputs, DataType const* inputTypes, int32_t nbInputs) const noexcept
+{
+    ASSERT(inputTypes != nullptr);
+    ASSERT(nbInputs == 1);
+    ASSERT(nbOutputs == 1);
+    outputTypes[0] = inputTypes[0];
+    return 0;
+}
+
+int32_t HardmaxPlugin::getOutputShapes(DimsExprs const* inputs, int32_t nbInputs, DimsExprs const* shapeInputs,
+    int32_t nbShapeInputs, DimsExprs* outputs, int32_t nbOutputs, IExprBuilder& exprBuilder) noexcept
+{
+    ASSERT(nbInputs == 1);
+    ASSERT(nbOutputs == 1);
+    outputs[0] = inputs[0];
+    return 0;
+}
+
+bool HardmaxPlugin::supportsFormatCombination(
+    int32_t pos, DynamicPluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept
+{
+    ASSERT(inOut && pos < (nbInputs + nbOutputs));
+
+    // Type changes are not allowed
+    if (inOut[0].desc.type != inOut[pos].desc.type)
+    {
+        return false;
+    }
+
+    return inOut[pos].desc.type == DataType::kFLOAT && inOut[pos].desc.format == PluginFormat::kLINEAR;
+}
+
+int32_t HardmaxPlugin::configurePlugin(
+    DynamicPluginTensorDesc const* in, int32_t nbInputs, DynamicPluginTensorDesc const* out, int32_t nbOutputs) noexcept
+{
+    ASSERT(nbInputs == 1);
+    ASSERT(nbOutputs == 1);
+
+    Dims const& inDims = in[0].desc.dims;
+
+    // Normalize negative axis to positive
+    if (mAxis < 0)
+    {
+        mAxis += inDims.nbDims;
+        ASSERT(mAxis >= 0);
+    }
+    ASSERT(inDims.nbDims > mAxis);
+
+    return 0;
+}
+
+size_t HardmaxPlugin::getWorkspaceSize(DynamicPluginTensorDesc const* inputs, int32_t nbInputs,
+    DynamicPluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept
+{
+    ASSERT(mAxis >= 0);
+    // Two arrays are needed:
+    // 1. For the contents of the working axis
+    // 2. For an array of 1's
+    return 2 * inputs[0].max.d[mAxis] * sizeof(float);
+}
+
+// IPluginV3OneRuntime methods
+
+int32_t HardmaxPlugin::onShapeChange(
+    PluginTensorDesc const* in, int32_t nbInputs, PluginTensorDesc const* out, int32_t nbOutputs) noexcept
+{
+    ASSERT(nbInputs == 1);
+    ASSERT(nbOutputs == 1);
+
+    Dims const& inDims = in[0].dims;
+
+    // Axis should already be normalized by configurePlugin, but handle it regardless to be safe.
+    if (mAxis < 0)
+    {
+        mAxis += inDims.nbDims;
+        ASSERT(mAxis >= 0);
+    }
+    ASSERT(inDims.nbDims > mAxis);
+
+    mDimProductOuter = samplesCommon::volume(inDims, 0, mAxis);
+    mAxisSize = inDims.d[mAxis];
+    mDimProductInner = samplesCommon::volume(inDims, mAxis + 1, inDims.nbDims);
+
+    return 0;
+}
+
+int32_t HardmaxPlugin::enqueue(PluginTensorDesc const* inputDesc, PluginTensorDesc const* outputDesc,
+    void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept
+{
+    if (inputDesc[0].type != DataType::kFLOAT)
+    {
+        return -1;
+    }
+
+    CUBLAS_CALL(cublasSetStream(mCublas, stream));
+
+    auto const* data = static_cast<float const*>(inputs[0]);
+    auto* result = static_cast<float*>(outputs[0]);
+
+    // Make sure output is initialized to all 0's.
+    // Later we will set the correct outputs to be 1's and not touch the rest.
+    CUDA_CALL(cudaMemsetAsync(result, 0, mDimProductOuter * mDimProductInner * mAxisSize * sizeof(float), stream));
+
+    // We use the workspace in the case that the first call to 'cublasIsamax' is insufficient.
+    // The first half of the workspace we use to copy the values of the axis into, so that we can
+    // subtract out the minimum value and call 'cublasIsamax' again. See the comment below.
+    // The second half of the workspace will be a costant array of 1's, necessary for our cublasSaxpy call.
+    auto* const axisFlat = static_cast<float* const>(workspace);
+    float* const ones = axisFlat + mAxisSize;
+    float const one = 1.0F;
+    CUDRIVER_CALL(cuMemsetD32Async(CUdeviceptr(ones), *reinterpret_cast<int const*>(&one), mAxisSize, stream));
+
+    // This plugin works by parallelizing the argmax operation along a single axis.
+    // This is efficient when the axis size is very large compared to the other dimensions.
+    //
+    // Consider an input shape (1, 512, 3) with axis = 1. This plugin will perform well because
+    // the work which is parallelized is over the large 512-element-long axis, and the work that is done
+    // serially is over the small 1-element-long and 3-element-long axes.
+    //
+    // However, when the axis size is small compared to the other dimensions, this plugin will be very
+    // inefficient. If the input shape is (1, 512, 3) and the hardmax is over axis = 2, then
+    // the work is parallelized over the small 3-element-long axis and the work is done serially over
+    // the large 512-element-long axis. A smarter plugin would try to recognize this and parallelize
+    // the work which would take longest.
+    for (int32_t outer = 0; outer < mDimProductOuter; outer++)
+    {
+        for (int32_t inner = 0; inner < mDimProductInner; inner++)
+        {
+            int32_t const axesOffset = outer * mDimProductInner * mAxisSize + inner;
+            float const* arr = &data[axesOffset];
+            int32_t const stride = mDimProductInner;
+            int32_t argmaxResult;
+            CUBLAS_CALL(cublasIsamax(mCublas, mAxisSize, arr, stride, &argmaxResult));
+
+            // cublasIsamax returns 1-indexed so convert to 0-indexed
+            argmaxResult--;
+
+            // cublasIsamax returns the index of the element with the highest absolute value.
+            // If this element is positive, then we know it is also the max.
+            // However, if it is negative, we need to
+            //      1) Copy the axis into our workspace
+            //      2) Subtract the minimum value we found from our array. This ensures that
+            //         none of the values are negative, and that the largest element remains
+            //         the largest element.
+            //      3) Use cublasIsamax to find the largest element again.
+            // NOTE: We are using cudaMemcpy instead of cudaMemcpyAsync because we need to know
+            //       maxAbsValue before proceeding. However, using synchronous rather than
+            //       asynchronous calls inside of enqueue() hurts performance.
+            //       This could be fixed by implementing the functionality of this plugin with a kernel
+            //       instead of relying only on cuBLAS.
+            float maxAbsValue;
+            CUDA_CALL(cudaMemcpy(&maxAbsValue, &arr[argmaxResult * stride], sizeof(float), cudaMemcpyDeviceToHost));
+            if (maxAbsValue < 0)
+            {
+                float negMinValue = -maxAbsValue;
+                CUBLAS_CALL(cublasScopy(mCublas, mAxisSize, arr, stride, axisFlat, 1));
+                CUBLAS_CALL(cublasSaxpy(mCublas, mAxisSize, &negMinValue, ones, 1, axisFlat, 1));
+                CUBLAS_CALL(cublasIsamax(mCublas, mAxisSize, axisFlat, 1, &argmaxResult));
+                argmaxResult--;
+            }
+
+            CUDA_CALL(cudaMemcpyAsync(
+                &result[axesOffset + argmaxResult * stride], &one, sizeof(float), cudaMemcpyHostToDevice, stream));
+        }
+    }
+    return cudaPeekAtLastError();
+}
+
+IPluginV3* HardmaxPlugin::attachToContext(IPluginResourceContext* context) noexcept
+{
+    auto* cloned = static_cast<HardmaxPlugin*>(clone());
+    if (cloned == nullptr)
+    {
+        return nullptr;
+    }
+    cublasStatus_t ret = cublasCreate(&cloned->mCublas);
+    ASSERT(ret == CUBLAS_STATUS_SUCCESS && cloned->mCublas != nullptr && "Failed to create cublasHandle_t.");
+    return cloned;
+}
+
+PluginFieldCollection const* HardmaxPlugin::getFieldsToSerialize() noexcept
+{
+    mDataToSerialize.clear();
+    mDataToSerialize.emplace_back("axis", &mAxis, PluginFieldType::kINT32, 1);
+    mFCToSerialize.nbFields = mDataToSerialize.size();
+    mFCToSerialize.fields = mDataToSerialize.data();
+    return &mFCToSerialize;
+}
+
+void HardmaxPlugin::setPluginNamespace(char const* libNamespace) noexcept
+{
+    ASSERT(libNamespace != nullptr);
+    mNamespace = libNamespace;
+}
+
+// HardmaxPluginCreator methods
+
+HardmaxPluginCreator::HardmaxPluginCreator()
+{
+    mPluginAttributes.clear();
+
+    // Consistent with the ONNX model attr fields
+    static auto const axisField = PluginField("axis", nullptr, PluginFieldType::kINT32, 1);
+    mPluginAttributes.emplace_back(axisField);
+
+    mFC.nbFields = mPluginAttributes.size();
+    mFC.fields = mPluginAttributes.data();
+}
+
+char const* HardmaxPluginCreator::getPluginName() const noexcept
+{
+    return kHARDMAX_NAME;
+}
+
+char const* HardmaxPluginCreator::getPluginVersion() const noexcept
+{
+    return kHARDMAX_VERSION;
+}
+
+PluginFieldCollection const* HardmaxPluginCreator::getFieldNames() noexcept
+{
+    return &mFC;
+}
+
+char const* HardmaxPluginCreator::getPluginNamespace() const noexcept
+{
+    return mNamespace.c_str();
+}
+
+void HardmaxPluginCreator::setPluginNamespace(char const* libNamespace) noexcept
+{
+    ASSERT(libNamespace != nullptr);
+    mNamespace = libNamespace;
+}
+
+IPluginV3* HardmaxPluginCreator::createPlugin(
+    char const* name, PluginFieldCollection const* fc, TensorRTPhase phase) noexcept
+{
+    using namespace std::string_view_literals;
+    // Set default value
+    int32_t axis = -1;
+
+    for (int32_t i = 0; i < fc->nbFields; i++)
+    {
+        if (fc->fields[i].name == "axis"sv)
+        {
+            ASSERT(fc->fields[i].type == PluginFieldType::kINT32);
+            axis = *static_cast<int32_t const*>(fc->fields[i].data);
+        }
+    }
+
+    auto plugin = std::make_unique<HardmaxPlugin>(axis);
+    plugin->setPluginNamespace(mNamespace.c_str());
+
+    return plugin.release();
+}

samples/python/onnx_custom_plugin/plugin/customHardmaxPlugin.h

@@ -0,0 +1,131 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2026 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.
+ */
+
+#ifndef TRT_HARDMAX_PLUGIN_H
+#define TRT_HARDMAX_PLUGIN_H
+
+#include "NvInferPlugin.h"
+#include <cublas_v2.h>
+#include <string>
+#include <vector>
+
+// This sample demonstrates how to implement a TensorRT plugin using the IPluginV3 interface.
+
+class HardmaxPlugin final : public nvinfer1::IPluginV3,
+                            public nvinfer1::IPluginV3OneCore,
+                            public nvinfer1::IPluginV3OneBuild,
+                            public nvinfer1::IPluginV3OneRuntime
+{
+public:
+    HardmaxPlugin() = delete;
+    HardmaxPlugin(int32_t axis);
+    HardmaxPlugin(HardmaxPlugin const& other);
+    ~HardmaxPlugin() override;
+
+    // IPluginV3 methods
+    nvinfer1::IPluginCapability* getCapabilityInterface(nvinfer1::PluginCapabilityType type) noexcept override;
+    nvinfer1::IPluginV3* clone() noexcept override;
+
+    // IPluginV3OneCore methods
+    char const* getPluginName() const noexcept override;
+    char const* getPluginVersion() const noexcept override;
+    char const* getPluginNamespace() const noexcept override;
+
+    // IPluginV3OneBuild methods
+    int32_t getNbOutputs() const noexcept override;
+
+    int32_t configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int32_t nbInputs,
+        nvinfer1::DynamicPluginTensorDesc const* out, int32_t nbOutputs) noexcept override;
+
+    bool supportsFormatCombination(int32_t pos, nvinfer1::DynamicPluginTensorDesc const* inOut, int32_t nbInputs,
+        int32_t nbOutputs) noexcept override;
+
+    int32_t getOutputDataTypes(nvinfer1::DataType* outputTypes, int32_t nbOutputs, nvinfer1::DataType const* inputTypes,
+        int32_t nbInputs) const noexcept override;
+
+    int32_t getOutputShapes(nvinfer1::DimsExprs const* inputs, int32_t nbInputs, nvinfer1::DimsExprs const* shapeInputs,
+        int32_t nbShapeInputs, nvinfer1::DimsExprs* outputs, int32_t nbOutputs,
+        nvinfer1::IExprBuilder& exprBuilder) noexcept override;
+
+    size_t getWorkspaceSize(nvinfer1::DynamicPluginTensorDesc const* inputs, int32_t nbInputs,
+        nvinfer1::DynamicPluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept override;
+
+    // IPluginV3OneRuntime methods
+    int32_t enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc,
+        void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override;
+
+    int32_t onShapeChange(nvinfer1::PluginTensorDesc const* in, int32_t nbInputs, nvinfer1::PluginTensorDesc const* out,
+        int32_t nbOutputs) noexcept override;
+
+    nvinfer1::IPluginV3* attachToContext(nvinfer1::IPluginResourceContext* context) noexcept override;
+
+    nvinfer1::PluginFieldCollection const* getFieldsToSerialize() noexcept override;
+
+    void setPluginNamespace(char const* pluginNamespace) noexcept;
+
+private:
+    std::string mNamespace;
+
+    // Number of elements in the axis along which hardmax is performed.
+    int32_t mAxisSize{0};
+
+    // Product of dimensions before and after mAxis.
+    // For example, if the input dimensions are [3, 4, 5, 6, 7] and mAxis = 2,
+    // then mDimProductOuter = 12 and mDimProductInner = 42.
+    int32_t mDimProductOuter{1};
+    int32_t mDimProductInner{1};
+
+    cublasHandle_t mCublas{nullptr};
+
+    // Attributes
+    // Axis along which to perform hardmax.
+    // Can be negative initially, but once configurePlugin() is called it will
+    // be converted to a positive axis.
+    int32_t mAxis{-1};
+
+    // Serialization helpers
+    std::vector<nvinfer1::PluginField> mDataToSerialize;
+    nvinfer1::PluginFieldCollection mFCToSerialize;
+};
+
+class HardmaxPluginCreator : public nvinfer1::IPluginCreatorV3One
+{
+public:
+    HardmaxPluginCreator();
+
+    ~HardmaxPluginCreator() override = default;
+
+    char const* getPluginName() const noexcept override;
+
+    char const* getPluginVersion() const noexcept override;
+
+    nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override;
+
+    nvinfer1::IPluginV3* createPlugin(
+        char const* name, nvinfer1::PluginFieldCollection const* fc, nvinfer1::TensorRTPhase phase) noexcept override;
+
+    void setPluginNamespace(char const* pluginNamespace) noexcept;
+
+    char const* getPluginNamespace() const noexcept override;
+
+private:
+    nvinfer1::PluginFieldCollection mFC;
+    std::vector<nvinfer1::PluginField> mPluginAttributes;
+    std::string mNamespace;
+};
+
+#endif // TRT_HARDMAX_PLUGIN_H

samples/python/onnx_custom_plugin/requirements.txt

@@ -0,0 +1,11 @@
+nltk==3.9.1
+onnx==1.18.0
+--extra-index-url https://pypi.ngc.nvidia.com
+onnx-graphsurgeon>=0.3.20
+wget>=3.2
+cuda-python==12.9.0
+pywin32; platform_system == "Windows"
+pyyaml==6.0.3
+requests==2.32.4
+tqdm==4.66.4
+numpy==1.26.4

samples/python/onnx_custom_plugin/sample.py

@@ -0,0 +1,193 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2026 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 sys
+
+import tensorrt as trt
+
+from model import TRT_MODEL_PATH
+from load_plugin_lib import load_plugin_lib
+
+# ../common.py
+parent_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), os.pardir)
+sys.path.insert(1, parent_dir)
+import common
+
+# Reuse some BiDAF-specific methods
+# ../engine_refit_onnx_bidaf/data_processing.py
+sys.path.insert(1, os.path.join(parent_dir, "engine_refit_onnx_bidaf"))
+from engine_refit_onnx_bidaf.data_processing import preprocess, get_inputs
+
+# Maxmimum number of words in context or query text.
+# Used in optimization profile when building engine.
+# Adjustable.
+MAX_TEXT_LENGTH = 64
+
+WORKING_DIR = os.environ.get("TRT_WORKING_DIR") or os.path.dirname(
+    os.path.realpath(__file__)
+)
+
+# Path to which trained model will be saved (check README.md)
+ENGINE_FILE_PATH = os.path.join(WORKING_DIR, "bidaf.trt")
+
+# Define global logger object (it should be a singleton,
+# available for TensorRT from anywhere in code).
+# You can set the logger severity higher to suppress messages
+# (or lower to display more messages)
+TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
+
+
+# Builds TensorRT Engine
+def build_engine(model_path):
+
+    builder = trt.Builder(TRT_LOGGER)
+    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED))
+    config = builder.create_builder_config()
+    parser = trt.OnnxParser(network, TRT_LOGGER)
+    runtime = trt.Runtime(TRT_LOGGER)
+
+    # Parse model file
+    print("Loading ONNX file from path {}...".format(model_path))
+    with open(model_path, "rb") as model:
+        print("Beginning ONNX file parsing")
+        if not parser.parse(model.read()):
+            print("ERROR: Failed to parse the ONNX file.")
+            for error in range(parser.num_errors):
+                print(parser.get_error(error))
+            return None
+    print("Completed parsing of ONNX file")
+
+    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, common.GiB(1))
+
+    # The input text length is variable, so we need to specify an optimization profile.
+    profile = builder.create_optimization_profile()
+    for i in range(network.num_inputs):
+        input = network.get_input(i)
+        assert input.shape[0] == -1
+        min_shape = [1] + list(input.shape[1:])
+        opt_shape = [8] + list(input.shape[1:])
+        max_shape = [MAX_TEXT_LENGTH] + list(input.shape[1:])
+        profile.set_shape(input.name, min_shape, opt_shape, max_shape)
+    config.add_optimization_profile(profile)
+
+    print("Building TensorRT engine. This may take a few minutes.")
+    plan = builder.build_serialized_network(network, config)
+    engine = runtime.deserialize_cuda_engine(plan)
+    with open(ENGINE_FILE_PATH, "wb") as f:
+        f.write(plan)
+    return engine
+
+
+def load_test_case(inputs, context_text, query_text, trt_context):
+    # Part 1: Specify Input shapes
+    cw, cc = preprocess(context_text)
+    qw, qc = preprocess(query_text)
+    for arr in (cw, cc, qw, qc):
+        assert arr.shape[0] <= MAX_TEXT_LENGTH, (
+            "Input context or query is too long! "
+            + "Either decrease the input length or increase MAX_TEXT_LENGTH"
+        )
+    trt_context.set_input_shape("CategoryMapper_4", cw.shape)
+    trt_context.set_input_shape("CategoryMapper_5", cc.shape)
+    trt_context.set_input_shape("CategoryMapper_6", qw.shape)
+    trt_context.set_input_shape("CategoryMapper_7", qc.shape)
+
+    # Part 2: load input data
+    cw_flat, cc_flat, qw_flat, qc_flat = get_inputs(context_text, query_text)
+    for i, arr in enumerate([cw_flat, cc_flat, qw_flat, qc_flat]):
+        inputs[i].host = arr
+
+
+def main():
+    # Load the shared object file containing the Hardmax plugin implementation.
+    # By doing this, you will also register the Hardmax plugin with the TensorRT
+    # PluginRegistry through use of the macro REGISTER_TENSORRT_PLUGIN present
+    # in the plugin implementation. Refer to plugin/customHardmaxPlugin.cpp for more details.
+    load_plugin_lib()
+
+    # Load pretrained model
+    if not os.path.isfile(TRT_MODEL_PATH):
+        raise IOError(
+            "\n{}\n{}\n{}\n".format(
+                "Failed to load model file ({}).".format(TRT_MODEL_PATH),
+                "Please use 'python3 model.py' to generate the ONNX model.",
+                "For more information, see README.md",
+            )
+        )
+
+    if os.path.exists(ENGINE_FILE_PATH):
+        print(f"Loading saved TRT engine from {ENGINE_FILE_PATH}")
+        with open(ENGINE_FILE_PATH, "rb") as f:
+            runtime = trt.Runtime(TRT_LOGGER)
+            runtime.max_threads = 10
+            engine = runtime.deserialize_cuda_engine(f.read())
+    else:
+        print("Engine plan not saved. Building new engine...")
+        engine = build_engine(TRT_MODEL_PATH)
+
+    inputs, outputs, bindings = common.allocate_buffers(engine, profile_idx=0)
+
+    testcases = [
+        (
+            "Garry the lion is 5 years old. He lives in the savanna.",
+            "Where does the lion live?",
+        ),
+        ("A quick brown fox jumps over the lazy dog.", "What color is the fox?"),
+    ]
+
+    print("\n=== Testing ===")
+
+    interactive = "--interactive" in sys.argv
+    if interactive:
+        context_text = input("Enter context: ")
+        query_text = input("Enter query: ")
+        testcases = [(context_text, query_text)]
+
+    trt_context = engine.create_execution_context()
+    
+    # Use context manager for proper stream lifecycle management
+    with common.CudaStreamContext() as stream:
+        for context_text, query_text in testcases:
+
+            context_words, _ = preprocess(context_text)
+
+            load_test_case(inputs, context_text, query_text, trt_context)
+            if not interactive:
+                print(f"Input context: {context_text}")
+                print(f"Input query: {query_text}")
+            trt_outputs = common.do_inference(
+                trt_context,
+                engine=engine,
+                bindings=bindings,
+                inputs=inputs,
+                outputs=outputs,
+                stream=stream,
+            )
+            start = trt_outputs[1].item()
+            end = trt_outputs[0].item()
+            answer = context_words[start : end + 1].flatten()
+            print(f"Model prediction: ", " ".join(answer))
+            print()
+        
+        # Note: free_buffers no longer needs stream parameter
+        common.free_buffers(inputs, outputs)
+    print("Passed")
+
+
+if __name__ == "__main__":
+    main()

samples/python/onnx_custom_plugin/test_custom_hardmax_plugin.py

@@ -0,0 +1,102 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2026 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 os
+import sys
+import tensorrt as trt
+
+# ../common.py
+parent_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), os.pardir)
+sys.path.insert(1, parent_dir)
+import common
+
+from load_plugin_lib import load_plugin_lib
+
+TRT_LOGGER = trt.Logger(trt.Logger.ERROR)
+
+
+def hardmax_reference_impl(arr, axis):
+    one_hot = np.zeros(arr.shape, dtype=arr.dtype)
+    argmax = np.expand_dims(np.argmax(arr, axis), axis)
+    np.put_along_axis(one_hot, argmax, 1, axis=axis)
+    return one_hot
+
+
+def make_trt_network_and_engine(input_shape, axis):
+    registry = trt.get_plugin_registry()
+    plugin_creator = registry.get_creator("CustomHardmax", "1", "")
+    axis_buffer = np.array([axis])
+    axis_attr = trt.PluginField("axis", axis_buffer, type=trt.PluginFieldType.INT32)
+    field_collection = trt.PluginFieldCollection([axis_attr])
+    plugin = plugin_creator.create_plugin(
+        name="CustomHardmax", field_collection=field_collection, phase=trt.TensorRTPhase.BUILD
+    )
+
+    builder = trt.Builder(TRT_LOGGER)
+    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED))
+    config = builder.create_builder_config()
+    runtime = trt.Runtime(TRT_LOGGER)
+
+    input_layer = network.add_input(
+        name="input_layer", dtype=trt.float32, shape=input_shape
+    )
+    hardmax = network.add_plugin_v3(inputs=[input_layer], shape_inputs=[], plugin=plugin)
+    network.mark_output(hardmax.get_output(0))
+
+    plan = builder.build_serialized_network(network, config)
+    engine = runtime.deserialize_cuda_engine(plan)
+
+    return engine
+
+
+def custom_plugin_impl(input_arr, engine):
+    inputs, outputs, bindings = common.allocate_buffers(engine)
+    context = engine.create_execution_context()
+    inputs[0].host = input_arr.astype(trt.nptype(trt.float32))
+    
+    # Use context manager for proper stream lifecycle management
+    with common.CudaStreamContext() as stream:
+        trt_outputs = common.do_inference(
+            context,
+            engine=engine,
+            bindings=bindings,
+            inputs=inputs,
+            outputs=outputs,
+            stream=stream,
+        )
+        output = trt_outputs[0].copy()
+    common.free_buffers(inputs, outputs)
+    return output
+
+
+def main():
+    load_plugin_lib()
+    for num_dims in range(1, 8):
+        for axis in range(-num_dims, num_dims):
+            shape = np.random.randint(1, 4, size=num_dims)
+            arr = np.random.rand(*shape)
+            arr = (arr - 0.5) * 200
+            engine = make_trt_network_and_engine(shape, axis)
+            res1 = hardmax_reference_impl(arr, axis)
+            res2 = custom_plugin_impl(arr, engine).reshape(res1.shape)
+            assert np.all(res1 == res2), f"Test failed for shape={shape}, axis={axis}"
+    print("Passed")
+
+
+if __name__ == "__main__":
+    main()

samples/python/onnx_packnet/README.md

@@ -0,0 +1,112 @@
+# TensorRT Inference of ONNX models with custom layers.
+
+**Table Of Contents**
+- [Description](#description)
+- [How does this sample work?](#how-does-this-sample-work)
+- [Prerequisites](#prerequisites)
+- [Running the sample](#running-the-sample)
+    * [Cloning the packnet repository](#cloning-the-packnet-repository)
+    * [Conversion to ONNX](#conversion-to-onnx)
+    * [Inference with TensorRT](#inference-with-tensorrt)
+	* [Sample `--help` options](#sample-help-options)
+- [Additional resources](#additional-resources)
+- [License](#license)
+- [Changelog](#changelog)
+- [Known issues](#known-issues)
+
+## Description
+
+This sample, samplePackNet, is a Python sample which uses TensorRT to perform inference with PackNet network. PackNet is a self-supervised monocular depth estimation network used in autonomous driving.
+
+
+## How does this sample work?
+
+This sample converts the Pytorch graph into ONNX and uses ONNX-parser included in TensorRT to parse the ONNX graph. The sample also demonstrates
+
+* Use of custom layers (plugins) in ONNX graph. These plugins would be automatically registered in TensorRT by using `REGISTER_TENSORRT_PLUGIN` API.
+* Use of ONNX-graphsurgeon (ONNX-GS) API to modify layers or subgraphs in the ONNX graph. For this network, we transform Group Normalization, upsample and pad layers to remove unnecessary
+  nodes for inference with TensorRT.
+
+
+## Prerequisites
+
+1. Upgrade pip version and install the sample dependencies.
+    ```bash
+    pip3 install --upgrade pip
+    pip3 install -r requirements.txt
+    ```
+
+On PowerPC systems, you will need to manually install PyTorch using IBM's [PowerAI](https://www.ibm.com/support/knowledgecenter/SS5SF7_1.6.0/navigation/pai_install.htm).
+
+
+## Running the sample
+
+### Preparing packnet
+
+Clone the [packnet](https://github.com/TRI-ML/packnet-sfm) repository and update `PYTHONPATH`.
+
+```
+git clone https://github.com/TRI-ML/packnet-sfm.git packnet-sfm
+pushd packnet-sfm && git checkout tags/v0.1.2 && popd
+export PYTHONPATH=$PWD/packnet-sfm # Note on Windows, the export command is: set PYTHONPATH=%cd%\packnet-sfm
+```
+
+### Conversion to ONNX
+Run the following command to convert the Packnet pytorch network to ONNX graph. This step also includes handling custom layers (Group Normalization) and using ONNX-GS to modify upsample and pad layers.
+
+```
+python3 convert_to_onnx.py --output model.onnx
+```
+
+### Inference with TensorRT
+
+Once the ONNX graph is generated, use `trtexec` tool (located in `bin` directory of TensorRT package) to perform inference on a random input image.
+
+```
+trtexec --onnx=model.onnx
+```
+
+Please refer to `trtexec` tool for more commandline options.
+
+### Sample --help options
+
+To see the full list of available options and their descriptions, use the `-h` or `--help` command line option. For example:
+```
+convert_to_onnx.py -h
+```
+
+# Additional resources
+
+The following resources provide a deeper understanding about PackNet network and importing a model into TensorRT using Python:
+
+**PackNet**
+- [3D Packing for Self-Supervised Monocular Depth Estimation](https://arxiv.org/pdf/1905.02693.pdf)
+- [TRI-ML Monocular Depth Estimation Repository](https://github.com/TRI-ML/packnet-sfm)
+
+**Parsers**
+- [ONNX Parser](https://docs.nvidia.com/deeplearning/sdk/tensorrt-api/python_api/parsers/Onnx/pyOnnx.html)
+
+**Documentation**
+- [Introduction To NVIDIA’s TensorRT Samples](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sample-support-guide/index.html#samples)
+- [Working With TensorRT Using The Python API](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#python_topics)
+- [Importing A Model Using A Parser In Python](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#import_model_python)
+- [NVIDIA’s TensorRT Documentation Library](https://docs.nvidia.com/deeplearning/sdk/tensorrt-archived/index.html)
+
+# License
+
+For terms and conditions for use, reproduction, and distribution, see the [TensorRT Software License Agreement](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sla/index.html) documentation.
+
+# Changelog
+
+August 2025:
+  - Removed support for Python versions < 3.10.
+
+August 2023:
+  - Update ONNX version support to 1.14.0
+  - Removed support for Python versions < 3.8.
+August 2021: Update sample to work with latest torch version
+June 2020: Initial release of this sample
+
+# Known issues
+
+There are no known issues in this sample

samples/python/onnx_packnet/convert_to_onnx.py

@@ -0,0 +1,100 @@
+#!/usr/bin/env python3
+#
+# 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 onnx
+import torch
+import numpy as np
+import argparse
+import onnx_graphsurgeon as gs
+from post_processing import *
+from packnet_sfm.networks.depth.PackNet01 import PackNet01
+
+
+def post_process_packnet(model_file, opset=11):
+    """
+    Use ONNX graph surgeon to replace upsample and instance normalization nodes. Refer to post_processing.py for details.
+    Args:
+        model_file : Path to ONNX file
+    """
+    # Load the packnet graph
+    graph = gs.import_onnx(onnx.load(model_file))
+
+    if opset >= 11:
+        graph = process_pad_nodes(graph)
+
+    # Replace the subgraph of upsample with a single node with input and scale factor.
+    if torch.__version__ < "1.5.0":
+        graph = process_upsample_nodes(graph, opset)
+
+    # Convert the group normalization subgraph into a single plugin node.
+    graph = process_groupnorm_nodes(graph)
+
+    # Remove unused nodes, and topologically sort the graph.
+    graph.cleanup().toposort()
+
+    # Export the onnx graph from graphsurgeon
+    onnx.save_model(gs.export_onnx(graph), model_file)
+
+    print("Saving the ONNX model to {}".format(model_file))
+
+
+def build_packnet(model_file, args):
+    """
+    Construct the packnet network and export it to ONNX
+    """
+    input_pyt = torch.randn((1, 3, 192, 640), requires_grad=False)
+
+    # Build the model
+    model_pyt = PackNet01(version="1A")
+
+    # Convert the model into ONNX
+    torch.onnx.export(
+        model_pyt, input_pyt, model_file, verbose=args.verbose, opset_version=args.opset
+    )
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Exports PackNet01 to ONNX, and post-processes it to insert TensorRT plugins"
+    )
+    parser.add_argument(
+        "-o",
+        "--output",
+        help="Path to save the generated ONNX model",
+        default="model.onnx",
+    )
+    parser.add_argument(
+        "-op", "--opset", type=int, help="ONNX opset to use", default=11
+    )
+    parser.add_argument(
+        "-v",
+        "--verbose",
+        action="store_true",
+        help="Flag to enable verbose logging for torch.onnx.export",
+    )
+    args = parser.parse_args()
+
+    # Construct the packnet graph and generate the onnx graph
+    build_packnet(args.output, args)
+
+    # Perform post processing on Instance Normalization and upsampling nodes and create a new ONNX graph
+    post_process_packnet(args.output, args.opset)
+
+
+if __name__ == "__main__":
+    main()

samples/python/onnx_packnet/download.yml

@@ -0,0 +1,21 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 2020-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.
+#
+sample: onnx_packnet
+files:
+  - path: samples/python/onnx_packnet/packnet-sfm-0.1.2.zip
+    url: https://github.com/TRI-ML/packnet-sfm/archive/v0.1.2.zip
+    checksum: 7a73db591d3955ccf407910cd928d9c0

samples/python/onnx_packnet/post_processing.py

@@ -0,0 +1,248 @@
+#!/usr/bin/env python3
+#
+# 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 onnx_graphsurgeon as gs
+import argparse
+import onnx
+import numpy as np
+import torch
+
+
+# Pad layer subgraph structure in ONNX (specific to opset 11):
+#               Constant
+#                  |
+#                Shape
+#                  |
+#         Mul   Gather
+#          \     /
+#            Sub
+#             |
+#       ConstantOfShape
+#             |
+#          Concat
+#             |
+#          Reshape
+#             |
+#           Slice
+#             |
+#          Transpose
+#             |
+#          Reshape
+#             |
+#    Input  Cast  Constant
+#       \     |    /
+#            Pad
+def process_pad_nodes(graph):
+    """
+    Fold the pad subgraph into a single layer with pad values as input
+      Input
+       |
+      Pad
+       |
+      Conv
+    """
+    pad_nodes = [node for node in graph.nodes if node.op == "Pad"]
+    for node in pad_nodes:
+        fold_pad_inputs(node, graph)
+
+    return graph
+
+
+def fold_pad_inputs(node, graph):
+    # Gather the amount of padding in each dimension from pytorch graph.
+    if torch.__version__ < "1.5.0":
+        pad_values_pyt = (
+            node.i(1).i(0).i(0).i(0).i(0).i(0).i(0).i(0).attrs["value"].values
+        )
+    elif torch.__version__ < "2.0.0":
+        pad_values_pyt = node.i(1).i(0).i(0).i(0).i(0).i(0).inputs[0].values
+    else:
+        pad_values_pyt = node.i(1).i(0).i(0).i(0).i(0).i(0).i(0).attrs["value"].values
+
+    # Assumption a 4d input tensor
+    onnx_pad_values = [0] * 4 * 2  # 4d tensor and 2 sides padding for each dimension
+    j = 3
+    for i in range(0, len(pad_values_pyt), 2):
+        onnx_pad_values[j] = pad_values_pyt[i]
+        onnx_pad_values[j + 4] = pad_values_pyt[i + 1]
+        j -= 1
+
+    # Change the existing pad tensor to the new onnx_pad values tensor
+    pads_folded_tensor = gs.Constant(
+        name=node.inputs[1].name, values=np.array(onnx_pad_values)
+    )
+    node.inputs[1] = pads_folded_tensor
+
+
+# Pytorch-exported Upsample structure in ONNX:
+#        Mul        Mul
+#         |          |
+#        Cast       Cast
+#         |          |
+#        Floor      Floor
+#         |          |
+#      Unsqueeze  Unsqueeze
+#         \         /
+#           Concat
+#             |
+#            Cast    Cast
+#              \      /
+#                Div
+#                 |
+#     Input     Concat
+#       \         /
+#         Upsample
+def process_upsample_nodes(graph, opset=11):
+    """
+    Replace the upsample structure with structure below
+      Conv   scale_factor
+       |      /
+      Upsample
+       |
+      ReLU
+    """
+    if opset >= 11:
+        upsample_layer_name = "Resize"
+    else:
+        upsample_layer_name = "Upsample"
+
+    upsample_nodes = [node for node in graph.nodes if node.op == upsample_layer_name]
+    for node in upsample_nodes:
+        fold_upsample_inputs(node, graph, opset)
+
+    return graph
+
+
+def fold_upsample_inputs(upsample, graph, opset=11):
+    """
+    Inplace transformation of the graph. The upsample subgraph is collapsed
+    to single upsample node with input and scale factor (constant tensor).
+    Args:
+        upsample: upsample node in the original graph.
+        graph: graph object.
+    """
+
+    if opset == 9:
+        # Gather the scale factor from mul op in the upsample input subgraph
+        scale_factor = (
+            upsample.i(1).i(1).i(0).i(0).i(0).i(0).i(0).i(0).i(1).attrs["value"].values
+        )
+
+        # Create the new scales tensor
+        scales = np.array([1.0, 1.0, scale_factor, scale_factor], dtype=np.float32)
+        scale_tensor = gs.Constant(name=upsample.inputs[-1].name, values=scales)
+
+        # Change the last input to the node to the new constant scales tensor.
+        upsample.inputs[-1] = scale_tensor
+    else:
+        # In opset 11, upsample layer is exported as Resize. We will transform this Resize layer into an Upsample layer
+        # and collapse the input
+        sizes_tensor_name = upsample.inputs[3].name
+
+        # Create the new scales tensor
+        scale_factor = (
+            upsample.i(3).i(1).i().i().i().i().i(0).i(1).attrs["value"].values
+        )
+        scales = np.array([1.0, 1.0, scale_factor, scale_factor], dtype=np.float32)
+        scale_tensor = gs.Constant(name=sizes_tensor_name, values=scales)
+
+        # Rename the Resize op to upsample and add the data and scales as inputs to the upsample layer.
+        input_tensor = upsample.inputs[0]
+        upsample.inputs = [input_tensor, scale_tensor]
+        upsample.op = "Upsample"
+
+
+# Pytorch-exported GroupNorm subgraph in ONNX:
+# Conv
+#   |
+# Reshape    Scale    Bias
+#     \       |       /
+# InstanceNormalization
+#         |
+#      Reshape    Unsqueeze
+#          \      /
+#        Mul (scale)   Unsqueeze
+#           \         /
+#           Add (bias)
+#              |
+#            ReLU
+def process_groupnorm_nodes(graph):
+    """
+    Gather the instance normalization nodes and the rest of the subgraph
+    and convert into a single group normalization node.
+    """
+    instancenorms = [node for node in graph.nodes if node.op == "InstanceNormalization"]
+    for node in instancenorms:
+        convert_to_groupnorm(node, graph)
+
+    return graph
+
+
+def retrieve_attrs(instancenorm):
+    """
+    Gather the required attributes for the GroupNorm plugin from the subgraph.
+    Args:
+        instancenorm: Instance Normalization node in the graph.
+    """
+    attrs = {}
+    # The 2nd dimension of the Reshape shape is the number of groups
+    attrs["num_groups"] = instancenorm.i().i(1).attrs["value"].values[1]
+    attrs["eps"] = instancenorm.attrs["epsilon"]
+
+    # 1 is the default plugin version the parser will search for, and therefore can be omitted,
+    # but we include it here for illustrative purposes.
+    attrs["plugin_version"] = "1"
+
+    # "" is the default plugin namespace the parser will use, included here for illustrative purposes
+    attrs["plugin_namespace"] = ""
+
+    return attrs
+
+
+def convert_to_groupnorm(instancenorm, graph):
+    """
+    Convert the Pytorch-exported GroupNorm subgraph to the subgraph below
+    Conv
+      |
+    GroupNorm
+      |
+    ReLU
+    Attributes:
+        instancenorm: Instance Normalization node in the graph.
+        graph: Input graph object
+    """
+    # Retrieve the instancenorm attributes and create the replacement node
+    attrs = retrieve_attrs(instancenorm)
+    groupnorm = gs.Node(op="GroupNormalizationPlugin", attrs=attrs)
+    graph.nodes.append(groupnorm)
+
+    # The plugin needs to receive an input from the Conv node, and output to the ReLU node
+    conv_output_tensor = instancenorm.i().inputs[0]  # Output of Conv
+    relu_input_tensor = instancenorm.o().o().o().outputs[0]  # Output of Add
+
+    # Reconnect inputs/outputs to the groupnorm plugin
+    conv_output_tensor.outputs[0] = groupnorm
+    relu_input_tensor.inputs[0] = groupnorm
+
+    # Add scale and bias constant tensors to group norm plugin
+    if torch.__version__ < "1.5.0":
+        groupnorm.inputs.append(instancenorm.o().o().i(1).inputs[0])
+        groupnorm.inputs.append(instancenorm.o().o().o().i(1).inputs[0])
+    else:
+        groupnorm.inputs.append(instancenorm.o().o().inputs[1])
+        groupnorm.inputs.append(instancenorm.o().o().o().inputs[1])

samples/python/onnx_packnet/requirements.txt

@@ -0,0 +1,9 @@
+onnx==1.18.0
+--extra-index-url https://pypi.ngc.nvidia.com
+onnx-graphsurgeon>=0.3.20
+torch
+torchvision
+pyyaml==6.0.3
+requests==2.32.4
+tqdm==4.66.4
+numpy==1.26.4

samples/python/plugin_utils.py

@@ -0,0 +1,135 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2026 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 cuda.bindings import driver as cuda, runtime as cudart, nvrtc
+
+import numpy as np
+import os
+from common_runtime import cuda_call, create_cuda_context, cuda_init, cuda_get_device, cuda_memcpy_htod
+import argparse
+import threading
+
+import tensorrt as trt
+import cupy as cp
+
+
+def parseArgs():
+    parser = argparse.ArgumentParser(
+        description="Options for Circular Padding plugin C++ example"
+    )
+    parser.add_argument(
+        "--precision",
+        type=str,
+        default="fp32",
+        choices=["fp32", "fp16"],
+        help="Precision to use for plugin",
+    )
+
+    return parser.parse_args()
+
+
+def volume(d):
+    return np.prod(d)
+
+
+def getComputeCapacity(devID):
+    major = cuda_call(cudart.cudaDeviceGetAttribute(cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMajor, devID))
+    minor = cuda_call(cudart.cudaDeviceGetAttribute(cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMinor, devID))
+    return (major, minor)
+
+
+# Taken from https://github.com/NVIDIA/cuda-python/blob/main/examples/common/common.py
+class KernelHelper:
+    def __init__(self, code, devID):
+        prog = cuda_call(
+            nvrtc.nvrtcCreateProgram(str.encode(code), b"sourceCode.cu", 0, [], [])
+        )
+        cuda_root = None
+        for env_name in ("CUDA_PATH", "CUDA_HOME"):
+            cand = os.getenv(env_name)
+            if cand and os.path.isfile(os.path.join(cand, "include", "cuda_fp16.h")):
+                cuda_root = cand
+                break
+        if cuda_root is None:
+            raise RuntimeError(
+                "Neither CUDA_PATH nor CUDA_HOME points at a CUDA install containing include/cuda_fp16.h"
+            )
+        include_dirs = os.path.join(cuda_root, "include")
+
+        # Initialize CUDA
+        cuda_call(cudart.cudaFree(0))
+
+        major, minor = getComputeCapacity(devID)
+        _, nvrtc_minor = cuda_call(nvrtc.nvrtcVersion())
+        use_cubin = nvrtc_minor >= 1
+        prefix = "sm" if use_cubin else "compute"
+        arch_arg = bytes(f"--gpu-architecture={prefix}_{major}{minor}", "ascii")
+
+        try:
+            opts = [
+                b"--fmad=true",
+                arch_arg,
+                ('-I' + include_dirs).encode("UTF-8"),
+                b"--std=c++11",
+                b"-default-device",
+            ]
+            cuda_call(nvrtc.nvrtcCompileProgram(prog, len(opts), opts))
+        except RuntimeError as err:
+            logSize = cuda_call(nvrtc.nvrtcGetProgramLogSize(prog))
+            log = b" " * logSize
+            cuda_call(nvrtc.nvrtcGetProgramLog(prog, log))
+            print(log.decode())
+            print(err)
+            exit(-1)
+
+        if use_cubin:
+            dataSize = cuda_call(nvrtc.nvrtcGetCUBINSize(prog))
+            data = b" " * dataSize
+            cuda_call(nvrtc.nvrtcGetCUBIN(prog, data))
+        else:
+            dataSize = cuda_call(nvrtc.nvrtcGetPTXSize(prog))
+            data = b" " * dataSize
+            cuda_call(nvrtc.nvrtcGetPTX(prog, data))
+
+        self.module = cuda_call(cuda.cuModuleLoadData(np.char.array(data)))
+
+    def getFunction(self, name):
+        return cuda_call(cuda.cuModuleGetFunction(self.module, name))
+
+
+class CudaCtxManager(trt.IPluginResource):
+    def __init__(self, device=None):
+        trt.IPluginResource.__init__(self)
+        self.device = device
+        self.cuda_ctx = None
+
+    def clone(self):
+        cloned = CudaCtxManager()
+        cloned.__dict__.update(self.__dict__)
+        # Delay the CUDA ctx creation until clone()
+        # since only a cloned resource is registered by TRT
+        cloned.cuda_ctx = create_cuda_context(self.device)
+        return cloned
+
+    def release(self):
+        cuda_call(cuda.cuCtxDestroy(self.cuda_ctx))
+
+class UnownedMemory:
+    def __init__(self, ptr, shape, dtype):
+        mem = cp.cuda.UnownedMemory(ptr, volume(shape) * cp.dtype(dtype).itemsize, self)
+        cupy_ptr = cp.cuda.MemoryPointer(mem, 0)
+        self.d = cp.ndarray(shape, dtype=dtype, memptr=cupy_ptr)

samples/python/python_plugin/CMakeLists.txt

@@ -0,0 +1,75 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2026 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.
+#
+
+cmake_minimum_required(VERSION 3.31 FATAL_ERROR)
+
+# Default to building for the local GPU. Must be set BEFORE project(),
+# otherwise project()'s CUDA language enablement initializes
+# CMAKE_CUDA_ARCHITECTURES to the toolkit's lowest supported SM
+# (typically sm_75), and any later `if(NOT DEFINED)` guard misfires.
+# `all` is not a substitute: it caps at sm_90 for CUDA >= 12.0
+# (https://github.com/Kitware/CMake/blob/v3.31.6/Modules/Internal/CMakeCUDAArchitecturesAll.cmake#L65-L77),
+# missing Blackwell/Drive-Thor/Spark, and the PTX fallback fails on
+# embedded/safety drivers (e.g. DriveOS QNX) that lack JIT.
+if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES)
+    set(CMAKE_CUDA_ARCHITECTURES native)
+endif()
+
+project(CircPadPlugin LANGUAGES CXX CUDA)
+
+find_package(CUDAToolkit REQUIRED)
+
+if(NOT MSVC)
+    # Enable all compile warnings
+    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall -Wno-long-long -pedantic -Wno-deprecated-declarations")
+    set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -Xcompiler -Wno-deprecated-declarations")
+endif()
+
+# Sets variable to a value if variable is unset.
+macro(set_ifndef var val)
+    if(NOT ${var})
+        set(${var} ${val})
+    endif()
+    message(STATUS "Configurable variable ${var} set to ${${var}}")
+endmacro()
+
+# -------- CONFIGURATION --------
+if(NOT MSVC)
+    set_ifndef(TRT_LIB /usr/lib/x86_64-linux-gnu)
+    set_ifndef(TRT_INCLUDE /usr/include/x86_64-linux-gnu)
+endif()
+
+message("\nThe following variables are derived from the values of the previous variables unless provided explicitly:\n")
+
+find_library(
+    _NVINFER_LIB nvinfer
+    HINTS ${TRT_LIB}
+    PATH_SUFFIXES lib lib64)
+set_ifndef(NVINFER_LIB ${_NVINFER_LIB})
+
+# -------- BUILDING --------
+
+add_library(circ_pad_plugin SHARED ${CMAKE_SOURCE_DIR}/circ_plugin_cpp/circ_pad_plugin.cu)
+target_include_directories(
+    circ_pad_plugin
+    PUBLIC ${CUDAToolkit_INCLUDE_DIRS}
+    PUBLIC ${TRT_INCLUDE})
+
+set_property(TARGET circ_pad_plugin PROPERTY CUDA_STANDARD 17)
+
+target_link_libraries(circ_pad_plugin PRIVATE ${NVINFER_LIB})
+target_link_libraries(circ_pad_plugin PRIVATE CUDA::cuda_driver)

samples/python/python_plugin/README.md

@@ -0,0 +1,171 @@
+# Python-based TRT Plugins
+
+This is a sample to showcase Python-based plugin definitions in TRT. No changes to existing TRT APIs have been made
+to deliver this feature, so using the updated bindings should not break any existing code.
+
+## Introduction
+
+Until TRT 9.1, plugin implementations could only be done through the TRT C++ API. To use a plugin in a Python app, one had to
+ - Implement plugin in C++ and build into a shared library
+ - Load plugin lib and register plugin creator (statically or dynamically)
+ - Retrieve plugin creator and create plugin instance through the respective Python API
+
+The following design considerations were followed in creating bindings to allow Python-based plugin definitions:
+ - Zero additional C++ code shall be required to implement, integrate and run a plugin within TensorRT
+ - Offer the flexibility to implement the kernel(s) for the plugin through any method of choice
+   - Many libraries have sprung up to provide CUDA kernel support with AOT/JIT compilation
+     - Numba, OpenAI Triton, CuPy etc.
+   - Could even do without explicit kernels (e.g. leverage PyTorch functional op)
+
+ - Will only support `IPluginV2DynamicExt` and `IPluginV3`-based plugins
+   - Other plugin interfaces (except `IPluginV2IOExt`) are deprecated since TRT 8.5
+
+With these bindings, plugins can be implemented and integrated to TRT purely with Python.
+
+## Setting Up The Build Environment
+
+To build and install the bindings, follow the instructions in `$TRT_OSSPATH/python/README.md`.
+
+Then install the requisite packages
+```bash
+cd $TRT_OSSPATH/samples/python/trt_python_plugin
+pip3 install -r requirements.txt
+```
+Install `cupy-cuda11x` instead if testing on a CUDA 11.x environment.
+
+# TensorRT Plugin API for Python
+
+Implementing a TRT plugin in Python is similar to C++ in that implementation of `IPluginV2DynamicExt`+`IPluginCreator` or `IPluginV3`+`IPluginCreatorV3One` is necessary. Refer to the TensorRT Python API reference for a concise description.
+
+## Differences in C++ and Python APIs for `IPluginV2DynamicExt`
+The interface methods in Python have mostly similar APIs to their C++ counterparts, except for `serialize()` and `enqueue()`.
+ - While the C++ API for `serialize()` is `void serialize (void *buffer)` where the plugin writes to the passed-in `buffer`, the Python API is `serialize(self) -> bytes`, where the implementation of the method is expected to return a bytes object containing a serialized representation of the plugin object.
+ - In `enqueue()`, the device pointers for input and output tensors are passed as their `intptr_t` casts. Since these buffers are created and owned by TRT, care must be taken when writing to them from the Python side.
+  - No bindings yet for `attachToContext()` and `detachFromContext()` which are not pure virtual.
+
+# Running the sample: Circular padding plugin
+
+This sample contains a circular padding plugin, where the `enqueue` has been implemented with various frameworks for writing kernels or executing GPU ops (torch).
+
+Each script accepts a command-line argument to choose precision from either FP32 or FP16. e.g.
+```bash
+python3 circ_pad_plugin_cuda_python.py --precision fp32 # fp32 or fp16
+```
+
+## Circular padding
+
+Circular padding is useful for ops like circular convolution in deep learning. The following image denotes how the original image (red) is circular padded once (green) and twice (blue):
+
+![alt text](circ_pad_example.png "Circular padding example")
+
+The plugin shall have the following characteristics:
+ - Input: 4-dimensional input (e.g. NxCxHxW)
+ - Attribute(s): m-dimensional parameter `pads` where $m$ is even and $m/2 \le 4$. `pads` denotes the amount of padding to apply before and after each of the $m/2$ last dimensions of the input tensor.
+ - Output: Padded tensor. Shape depends on `pads`.
+
+## Baseline: Using a C++ plugin
+
+To establish a baseline, we first demonstrate a C++ plugin implementing circular padding. The relevant files can be found in the `circ_plugin_cpp` folder: the included `CMakeLists.txt` can be used to build the shared library `libcirc_pad_plugin.so` / `circ_pad_plugin.dll`.
+
+```bash
+cd $TRT_OSSPATH/samples/python/trt_python_plugin
+mkdir build && pushd build
+cmake .. && make -j
+popd
+python3 circ_pad_plugin_cpp.py --plugin-lib build/libcirc_pad_plugin.so
+```
+
+## Python plugin: cuda-python
+
+The cuda-python based implementation can be found in `circ_pad_plugin_cuda_python.py`. `cuda.nvrtc` is used to JIT compile a C/C++-based kernel, which is provided as a string. The compiled kernel is then launched through cuda-python's `cuda.cuLaunchKernel`.
+
+`circ_pad_plugin_cuda_python.py` demonstrates an ONNX-based workflow: `circ_pad_plugin_inetdef_cuda_python.py` demonstrates a workflow where the model is constructed through `INetworkDefinition`.
+
+## Python plugin: CuPy
+
+The CuPy-based implementation can be found in `circ_pad_plugin_cupy.py`. CuPy's `RawKernel` class has been used to provide the C/C++-based kernel implementation as a string. CuPy will JIT compile the kernel.
+
+## Python plugin: Triton (valid only on Linux)
+
+The same plugin can be implemented with a Triton-based kernel as well. The only other change would be to `enqueue`. The entire implementation can be found in `circ_pad_plugin_triton.py`.
+
+Some remarks:
+ - Triton also allows for JIT-able kernels.
+ - CuPy device arrays cannot be passed into Triton kernels directly -- only Torch arrays are accepted. However, we can use `torch.as_tensor()` to get around this constraint.
+ - Triton does not seem to allow the specification of a CUDA stream.
+
+## Python plugin: Numba
+
+The Numba implementation can be found in `circ_pad_plugin_numba.py`. Some remarks:
+ - Numba also allows for JIT-able kernels.
+ - CuPy device arrays can be passed into Numba kernels without issue since CuPy arrays implement `__cuda_array_interface__`.
+
+## Python plugin: Torch
+
+The flexibility of the `enqueue()` interface means that it is not always necessary to implement a custom kernel. In this case, PyTorch's [torch.nn.functional.pad](https://pytorch.org/docs/stable/generated/torch.nn.functional.pad.html) offers the exact same capability we want, so we can use that inside `enqueue()`, as in `circ_pad_plugin_torch.py`.
+
+## Python plugin: Multi-tactic, Multi-plugin (based on IPluginV3)
+
+The entire implementation can be found in `circ_pad_plugin_multi_tactic.py`.
+
+### Custom tactics
+
+When multiple options are available to compute the same op, and it's not possible to reliably predict which one will be faster for the expected input shapes/types or the target platform,
+it is useful to ask TensorRT to time all available options during the build stage. In V2 plugins, TensorRT would only time different type/format combinations supported by the plugin, but
+V3 plugins allow users to specify any number of custom tactics to time also (in addition to type/format combinations).
+
+In this example, we specify two custom tactics: PyTorch's [torch.nn.functional.pad](https://pytorch.org/docs/stable/generated/torch.nn.functional.pad.html) and a custom kernel written using OpenAI Triton.
+
+It is possible to advertise tactics specific to a format combination. e.g. In this sample, we can support both tactics for FP32 I/O, and only support the OpenAI Triton tactic for FP16 I/O. To achieve this, return in `get_valid_tactics()` the set of tactics `T(f)` supported by the plugin for the format combination `f` indicated by the immediately preceding call to `configure_plugin()`. To enable this behavior in this sample, pass the flag `--per-format-tactics`.
+
+### Multiple plugins instances
+
+Imagine that you expect to have multiple instances of the same plugin in your network, which would operate on separate inputs, but where the input and output shapes/formats, as well
+as other determining plugin attributes would be the same. With V2 plugins, TensorRT would time all such plugin instances during the engine build -- however, this would be inefficient because the only salient difference between those instances are the values of the input tensors.
+
+To communicate to TensorRT that you would like the timing for similar plugin instances to be cached, V3 plugins allow for the specification of a timing cache ID. The timing cache ID
+should only capture timing determinants extraneous to plugin I/O, like their shapes and formats. Typically, this would be the values of any plugin attributes that might be different
+between the plugin instances.
+
+In this example,
+ - The shape of the `pads` parameter affects timing, but only as far as it affects the output shape. Therefore, the timing cache ID could be an empty string.
+ - We consider a scenario where there are two circular padding plugin instances with identical configurations. Therefore, only a single instance should be timed by TensorRT.
+   This can be verified by inspecting the log.
+
+# Limitations
+
+ - Plugins cannot be serialized into the engine (in contrast to `IBuilderConfig::setPluginsToSerialize()`)
+   - Plugin class and Plugin Creator class must exist in the module where the engine is deserialized
+ - The engine / ONNX model cannot be run from outside Python (e.g. with `trtexec`)
+   - This functionality is possible to implement but comes at the cost of embedding the Python interpreter to the TRT runtime / the binary loading the engine
+ - (For `IPluginV2DynamicExt` only) No bindings yet for `attachToContext()` and `detachFromContext()` which are not pure virtual.
+ - `circ_pad_plugin_torch.py` may work on aarch64 platforms but is unsupported.
+
+# FAQ
+
+1. What are the performance impacts of a Python-based plugin versus a C++ one?
+
+   In preliminary testing, the Python overhead was found to be very minimal to negligible. In fact, if the kernels were compiled AOT (instead of JIT) the CuPY and Triton
+   versions of the plugin were as performant as the C++ one. However, with Numba, there seems to be a significant kernel launch overhead.
+
+2. Can I deploy a TRT engine including a Python plugin in a runtime environment without Python?
+
+   No. There is no way to fully embed a Python plugin into the engine that allows for it to be executed without the need for Python during inference time.
+
+   This design principle is what allows for the `enqueue()` to be implemented in any framework of choice.
+
+# License
+
+For terms and conditions for use, reproduction, and distribution, see the [TensorRT Software License Agreement](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sla/index.html) documentation.
+
+# Changelog
+
+October 2025: Migrate to strongly typed APIs.
+  
+August 2025: Removed support for Python versions < 3.10.
+
+July 2023: Initial release of this sample
+
+# Known issues
+
+There are no known issues in this sample

samples/python/python_plugin/circ_pad_plugin_cpp.py

@@ -0,0 +1,97 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2025 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 argparse
+import onnx_graphsurgeon as gs
+import numpy as np
+import onnx
+import ctypes
+
+import tensorrt as trt
+from polygraphy.backend.trt import (
+    CreateConfig,
+    EngineFromNetwork,
+    NetworkFromOnnxPath,
+    TrtRunner,
+)
+
+
+def parseArgs():
+    parser = argparse.ArgumentParser(
+        description="Options for Circular Padding plugin C++ example"
+    )
+
+    parser.add_argument(
+        "--precision",
+        type=str,
+        default="fp32",
+        choices=["fp32", "fp16"],
+        help="Precision to use for plugin",
+    )
+    parser.add_argument(
+        "--plugin-lib",
+        type=str,
+        help="Path to the Circular Padding plugin lib",
+        required=True,
+    )
+
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+
+    args = parseArgs()
+
+    handle = ctypes.CDLL(args.plugin_lib)
+    if not handle:
+        raise RuntimeError("Could not load Circular Padding plugin library")
+
+    precision = np.float32 if args.precision == "fp32" else np.float16
+    inp_shape = (10, 3, 32, 32)
+    X = np.random.normal(size=inp_shape).astype(precision)
+
+    pads = (1, 1, 1, 1)
+
+    # create ONNX model
+    onnx_path = f"test_CircPadPlugin_cpp_{args.precision}.onnx"
+    inputA = gs.Variable(name="X", shape=inp_shape, dtype=precision)
+    Y = gs.Variable(name="Y", dtype=precision)
+    myPluginNode = gs.Node(
+        name="CircPadPlugin",
+        op="CircPadPlugin",
+        inputs=[inputA],
+        outputs=[Y],
+        attrs={"pads": pads},
+    )
+    graph = gs.Graph(nodes=[myPluginNode], inputs=[inputA], outputs=[Y], opset=16)
+    onnx.save(gs.export_onnx(graph), onnx_path)
+
+    # build engine
+    build_engine = EngineFromNetwork(
+        NetworkFromOnnxPath(onnx_path, strongly_typed=True), CreateConfig()
+    )
+
+    Y_ref = np.pad(X, [[0, 0], [0, 0], [pads[0], pads[1]], [pads[2], pads[3]]], "wrap")
+    # Run
+    with TrtRunner(build_engine, "trt_runner") as runner:
+        outputs = runner.infer({"X": X})
+        Y = outputs["Y"]
+
+        if np.allclose(Y, Y_ref):
+            print("Inference result correct!")
+        else:
+            print("Inference result incorrect!")

samples/python/python_plugin/circ_pad_plugin_cuda_python.py

@@ -0,0 +1,364 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2025 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 onnx_graphsurgeon as gs
+import numpy as np
+import onnx
+import sys
+import os
+
+import tensorrt as trt
+from polygraphy.backend.trt import (
+    CreateConfig,
+    EngineFromNetwork,
+    NetworkFromOnnxPath,
+    TrtRunner,
+)
+from polygraphy.json import to_json, from_json
+
+sys.path.insert(1, os.path.join(os.path.dirname(os.path.realpath(__file__)), os.pardir))
+from plugin_utils import cuda_call, KernelHelper, parseArgs, CudaCtxManager, cuda_init, cuda_get_device, cuda_memcpy_htod
+import common_runtime as common
+from cuda.bindings import driver as cuda
+
+circ_pad_half_kernel = r"""
+#include <cuda_fp16.h>
+extern "C" __global__
+void circ_pad_half(half const* X, int const* all_pads, int const* orig_dims, half* Y, int const* Y_shape, int Y_len) {
+    int index = blockIdx.x * blockDim.x + threadIdx.x;
+    int stride = blockDim.x * gridDim.x;
+
+    for(int i = index; i < Y_len; i += stride)
+    {
+        int i3 = i % Y_shape[3];
+        int i2 = (i / Y_shape[3]) % Y_shape[2];
+        int i1 = (i / Y_shape[3] / Y_shape[2]) % Y_shape[1];
+        int i0 = i / Y_shape[3] / Y_shape[2] / Y_shape[1];
+
+        int j0 = (i0 - all_pads[0] + orig_dims[0]) % orig_dims[0];
+        int j1 = (i1 - all_pads[2] + orig_dims[1]) % orig_dims[1];
+        int j2 = (i2 - all_pads[4] + orig_dims[2]) % orig_dims[2];
+        int j3 = (i3 - all_pads[6] + orig_dims[3]) % orig_dims[3];
+
+        Y[i] = X[
+            orig_dims[3] * orig_dims[2] * orig_dims[1] * j0
+            + orig_dims[3] * orig_dims[2] * j1
+            + orig_dims[3] * j2
+            + j3
+        ];
+    }
+}
+"""
+
+circ_pad_float_kernel = r"""
+extern "C" __global__
+void circ_pad_float(float const* X, int const* all_pads, int const* orig_dims, float* Y, int const* Y_shape, int Y_len) {
+    int index = blockIdx.x * blockDim.x + threadIdx.x;
+    int stride = blockDim.x * gridDim.x;
+
+    for(int i = index; i < Y_len; i += stride)
+    {
+        int i3 = i % Y_shape[3];
+        int i2 = (i / Y_shape[3]) % Y_shape[2];
+        int i1 = (i / Y_shape[3] / Y_shape[2]) % Y_shape[1];
+        int i0 = i / Y_shape[3] / Y_shape[2] / Y_shape[1];
+
+        int j0 = (i0 - all_pads[0] + orig_dims[0]) % orig_dims[0];
+        int j1 = (i1 - all_pads[2] + orig_dims[1]) % orig_dims[1];
+        int j2 = (i2 - all_pads[4] + orig_dims[2]) % orig_dims[2];
+        int j3 = (i3 - all_pads[6] + orig_dims[3]) % orig_dims[3];
+
+        Y[i] = X[
+            orig_dims[3] * orig_dims[2] * orig_dims[1] * j0
+            + orig_dims[3] * orig_dims[2] * j1
+            + orig_dims[3] * j2
+            + j3
+        ];
+    }
+}
+"""
+
+
+class CircPadPlugin(trt.IPluginV2DynamicExt):
+    def __init__(self, fc=None):
+        trt.IPluginV2DynamicExt.__init__(self)
+        self.pads = []
+        self.X_shape = []
+        self.N = 0
+
+        self.all_pads_d = None
+        self.orig_dims_d = None
+        self.Y_shape_d = None
+
+        self.num_outputs = 1
+        self.plugin_namespace = ""
+        self.plugin_type = "CircPadPlugin"
+        self.plugin_version = "1"
+
+        self.cuDevice = None
+
+        if fc is not None:
+            assert set([f.name for f in fc]) == set(
+                ["pads", "N"]
+            ), "Field collection invalid"
+            for f in fc:
+                if f.name == "pads":
+                    self.pads = f.data
+                elif f.name == "N":
+                    self.N = int(f.data)
+
+    def initialize(self):
+        self.cuDevice = cuda_get_device(0)
+        trt.get_plugin_registry().acquire_plugin_resource(
+            "cuda_ctx", CudaCtxManager(self.cuDevice)
+        )
+        self.all_pads_d = common.DeviceMem(np.int32().itemsize * self.N * 2)
+        self.orig_dims_d = common.DeviceMem(np.int32().itemsize * self.N)
+        self.Y_shape_d = common.DeviceMem(np.int32().itemsize * self.N)
+
+    def get_output_datatype(self, index, input_types):
+        return input_types[0]
+
+    def get_output_dimensions(self, output_index, inputs, exprBuilder):
+
+        output_dims = trt.DimsExprs(inputs[0])
+
+        for i in range(np.size(self.pads) // 2):
+            output_dims[len(output_dims) - i - 1] = exprBuilder.operation(
+                trt.DimensionOperation.SUM,
+                inputs[0][len(output_dims) - i - 1],
+                exprBuilder.constant(self.pads[i * 2] + self.pads[i * 2 + 1]),
+            )
+
+        return output_dims
+
+    def serialize(self):
+        return to_json({"pads": self.pads, "N": self.N})
+
+    def configure_plugin(self, inp, out):
+        X_dims = inp[0].desc.dims
+        self.X_shape = np.zeros((len(X_dims),))
+        for i in range(len(X_dims)):
+            self.X_shape[i] = X_dims[i]
+
+        all_pads = np.zeros((self.N * 2,), dtype=np.int32)
+        orig_dims = np.array(self.X_shape, dtype=np.int32)
+        out_dims = np.array(self.X_shape, dtype=np.int32)
+
+        for i in range(np.size(self.pads) // 2):
+            out_dims[self.N - i - 1] += self.pads[i * 2] + self.pads[i * 2 + 1]
+            all_pads[self.N * 2 - 2 * i - 2] = self.pads[i * 2]
+            all_pads[self.N * 2 - 2 * i - 1] = self.pads[i * 2 + 1]
+
+        # Copy vectors from host memory to device memory
+        if self.all_pads_d:
+            cuda_memcpy_htod(self.all_pads_d.device_ptr, all_pads)
+        if self.orig_dims_d:
+            cuda_memcpy_htod(self.orig_dims_d.device_ptr, orig_dims)
+        if self.Y_shape_d:
+            cuda_memcpy_htod(self.Y_shape_d.device_ptr, out_dims)
+
+        self.Y_len_d = np.prod(out_dims)
+
+    def supports_format_combination(self, pos, in_out, num_inputs):
+        assert num_inputs == 1
+        assert pos < len(in_out)
+
+        desc = in_out[pos]
+        if desc.format != trt.TensorFormat.LINEAR:
+            return False
+
+        # first input should be float16 or float32
+        if pos == 0:
+            return desc.type == trt.DataType.FLOAT or desc.type == trt.DataType.HALF
+
+        # output should have the same type as the input
+        if pos == 1:
+            return in_out[0].type == desc.type
+
+        assert False
+
+    def enqueue(self, input_desc, output_desc, inputs, outputs, workspace, stream):
+
+        inp_dtype = trt.nptype(input_desc[0].type)
+
+        blockSize = 256
+        numBlocks = int((np.prod(np.array(self.X_shape)) + blockSize - 1) // blockSize)
+
+        da = np.array([inputs[0]], dtype=np.uint64)
+        dc = np.array([outputs[0]], dtype=np.uint64)
+
+        d_all_pads = np.array([int(self.all_pads_d.device_ptr)], dtype=np.uint64)
+        d_orig_dims = np.array([int(self.orig_dims_d.device_ptr)], dtype=np.uint64)
+        d_Y_shape = np.array([int(self.Y_shape_d.device_ptr)], dtype=np.uint64)
+        Y_len = np.array(self.Y_len_d, dtype=np.uint32)
+
+        args = [da, d_all_pads, d_orig_dims, dc, d_Y_shape, Y_len]
+        kernelArgs = np.array([arg.ctypes.data for arg in args], dtype=np.uint64)
+
+        stream_ptr = np.array([stream], dtype=np.uint64)
+
+        if inp_dtype == np.float32:
+            kernelHelper = KernelHelper(circ_pad_float_kernel, int(self.cuDevice))
+            _circ_pad_float_kernel = kernelHelper.getFunction(b"circ_pad_float")
+            cuda_call(
+                cuda.cuLaunchKernel(
+                    _circ_pad_float_kernel,
+                    numBlocks,
+                    1,
+                    1,
+                    blockSize,
+                    1,
+                    1,
+                    0,
+                    stream_ptr,
+                    kernelArgs,
+                    0,
+                )
+            )
+        elif inp_dtype == np.float16:
+            kernelHelper = KernelHelper(circ_pad_half_kernel, int(self.cuDevice))
+            _circ_pad_half_kernel = kernelHelper.getFunction(b"circ_pad_half")
+            cuda_call(
+                cuda.cuLaunchKernel(
+                    _circ_pad_half_kernel,
+                    numBlocks,
+                    1,
+                    1,
+                    blockSize,
+                    1,
+                    1,
+                    0,
+                    stream_ptr,
+                    kernelArgs,
+                    0,
+                )
+            )
+        else:
+            raise ValueError("inp_dtype not valid")
+
+    def clone(self):
+        cloned_plugin = CircPadPlugin()
+        cloned_plugin.__dict__.update(self.__dict__)
+        return cloned_plugin
+
+    def terminate(self):
+        # Release DeviceMem objects - automatic cleanup via __del__ when reference count reaches 0
+        self.all_pads_d = None
+        self.orig_dims_d = None
+        self.Y_shape_d = None
+
+        trt.get_plugin_registry().release_plugin_resource("cuda_ctx")
+
+    #
+    # The following defaults take effect since the respective methods are not overriden
+    #
+
+    # def get_serialization_size(self):
+    #     return len(to_json({"pads": self.pads}))
+
+    # def get_workspace_size(self, input_desc, output_desc):
+    #     return 0
+
+    # def destroy(self):
+    #     pass
+
+
+class CircPadPluginCreator(trt.IPluginCreator):
+    def __init__(self):
+        trt.IPluginCreator.__init__(self)
+        self.name = "CircPadPlugin"
+        self.plugin_namespace = ""
+        self.plugin_version = "1"
+        self.field_names = trt.PluginFieldCollection(
+            [
+                trt.PluginField("pads", np.array([]), trt.PluginFieldType.INT32),
+                trt.PluginField("N", np.array([]), trt.PluginFieldType.INT32),
+            ]
+        )
+
+    def create_plugin(self, name, fc):
+        return CircPadPlugin(fc)
+
+    def deserialize_plugin(self, name, data):
+        deserialized = CircPadPlugin()
+        j = dict(from_json(data))
+        deserialized.__dict__.update(j)
+        return deserialized
+
+
+if __name__ == "__main__":
+
+    args = parseArgs()
+
+    # Initialize CUDA Driver API
+    cuda_init()
+
+    # Retrieve handle for device 0
+    cuDevice = cuda_get_device(0)
+
+    plg_registry = trt.get_plugin_registry()
+
+    # Create context
+    plg_registry.acquire_plugin_resource("cuda_ctx", CudaCtxManager(cuDevice))
+
+    precision = np.float32 if args.precision == "fp32" else np.float16
+
+    inp_shape = (100, 2, 32, 32)
+    X = np.random.normal(size=inp_shape).astype(precision)
+
+    pads = (1, 1, 1, 1)
+
+    # Load standard plugins
+    TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
+    trt.init_libnvinfer_plugins(TRT_LOGGER, namespace="")
+
+    # Register plugin creator
+    my_plugin_creator = CircPadPluginCreator()
+    plg_registry.register_creator(my_plugin_creator, "")
+
+    # create ONNX model
+    onnx_path = f"test_CircPadPlugin_cuda_python_{args.precision}.onnx"
+    inputA = gs.Variable(name="X", shape=inp_shape, dtype=precision)
+    Y = gs.Variable(name="Y", dtype=precision)
+    myPluginNode = gs.Node(
+        name="CircPadPlugin",
+        op="CircPadPlugin",
+        inputs=[inputA],
+        outputs=[Y],
+        attrs={"pads": pads, "N": 4},
+    )
+    graph = gs.Graph(nodes=[myPluginNode], inputs=[inputA], outputs=[Y], opset=16)
+    onnx.save(gs.export_onnx(graph), onnx_path)
+
+    # build engine
+    build_engine = EngineFromNetwork(
+        NetworkFromOnnxPath(onnx_path, strongly_typed=True), CreateConfig()
+    )
+
+    Y_ref = np.pad(X, [[0, 0], [0, 0], [pads[0], pads[1]], [pads[2], pads[3]]], "wrap")
+    # Run
+    with TrtRunner(build_engine, "trt_runner") as runner:
+        outputs = runner.infer({"X": X})
+        Y = outputs["Y"]
+
+        if np.allclose(Y, Y_ref):
+            print("Inference result correct!")
+        else:
+            print("Inference result incorrect!")
+
+    plg_registry.release_plugin_resource("cuda_ctx")

samples/python/python_plugin/circ_pad_plugin_cupy.py

@@ -0,0 +1,324 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2025 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 onnx_graphsurgeon as gs
+import numpy as np
+import onnx
+import cupy as cp
+import time
+import pickle
+import sys
+import os
+
+import tensorrt as trt
+from polygraphy.backend.trt import (
+    CreateConfig,
+    EngineFromNetwork,
+    NetworkFromOnnxPath,
+    TrtRunner,
+)
+
+from polygraphy.json import to_json, from_json
+
+sys.path.insert(1, os.path.join(os.path.dirname(os.path.realpath(__file__)), os.pardir))
+from plugin_utils import volume, parseArgs
+
+circ_pad_half_kernel = cp.RawKernel(
+    r"""
+#include <cuda_fp16.h>
+extern "C" __global__
+void circ_pad_half(half const* X, int const* all_pads, int const* orig_dims, half* Y, int const* Y_shape, int const* Y_len) {
+    int index = blockIdx.x * blockDim.x + threadIdx.x;
+    int stride = blockDim.x * gridDim.x;
+
+    for(int i = index; i < *Y_len; i += stride)
+    {
+        int i3 = i % Y_shape[3];
+        int i2 = (i / Y_shape[3]) % Y_shape[2];
+        int i1 = (i / Y_shape[3] / Y_shape[2]) % Y_shape[1];
+        int i0 = i / Y_shape[3] / Y_shape[2] / Y_shape[1];
+
+        int j0 = (i0 - all_pads[0] + orig_dims[0]) % orig_dims[0];
+        int j1 = (i1 - all_pads[2] + orig_dims[1]) % orig_dims[1];
+        int j2 = (i2 - all_pads[4] + orig_dims[2]) % orig_dims[2];
+        int j3 = (i3 - all_pads[6] + orig_dims[3]) % orig_dims[3];
+
+        Y[i] = X[
+            orig_dims[3] * orig_dims[2] * orig_dims[1] * j0
+            + orig_dims[3] * orig_dims[2] * j1
+            + orig_dims[3] * j2
+            + j3
+        ];
+    }
+}
+""",
+    "circ_pad_half",
+)
+
+circ_pad_float_kernel = cp.RawKernel(
+    r"""
+extern "C" __global__
+void circ_pad_float(float const* X, int const* all_pads, int const* orig_dims, float* Y, int const* Y_shape, int const* Y_len) {
+    int index = blockIdx.x * blockDim.x + threadIdx.x;
+    int stride = blockDim.x * gridDim.x;
+
+    for(int i = index; i < *Y_len; i += stride)
+    {
+        int i3 = i % Y_shape[3];
+        int i2 = (i / Y_shape[3]) % Y_shape[2];
+        int i1 = (i / Y_shape[3] / Y_shape[2]) % Y_shape[1];
+        int i0 = i / Y_shape[3] / Y_shape[2] / Y_shape[1];
+
+        int j0 = (i0 - all_pads[0] + orig_dims[0]) % orig_dims[0];
+        int j1 = (i1 - all_pads[2] + orig_dims[1]) % orig_dims[1];
+        int j2 = (i2 - all_pads[4] + orig_dims[2]) % orig_dims[2];
+        int j3 = (i3 - all_pads[6] + orig_dims[3]) % orig_dims[3];
+
+        Y[i] = X[
+            orig_dims[3] * orig_dims[2] * orig_dims[1] * j0
+            + orig_dims[3] * orig_dims[2] * j1
+            + orig_dims[3] * j2
+            + j3
+        ];
+    }
+}
+""",
+    "circ_pad_float",
+)
+
+
+class CircPadPlugin(trt.IPluginV2DynamicExt):
+    def __init__(self, fc=None):
+        trt.IPluginV2DynamicExt.__init__(self)
+        self.pads = []
+        self.X_shape = []
+
+        self.num_outputs = 1
+        self.plugin_namespace = ""
+        self.plugin_type = "CircPadPlugin"
+        self.plugin_version = "1"
+
+        if fc is not None:
+            assert fc[0].name == "pads"
+            self.pads = fc[0].data
+
+    def get_output_datatype(self, index, input_types):
+        return input_types[0]
+
+    def get_output_dimensions(self, output_index, inputs, exprBuilder):
+
+        output_dims = trt.DimsExprs(inputs[0])
+
+        for i in range(np.size(self.pads) // 2):
+            output_dims[len(output_dims) - i - 1] = exprBuilder.operation(
+                trt.DimensionOperation.SUM,
+                inputs[0][len(output_dims) - i - 1],
+                exprBuilder.constant(self.pads[i * 2] + self.pads[i * 2 + 1]),
+            )
+
+        return output_dims
+
+    def serialize(self):
+        return to_json({"pads": self.pads})
+
+    def configure_plugin(self, inp, out):
+        X_dims = inp[0].desc.dims
+        self.X_shape = np.zeros((len(X_dims),))
+        for i in range(len(X_dims)):
+            self.X_shape[i] = X_dims[i]
+
+        N = len(self.X_shape)
+        all_pads = np.zeros((N * 2,))
+        orig_dims = np.array(self.X_shape)
+        out_dims = np.array(self.X_shape)
+
+        for i in range(np.size(pads) // 2):
+            out_dims[N - i - 1] += self.pads[i * 2] + self.pads[i * 2 + 1]
+            all_pads[N * 2 - 2 * i - 2] = self.pads[i * 2]
+            all_pads[N * 2 - 2 * i - 1] = self.pads[i * 2 + 1]
+
+        self.all_pads_d = cp.asarray(all_pads, dtype=cp.int32)
+        self.orig_dims_d = cp.asarray(orig_dims, dtype=cp.int32)
+        self.Y_shape_d = cp.asarray(out_dims, dtype=cp.int32)
+        self.Y_len_d = cp.array([np.prod(out_dims)], dtype=cp.int32)
+
+    def supports_format_combination(self, pos, in_out, num_inputs):
+        assert num_inputs == 1
+        assert pos < len(in_out)
+
+        desc = in_out[pos]
+        if desc.format != trt.TensorFormat.LINEAR:
+            return False
+
+        # first input should be float16 or float32
+        if pos == 0:
+            return desc.type == trt.DataType.FLOAT or desc.type == trt.DataType.HALF
+
+        # output should have the same type as the input
+        if pos == 1:
+            return in_out[0].type == desc.type
+
+        assert False
+
+    def enqueue(self, input_desc, output_desc, inputs, outputs, workspace, stream):
+
+        inp_dtype = trt.nptype(input_desc[0].type)
+
+        a_mem = cp.cuda.UnownedMemory(
+            inputs[0], volume(input_desc[0].dims) * cp.dtype(inp_dtype).itemsize, self
+        )
+        c_mem = cp.cuda.UnownedMemory(
+            outputs[0],
+            volume(output_desc[0].dims) * cp.dtype(inp_dtype).itemsize,
+            self,
+        )
+
+        a_ptr = cp.cuda.MemoryPointer(a_mem, 0)
+        c_ptr = cp.cuda.MemoryPointer(c_mem, 0)
+
+        a = cp.ndarray((volume(input_desc[0].dims)), dtype=inp_dtype, memptr=a_ptr)
+        c = cp.ndarray((volume(output_desc[0].dims)), dtype=inp_dtype, memptr=c_ptr)
+
+        cuda_stream = cp.cuda.ExternalStream(stream)
+
+        blockSize = 256
+        numBlocks = int((np.prod(np.array(self.X_shape)) + blockSize - 1) // blockSize)
+
+        with cuda_stream:
+            if inp_dtype == np.float32:
+                circ_pad_float_kernel(
+                    (numBlocks,),
+                    (blockSize,),
+                    (
+                        a,
+                        self.all_pads_d,
+                        self.orig_dims_d,
+                        c,
+                        self.Y_shape_d,
+                        self.Y_len_d,
+                    ),
+                )
+            elif inp_dtype == np.float16:
+                circ_pad_half_kernel(
+                    (numBlocks,),
+                    (blockSize,),
+                    (
+                        a,
+                        self.all_pads_d,
+                        self.orig_dims_d,
+                        c,
+                        self.Y_shape_d,
+                        self.Y_len_d,
+                    ),
+                )
+            else:
+                raise ValueError("inp_dtype not valid")
+
+    def clone(self):
+        cloned_plugin = CircPadPlugin()
+        cloned_plugin.__dict__.update(self.__dict__)
+        return cloned_plugin
+
+    #
+    # The following defaults take effect since the respective methods are not overriden
+    #
+
+    # def initialize(self):
+    #     pass
+
+    # def get_serialization_size(self):
+    #     return len(to_json({"pads": self.pads}))
+
+    # def get_workspace_size(self, input_desc, output_desc):
+    #     return 0
+
+    # def destroy(self):
+    #     pass
+
+    # def terminate(self):
+    #     pass
+
+
+class CircPadPluginCreator(trt.IPluginCreator):
+    def __init__(self):
+        trt.IPluginCreator.__init__(self)
+
+        self.name = "CircPadPlugin"
+        self.plugin_namespace = ""
+        self.plugin_version = "1"
+        self.field_names = trt.PluginFieldCollection(
+            [trt.PluginField("pads", np.array([]), trt.PluginFieldType.INT32)]
+        )
+
+    def create_plugin(self, name, fc):
+        return CircPadPlugin(fc)
+
+    def deserialize_plugin(self, name, data):
+        j = dict(from_json(data.decode("utf-8")))
+        deserialized = CircPadPlugin()
+        deserialized.__dict__.update(j)
+        return deserialized
+
+
+if __name__ == "__main__":
+
+    args = parseArgs()
+    precision = np.float32 if args.precision == "fp32" else np.float16
+
+    inp_shape = (100, 2, 32, 32)
+    X = np.random.normal(size=inp_shape).astype(precision)
+
+    pads = (1, 1, 1, 1)
+
+    # Load standard plugins
+    TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
+    trt.init_libnvinfer_plugins(TRT_LOGGER, namespace="")
+
+    # Register plugin creator
+    plg_registry = trt.get_plugin_registry()
+    my_plugin_creator = CircPadPluginCreator()
+    plg_registry.register_creator(my_plugin_creator, "")
+
+    # create ONNX model
+    onnx_path = f"test_CircPadPlugin_cupy_{args.precision}.onnx"
+    inputA = gs.Variable(name="X", shape=inp_shape, dtype=precision)
+    Y = gs.Variable(name="Y", dtype=precision)
+    myPluginNode = gs.Node(
+        name="CircPadPlugin",
+        op="CircPadPlugin",
+        inputs=[inputA],
+        outputs=[Y],
+        attrs={"pads": pads},
+    )
+    graph = gs.Graph(nodes=[myPluginNode], inputs=[inputA], outputs=[Y], opset=16)
+    onnx.save(gs.export_onnx(graph), onnx_path)
+
+    # build engine
+    build_engine = EngineFromNetwork(
+        NetworkFromOnnxPath(onnx_path, strongly_typed=True), CreateConfig()
+    )
+
+    Y_ref = np.pad(X, [[0, 0], [0, 0], [pads[0], pads[1]], [pads[2], pads[3]]], "wrap")
+    # Run
+    with TrtRunner(build_engine, "trt_runner") as runner:
+        outputs = runner.infer({"X": X})
+        Y = outputs["Y"]
+
+        if np.allclose(Y, Y_ref):
+            print("Inference result correct!")
+        else:
+            print("Inference result incorrect!")

samples/python/python_plugin/circ_pad_plugin_inetdef_cuda_python.py

@@ -0,0 +1,383 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2026 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 onnx_graphsurgeon as gs
+import numpy as np
+import sys
+import os
+
+import tensorrt as trt
+from polygraphy.backend.trt import (
+    CreateConfig,
+    TrtRunner,
+    create_network,
+    engine_from_network,
+)
+
+from polygraphy.json import to_json, from_json
+
+sys.path.insert(1, os.path.join(os.path.dirname(os.path.realpath(__file__)), os.pardir))
+from plugin_utils import cuda_call, KernelHelper, parseArgs, CudaCtxManager
+from cuda.bindings import driver as cuda
+
+circ_pad_half_kernel = r"""
+#include <cuda_fp16.h>
+extern "C" __global__
+void circ_pad_half(half const* X, int const* all_pads, int const* orig_dims, half* Y, int const* Y_shape, int Y_len) {
+    int index = blockIdx.x * blockDim.x + threadIdx.x;
+    int stride = blockDim.x * gridDim.x;
+
+    for(int i = index; i < Y_len; i += stride)
+    {
+        int i3 = i % Y_shape[3];
+        int i2 = (i / Y_shape[3]) % Y_shape[2];
+        int i1 = (i / Y_shape[3] / Y_shape[2]) % Y_shape[1];
+        int i0 = i / Y_shape[3] / Y_shape[2] / Y_shape[1];
+
+        int j0 = (i0 - all_pads[0] + orig_dims[0]) % orig_dims[0];
+        int j1 = (i1 - all_pads[2] + orig_dims[1]) % orig_dims[1];
+        int j2 = (i2 - all_pads[4] + orig_dims[2]) % orig_dims[2];
+        int j3 = (i3 - all_pads[6] + orig_dims[3]) % orig_dims[3];
+
+        Y[i] = X[
+            orig_dims[3] * orig_dims[2] * orig_dims[1] * j0
+            + orig_dims[3] * orig_dims[2] * j1
+            + orig_dims[3] * j2
+            + j3
+        ];
+    }
+}
+"""
+
+circ_pad_float_kernel = r"""
+extern "C" __global__
+void circ_pad_float(float const* X, int const* all_pads, int const* orig_dims, float* Y, int const* Y_shape, int Y_len) {
+    int index = blockIdx.x * blockDim.x + threadIdx.x;
+    int stride = blockDim.x * gridDim.x;
+
+    for(int i = index; i < Y_len; i += stride)
+    {
+        int i3 = i % Y_shape[3];
+        int i2 = (i / Y_shape[3]) % Y_shape[2];
+        int i1 = (i / Y_shape[3] / Y_shape[2]) % Y_shape[1];
+        int i0 = i / Y_shape[3] / Y_shape[2] / Y_shape[1];
+
+        int j0 = (i0 - all_pads[0] + orig_dims[0]) % orig_dims[0];
+        int j1 = (i1 - all_pads[2] + orig_dims[1]) % orig_dims[1];
+        int j2 = (i2 - all_pads[4] + orig_dims[2]) % orig_dims[2];
+        int j3 = (i3 - all_pads[6] + orig_dims[3]) % orig_dims[3];
+
+        Y[i] = X[
+            orig_dims[3] * orig_dims[2] * orig_dims[1] * j0
+            + orig_dims[3] * orig_dims[2] * j1
+            + orig_dims[3] * j2
+            + j3
+        ];
+    }
+}
+"""
+
+
+class CircPadPlugin(trt.IPluginV2DynamicExt):
+    def __init__(self, fc=None):
+        trt.IPluginV2DynamicExt.__init__(self)
+        self.pads = []
+        self.X_shape = []
+        self.N = 0
+
+        self.all_pads_d = None
+        self.orig_dims_d = None
+        self.Y_shape_d = None
+
+        self.num_outputs = 1
+        self.plugin_namespace = ""
+        self.plugin_type = "CircPadPlugin"
+        self.plugin_version = "1"
+
+        self.cuDevice = None
+
+        if fc is not None:
+            assert set([f.name for f in fc]) == set(
+                ["pads", "N"]
+            ), "Field collection invalid"
+            for f in fc:
+                if f.name == "pads":
+                    self.pads = f.data
+                elif f.name == "N":
+                    self.N = int(f.data)
+
+    def initialize(self):
+        self.cuDevice = cuda_call(cuda.cuDeviceGet(0))
+        trt.get_plugin_registry().acquire_plugin_resource(
+            "cuda_ctx", CudaCtxManager(self.cuDevice)
+        )
+        self.all_pads_d = cuda_call(
+            cuda.cuMemAlloc(np.int32().itemsize * self.N * 2)
+        )
+        self.orig_dims_d = cuda_call(
+            cuda.cuMemAlloc(np.int32().itemsize * self.N)
+        )
+        self.Y_shape_d = cuda_call(cuda.cuMemAlloc(np.int32().itemsize * self.N))
+
+    def get_output_datatype(self, index, input_types):
+        return input_types[0]
+
+    def get_output_dimensions(self, output_index, inputs, exprBuilder):
+
+        output_dims = trt.DimsExprs(inputs[0])
+
+        for i in range(np.size(self.pads) // 2):
+            output_dims[len(output_dims) - i - 1] = exprBuilder.operation(
+                trt.DimensionOperation.SUM,
+                inputs[0][len(output_dims) - i - 1],
+                exprBuilder.constant(self.pads[i * 2] + self.pads[i * 2 + 1]),
+            )
+
+        return output_dims
+
+    def serialize(self):
+        return to_json({"pads": self.pads, "N": self.N})
+
+    def configure_plugin(self, inp, out):
+        X_dims = inp[0].desc.dims
+        self.X_shape = np.zeros((len(X_dims),))
+        for i in range(len(X_dims)):
+            self.X_shape[i] = X_dims[i]
+
+        all_pads = np.zeros((self.N * 2,), dtype=np.int32)
+        orig_dims = np.array(self.X_shape, dtype=np.int32)
+        out_dims = np.array(self.X_shape, dtype=np.int32)
+
+        for i in range(np.size(self.pads) // 2):
+            out_dims[self.N - i - 1] += self.pads[i * 2] + self.pads[i * 2 + 1]
+            all_pads[self.N * 2 - 2 * i - 2] = self.pads[i * 2]
+            all_pads[self.N * 2 - 2 * i - 1] = self.pads[i * 2 + 1]
+
+        # Copy vectors from host memory to device memory
+        if self.all_pads_d:
+            cuda_call(
+                cuda.cuMemcpyHtoD(self.all_pads_d, all_pads, all_pads.nbytes)
+            )
+        if self.orig_dims_d:
+            cuda_call(
+                cuda.cuMemcpyHtoD(self.orig_dims_d, orig_dims, orig_dims.nbytes)
+            )
+        if self.Y_shape_d:
+            cuda_call(
+                cuda.cuMemcpyHtoD(self.Y_shape_d, out_dims, out_dims.nbytes)
+            )
+
+        self.Y_len_d = np.prod(out_dims)
+
+    def supports_format_combination(self, pos, in_out, num_inputs):
+        assert num_inputs == 1
+        assert pos < len(in_out)
+
+        desc = in_out[pos]
+        if desc.format != trt.TensorFormat.LINEAR:
+            return False
+
+        # first input should be float16 or float32
+        if pos == 0:
+            return desc.type == trt.DataType.FLOAT or desc.type == trt.DataType.HALF
+
+        # output should have the same type as the input
+        if pos == 1:
+            return in_out[0].type == desc.type
+
+        assert False
+
+    def enqueue(self, input_desc, output_desc, inputs, outputs, workspace, stream):
+
+        inp_dtype = trt.nptype(input_desc[0].type)
+
+        blockSize = 256
+        numBlocks = int((np.prod(np.array(self.X_shape)) + blockSize - 1) // blockSize)
+
+        da = np.array([inputs[0]], dtype=np.uint64)
+        dc = np.array([outputs[0]], dtype=np.uint64)
+
+        d_all_pads = np.array([int(self.all_pads_d)], dtype=np.uint64)
+        d_orig_dims = np.array([int(self.orig_dims_d)], dtype=np.uint64)
+        d_Y_shape = np.array([int(self.Y_shape_d)], dtype=np.uint64)
+        Y_len = np.array(self.Y_len_d, dtype=np.uint32)
+
+        args = [da, d_all_pads, d_orig_dims, dc, d_Y_shape, Y_len]
+        kernelArgs = np.array([arg.ctypes.data for arg in args], dtype=np.uint64)
+
+        stream_ptr = np.array([stream], dtype=np.uint64)
+
+        if inp_dtype == np.float32:
+            kernelHelper = KernelHelper(circ_pad_float_kernel, int(self.cuDevice))
+            _circ_pad_float_kernel = kernelHelper.getFunction(b"circ_pad_float")
+            cuda_call(
+                cuda.cuLaunchKernel(
+                    _circ_pad_float_kernel,
+                    numBlocks,
+                    1,
+                    1,
+                    blockSize,
+                    1,
+                    1,
+                    0,
+                    stream_ptr,
+                    kernelArgs,
+                    0,
+                )
+            )
+        elif inp_dtype == np.float16:
+            kernelHelper = KernelHelper(circ_pad_half_kernel, int(self.cuDevice))
+            _circ_pad_half_kernel = kernelHelper.getFunction(b"circ_pad_half")
+            cuda_call(
+                cuda.cuLaunchKernel(
+                    _circ_pad_half_kernel,
+                    numBlocks,
+                    1,
+                    1,
+                    blockSize,
+                    1,
+                    1,
+                    0,
+                    stream_ptr,
+                    kernelArgs,
+                    0,
+                )
+            )
+        else:
+            raise ValueError("inp_dtype not valid")
+
+    def clone(self):
+        cloned_plugin = CircPadPlugin()
+        cloned_plugin.__dict__.update(self.__dict__)
+        return cloned_plugin
+
+    def terminate(self):
+        if self.all_pads_d:
+            cuda_call(cuda.cuMemFree(self.all_pads_d))
+        if self.orig_dims_d:
+            cuda_call(cuda.cuMemFree(self.orig_dims_d))
+        if self.Y_shape_d:
+            cuda_call(cuda.cuMemFree(self.Y_shape_d))
+
+        plg_registry.release_plugin_resource("cuda_ctx")
+
+    #
+    # The following defaults take effect since the respective methods are not overriden
+    #
+
+    # def get_serialization_size(self):
+    #     return len(to_json({"pads": self.pads}))
+
+    # def get_workspace_size(self, input_desc, output_desc):
+    #     return 0
+
+    # def destroy(self):
+    #     pass
+
+
+class CircPadPluginCreator(trt.IPluginCreator):
+    def __init__(self):
+        trt.IPluginCreator.__init__(self)
+        self.name = "CircPadPlugin"
+        self.plugin_namespace = ""
+        self.plugin_version = "1"
+        self.field_names = trt.PluginFieldCollection(
+            [
+                trt.PluginField("pads", np.array([]), trt.PluginFieldType.INT32),
+                trt.PluginField("N", np.array([]), trt.PluginFieldType.INT32),
+            ]
+        )
+
+    def create_plugin(self, name, fc):
+        return CircPadPlugin(fc)
+
+    def deserialize_plugin(self, name, data):
+        deserialized = CircPadPlugin()
+        j = dict(from_json(data))
+        deserialized.__dict__.update(j)
+        return deserialized
+
+
+if __name__ == "__main__":
+
+    args = parseArgs()
+    precision = np.float32 if args.precision == "fp32" else np.float16
+
+    # Initialize CUDA Driver API
+    cuda_call(cuda.cuInit(0))
+
+    # Retrieve handle for device 0
+    cuDevice = cuda_call(cuda.cuDeviceGet(0))
+
+    plg_registry = trt.get_plugin_registry()
+
+    # Create context
+    plg_registry.acquire_plugin_resource("cuda_ctx", CudaCtxManager(cuDevice))
+
+    inp_shape = (100, 2, 32, 32)
+    X = np.random.normal(size=inp_shape).astype(precision)
+
+    pads = (1, 1, 1, 1)
+
+    TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
+    # Load standard plugins (if needed)
+    trt.init_libnvinfer_plugins(TRT_LOGGER, namespace="")
+
+    # Register plugin creator
+    my_plugin_creator = CircPadPluginCreator()
+    plg_registry.register_creator(my_plugin_creator, "")
+
+    # Create plugin object
+    builder, network = create_network(strongly_typed=True)
+    plg_creator = plg_registry.get_creator("CircPadPlugin", "1", "")
+    plugin_fields_list = [
+        trt.PluginField(
+            "pads", np.array(pads, dtype=np.int32), trt.PluginFieldType.INT32
+        ),
+        trt.PluginField("N", np.array([4], dtype=np.int32), trt.PluginFieldType.INT32),
+    ]
+    pfc = trt.PluginFieldCollection(plugin_fields_list)
+    plugin = plg_creator.create_plugin("CircPadPlugin", pfc)
+
+    # Populate network
+    input_X = network.add_input(
+        name="X",
+        dtype=trt.float32 if precision == np.float32 else trt.float16,
+        shape=X.shape,
+    )
+    out = network.add_plugin_v2([input_X], plugin)
+    out.get_output(0).name = "Y"
+    network.mark_output(tensor=out.get_output(0))
+
+    # Build engine
+    config = builder.create_builder_config()
+    engine = engine_from_network(
+        (builder, network), CreateConfig()
+    )
+
+    Y_ref = np.pad(X, [[0, 0], [0, 0], [pads[0], pads[1]], [pads[2], pads[3]]], "wrap")
+    # Run
+    with TrtRunner(engine, "trt_runner") as runner:
+        outputs = runner.infer({"X": X})
+        Y = outputs["Y"]
+
+        if np.allclose(Y, Y_ref):
+            print("Inference result correct!")
+        else:
+            print("Inference result incorrect!")
+
+    plg_registry.release_plugin_resource("cuda_ctx")

samples/python/python_plugin/circ_pad_plugin_multi_tactic.py

@@ -0,0 +1,376 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2025 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 onnx_graphsurgeon as gs
+import numpy as np
+import onnx
+import cupy as cp
+import logging
+import sys
+import os
+
+import tensorrt as trt
+from polygraphy.backend.trt import (
+    CreateConfig,
+    EngineFromNetwork,
+    NetworkFromOnnxPath,
+    TrtRunner,
+)
+
+import triton
+import triton.language as tl
+
+from enum import IntEnum
+
+from polygraphy.json import to_json, from_json
+import torch
+
+sys.path.insert(1, os.path.join(os.path.dirname(os.path.realpath(__file__)), os.pardir))
+from plugin_utils import volume, parseArgs
+
+import argparse
+
+logger = logging.getLogger("CircPadMultiTactic")
+
+class Tactic(IntEnum):
+    TORCH = 1
+    TRITON = 2
+
+@triton.jit
+def circ_pad(X,
+            all_pads_0, all_pads_2, all_pads_4, all_pads_6,
+            orig_dims_0, orig_dims_1, orig_dims_2, orig_dims_3,
+            Y,
+            Y_shape_1, Y_shape_2, Y_shape_3,
+            X_len, Y_len, BLOCK_SIZE: tl.constexpr,):
+    pid = tl.program_id(0)
+    i = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+
+    mask_y = i < Y_len
+
+    i3 = i % Y_shape_3
+    i2 = (i // Y_shape_3) % Y_shape_2
+    i1 = (i // Y_shape_3 // Y_shape_2) % Y_shape_1
+    i0 = i // Y_shape_3 // Y_shape_2 // Y_shape_1
+
+    j0 = (i0 - all_pads_0 + orig_dims_0) % orig_dims_0
+    j1 = (i1 - all_pads_2 + orig_dims_1) % orig_dims_1
+    j2 = (i2 - all_pads_4 + orig_dims_2) % orig_dims_2
+    j3 = (i3 - all_pads_6 + orig_dims_3) % orig_dims_3
+
+    load_idx = orig_dims_3 * orig_dims_2 * orig_dims_1 * j0 + orig_dims_3 * orig_dims_2 * j1 + orig_dims_3 * j2 + j3
+    mask_x = load_idx < X_len
+
+    x = tl.load(X + load_idx, mask=mask_x)
+
+    tl.store(Y + i, x, mask=mask_y)
+
+class CircPadPlugin(trt.IPluginV3, trt.IPluginV3OneCore, trt.IPluginV3OneBuild, trt.IPluginV3OneRuntime):
+    def __init__(self, fc=None, phase=None):
+        trt.IPluginV3.__init__(self)
+        trt.IPluginV3OneCore.__init__(self)
+        trt.IPluginV3OneBuild.__init__(self)
+        trt.IPluginV3OneRuntime.__init__(self)
+        self.pads = []
+        self.X_shape = []
+
+        self.per_format_tactics = (
+            False  # whether per-format tactics or global tactics should be used
+        )
+        self.curr_type = None  # format being timed currently by TRT auto-tuner
+
+        self.num_outputs = 1
+        self.plugin_namespace = ""
+        self.plugin_name = "CircPadPlugin"
+        self.plugin_version = "1"
+
+        # Set the timing cache ID to prevent unnecessary timing of second plugin instance
+        self.timing_cache_id = ""
+        
+        self.tactic = None
+
+        if fc is not None:
+            for f in fc:
+                if f.name == "pads":
+                    self.pads = f.data
+                elif f.name == "per_format_tactics":
+                    self.per_format_tactics = int(f.data)
+
+        if phase is not None:
+            self.phase = phase
+
+    def get_capability_interface(self, type):
+        return self
+
+    def get_output_data_types(self, input_types):
+        return [input_types[0]]
+
+    def get_output_shapes(self, inputs, shape_inputs, exprBuilder):
+        output_dims = trt.DimsExprs(inputs[0])
+
+        for i in range(np.size(self.pads) // 2):
+            output_dims[len(output_dims) - i - 1] = exprBuilder.operation(
+                trt.DimensionOperation.SUM,
+                inputs[0][len(output_dims) - i - 1],
+                exprBuilder.constant(self.pads[i * 2] + self.pads[i * 2 + 1]),
+            )
+
+        return [output_dims]
+
+    def get_fields_to_serialize(self):
+        return trt.PluginFieldCollection([
+            trt.PluginField("pads", self.pads, trt.PluginFieldType.INT32),
+            trt.PluginField(
+                "per_format_tactics",
+                np.array([self.per_format_tactics], dtype=np.int32),
+                trt.PluginFieldType.INT32,
+            ),
+        ])
+
+    def configure_plugin(self, inp, out):
+        assert inp[0].desc.type == trt.float32 or inp[0].desc.type == trt.float16
+        self.curr_type = inp[0].desc.type
+
+    def on_shape_change(self, inp, out):
+        if (
+            self.phase == trt.TensorRTPhase.RUNTIME
+            and self.per_format_tactics
+            and inp[0].type == trt.float16
+        ):
+            assert self.tactic == Tactic.TRITON
+
+        X_dims = inp[0].dims
+        self.X_shape = np.zeros((len(X_dims),))
+        for i in range(len(X_dims)):
+            self.X_shape[i] = X_dims[i]
+
+    def supports_format_combination(self, pos, in_out, num_inputs):
+        assert num_inputs == 1
+        assert pos < len(in_out)
+
+        desc = in_out[pos].desc
+        if desc.format != trt.TensorFormat.LINEAR:
+            return False
+
+        # first input should be float16 or float32
+        if pos == 0:
+            return desc.type == trt.DataType.FLOAT or desc.type == trt.DataType.HALF
+
+        # output should have the same type as the input
+        if pos == 1:
+            return in_out[0].desc.type == desc.type
+
+        assert False
+
+    def enqueue(self, input_desc, output_desc, inputs, outputs, workspace, stream):
+
+        inp_dtype = trt.nptype(input_desc[0].type)
+
+        a_mem = cp.cuda.UnownedMemory(
+            inputs[0], volume(input_desc[0].dims) * cp.dtype(inp_dtype).itemsize, self
+        )
+        c_mem = cp.cuda.UnownedMemory(
+            outputs[0],
+            volume(output_desc[0].dims) * cp.dtype(inp_dtype).itemsize,
+            self,
+        )
+
+        a_ptr = cp.cuda.MemoryPointer(a_mem, 0)
+        c_ptr = cp.cuda.MemoryPointer(c_mem, 0)
+
+        c_d = cp.ndarray((volume(output_desc[0].dims)), dtype=inp_dtype, memptr=c_ptr)
+        
+        if self.phase == trt.TensorRTPhase.BUILD:
+            logger.info(f"Timing tactic: {self.tactic}")
+
+        if self.tactic == Tactic.TORCH:
+            # Use PyTorch functional op - no need to write kernel
+            a_d = cp.ndarray(tuple(input_desc[0].dims), dtype=inp_dtype, memptr=a_ptr)
+            a_t = torch.as_tensor(a_d, device='cuda')
+            out = torch.nn.functional.pad(a_t, self.pads.tolist(), mode='circular')
+            cp.copyto(c_d, cp.reshape(cp.asarray(out), (-1,)))
+        elif self.tactic == Tactic.TRITON:
+            a_d = cp.ndarray((volume(input_desc[0].dims)), dtype=inp_dtype, memptr=a_ptr)
+            a_t = torch.as_tensor(a_d, device='cuda')
+            c_t = torch.as_tensor(c_d, device='cuda')
+
+            N = len(self.X_shape)
+            all_pads = np.zeros((N * 2,), dtype=np.int32)
+            orig_dims = np.array(self.X_shape, dtype=np.int32)
+            out_dims = np.array(self.X_shape, dtype=np.int32)
+
+            for i in range(np.size(pads) // 2):
+                out_dims[N - i - 1] += pads[i * 2] + pads[i * 2 + 1]
+                all_pads[N * 2 - 2 * i - 2] = pads[i * 2]
+                all_pads[N * 2 - 2 * i - 1] = pads[i * 2 + 1]
+
+            all_pads = all_pads.tolist()
+            orig_dims = orig_dims.tolist()
+            out_dims = out_dims.tolist()
+
+            blockSize = 256
+            numBlocks = tuple([int((np.prod(out_dims) + blockSize - 1) // blockSize)])
+
+            circ_pad[numBlocks](a_t,
+                all_pads[0], all_pads[2], all_pads[4], all_pads[6],
+                orig_dims[0], orig_dims[1], orig_dims[2], orig_dims[3],
+                c_t,
+                out_dims[1], out_dims[2], out_dims[3],
+                int(np.prod(orig_dims)), int(np.prod(out_dims)), BLOCK_SIZE=256
+            )
+        else:
+            raise RuntimeError("Invalid tactic")
+    
+    def attach_to_context(self, context):
+        return self.clone()
+    
+    def get_valid_tactics(self):
+        assert self.curr_type is not None
+        if self.per_format_tactics and self.curr_type == trt.float16:
+            return [int(Tactic.TRITON)]
+
+        return [int(Tactic.TORCH), int(Tactic.TRITON)]
+
+    def set_tactic(self, tactic):
+        self.tactic = Tactic(tactic)
+
+        if self.phase == trt.TensorRTPhase.RUNTIME:
+            logger.info(f"Best tactic chosen: {self.tactic}")
+
+    def clone(self):
+        cloned_plugin = CircPadPlugin()
+        cloned_plugin.__dict__.update(self.__dict__)
+        return cloned_plugin
+
+    # 
+    # The following defaults take effect since the respective methods are not overriden
+    #
+
+    # def get_workspace_size(self, input_desc, output_desc):
+    #     return 0
+    
+    # def destroy(self):
+    #     pass
+
+
+class CircPadPluginCreator(trt.IPluginCreatorV3One):
+    def __init__(self):
+        trt.IPluginCreatorV3One.__init__(self)
+        self.name = "CircPadPlugin"
+        self.plugin_namespace = ""
+        self.plugin_version = "1"
+        self.field_names = trt.PluginFieldCollection([
+            trt.PluginField("pads", np.array([]), trt.PluginFieldType.INT32),
+            trt.PluginField(
+                "per_format_tactics", np.array([]), trt.PluginFieldType.INT32
+            ),
+        ])
+
+    def create_plugin(self, name, fc, phase):
+        return CircPadPlugin(fc, phase)
+
+
+if __name__ == "__main__":
+    logging.basicConfig()
+    logger.setLevel(logging.INFO)
+
+    parser = argparse.ArgumentParser(
+        description="Options for Circular Padding plugin multi-tactic sample"
+    )
+
+    parser.add_argument(
+        "--precision",
+        type=str,
+        default="fp32",
+        choices=["fp32", "fp16"],
+        help="Precision to use for plugin",
+    )
+    parser.add_argument(
+        "--per-format-tactics",
+        action="store_true",
+        help="Whether per-format tactics or global tactics should be used",
+    )
+
+    args = parser.parse_args()
+
+    precision = np.float32 if args.precision == "fp32" else np.float16
+    is_tactics_per_format = 1 if args.per_format_tactics else 0
+
+    inp_shape = (10, 3, 32, 32)
+    X_A = np.random.normal(size=inp_shape).astype(precision)
+    X_B = np.random.normal(size=inp_shape).astype(precision)
+
+    pads = (1, 1, 1, 1)
+
+    # Register plugin creator
+    plg_registry = trt.get_plugin_registry()
+    my_plugin_creator = CircPadPluginCreator()
+    plg_registry.register_creator(my_plugin_creator, "")
+
+    # create ONNX model
+    onnx_path = f"test_CircPadPlugin_multi_tactic_{args.precision}.onnx"
+    inputA = gs.Variable(name="X_A", shape=inp_shape, dtype=precision)
+    inputB = gs.Variable(name="X_B", shape=inp_shape, dtype=precision)
+    Y_A = gs.Variable(name="Y_A", dtype=precision)
+    Y_B = gs.Variable(name="Y_B", dtype=precision)
+    myPluginNode_A = gs.Node(
+        name="CircPadPlugin_A",
+        op="CircPadPlugin",
+        inputs=[inputA],
+        outputs=[Y_A],
+        attrs={
+            "pads": pads,
+            "per_format_tactics": np.array([is_tactics_per_format], dtype=np.int32),
+        },
+    )
+    myPluginNode_B = gs.Node(
+        name="CircPadPlugin_B",
+        op="CircPadPlugin",
+        inputs=[inputB],
+        outputs=[Y_B],
+        attrs={
+            "pads": pads,
+            "per_format_tactics": np.array([is_tactics_per_format], dtype=np.int32),
+        },
+    )
+
+    graph = gs.Graph(nodes=[myPluginNode_A, myPluginNode_B], inputs=[inputA, inputB], outputs=[Y_A, Y_B], opset=16)
+    onnx.save(gs.export_onnx(graph), onnx_path)
+
+    # build engine
+    build_engine = EngineFromNetwork(
+        NetworkFromOnnxPath(onnx_path, strongly_typed=True), CreateConfig()
+    )
+
+    Y_A_ref = np.pad(X_A, [[0, 0], [0, 0], [pads[0], pads[1]], [pads[2], pads[3]]], "wrap")
+    Y_B_ref = np.pad(X_B, [[0, 0], [0, 0], [pads[0], pads[1]], [pads[2], pads[3]]], "wrap")
+
+    # Run
+    with TrtRunner(build_engine, "trt_runner")as runner:
+        outputs = runner.infer({"X_A": X_A, "X_B": X_B})
+        Y_A_out = outputs["Y_A"]
+        Y_B_out = outputs["Y_B"]
+
+        if np.allclose(Y_A_out, Y_A_ref):
+            print("Inference result A correct!")
+        else:
+            print("Inference result A incorrect!")
+
+        if np.allclose(Y_B_out, Y_B_ref):
+            print("Inference result B correct!")
+        else:
+            print("Inference result B incorrect!")

samples/python/python_plugin/circ_pad_plugin_numba.py

@@ -0,0 +1,257 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2025 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 onnx_graphsurgeon as gs
+import numpy as np
+import onnx
+import cupy as cp
+from numba import cuda
+import sys
+import os
+
+import tensorrt as trt
+from polygraphy.backend.trt import (
+    CreateConfig,
+    EngineFromNetwork,
+    NetworkFromOnnxPath,
+    TrtRunner,
+)
+
+from polygraphy.json import to_json, from_json
+
+sys.path.insert(1, os.path.join(os.path.dirname(os.path.realpath(__file__)), os.pardir))
+from plugin_utils import volume, parseArgs
+
+
+
+@cuda.jit
+def circ_pad(X, all_pads, orig_dims, Y, Y_shape, Y_len):
+    index = cuda.blockIdx.x * cuda.blockDim.x + cuda.threadIdx.x
+    stride = cuda.blockDim.x * cuda.gridDim.x
+
+    for i in range(index, Y_len, stride):
+        i3 = int(i % Y_shape[3])
+        i2 = int((i // Y_shape[3]) % Y_shape[2])
+        i1 = int((i // Y_shape[3] // Y_shape[2]) % Y_shape[1])
+        i0 = int(i // Y_shape[3] // Y_shape[2] // Y_shape[1])
+
+        j0 = int((i0 - all_pads[0]) % orig_dims[0])
+        j1 = int((i1 - all_pads[2]) % orig_dims[1])
+        j2 = int((i2 - all_pads[4]) % orig_dims[2])
+        j3 = int((i3 - all_pads[6]) % orig_dims[3])
+
+        Y[i] = X[
+            int(
+                orig_dims[3] * orig_dims[2] * orig_dims[1] * j0
+                + orig_dims[3] * orig_dims[2] * j1
+                + orig_dims[3] * j2
+                + j3
+            )
+        ]
+
+
+class CircPadPlugin(trt.IPluginV2DynamicExt):
+    def __init__(self, fc=None):
+        trt.IPluginV2DynamicExt.__init__(self)
+        self.pads = []
+        self.X_shape = []
+
+        self.num_outputs = 1
+        self.plugin_namespace = ""
+        self.plugin_type = "CircPadPlugin"
+        self.plugin_version = "1"
+
+        if fc is not None:
+            assert fc[0].name == "pads"
+            self.pads = fc[0].data
+
+    def get_output_datatype(self, index, input_types):
+        return input_types[0]
+
+    def get_output_dimensions(self, output_index, inputs, exprBuilder):
+
+        output_dims = trt.DimsExprs(inputs[0])
+
+        for i in range(np.size(self.pads) // 2):
+            output_dims[len(output_dims) - i - 1] = exprBuilder.operation(
+                trt.DimensionOperation.SUM,
+                inputs[0][len(output_dims) - i - 1],
+                exprBuilder.constant(self.pads[i * 2] + self.pads[i * 2 + 1]),
+            )
+
+        return output_dims
+
+    def serialize(self):
+        return to_json({"pads": self.pads})
+
+    def configure_plugin(self, inp, out):
+        X_dims = inp[0].desc.dims
+        self.X_shape = np.zeros((len(X_dims),))
+        for i in range(len(X_dims)):
+            self.X_shape[i] = X_dims[i]
+
+    def supports_format_combination(self, pos, in_out, num_inputs):
+        assert num_inputs == 1
+        assert pos < len(in_out)
+
+        desc = in_out[pos]
+        if desc.format != trt.TensorFormat.LINEAR:
+            return False
+
+        # first input should be float16 or float32
+        if pos == 0:
+            return desc.type == trt.DataType.FLOAT or desc.type == trt.DataType.HALF
+
+        # output should have the same type as the input
+        if pos == 1:
+            return in_out[0].type == desc.type
+
+        assert False
+
+    def enqueue(self, input_desc, output_desc, inputs, outputs, workspace, stream):
+
+        inp_dtype = trt.nptype(input_desc[0].type)
+
+        a_mem = cp.cuda.UnownedMemory(
+            inputs[0], volume(input_desc[0].dims) * cp.dtype(inp_dtype).itemsize, self
+        )
+        c_mem = cp.cuda.UnownedMemory(
+            outputs[0],
+            volume(output_desc[0].dims) * cp.dtype(inp_dtype).itemsize,
+            self,
+        )
+
+        a_ptr = cp.cuda.MemoryPointer(a_mem, 0)
+        c_ptr = cp.cuda.MemoryPointer(c_mem, 0)
+
+        a = cp.ndarray((volume(input_desc[0].dims)), dtype=inp_dtype, memptr=a_ptr)
+        c = cp.ndarray((volume(output_desc[0].dims)), dtype=inp_dtype, memptr=c_ptr)
+
+        numba_stream = cuda.external_stream(stream)
+
+        N = len(self.X_shape)
+        all_pads = np.zeros((N * 2,))
+        orig_dims = np.array(self.X_shape)
+        out_dims = np.array(self.X_shape)
+
+        for i in range(np.size(pads) // 2):
+            out_dims[N - i - 1] += pads[i * 2] + pads[i * 2 + 1]
+            all_pads[N * 2 - 2 * i - 2] = pads[i * 2]
+            all_pads[N * 2 - 2 * i - 1] = pads[i * 2 + 1]
+
+        all_pads_d = cp.asarray(all_pads)
+        orig_dims_d = cp.asarray(orig_dims)
+        Y_shape_d = cp.asarray(out_dims)
+
+        blockSize = 256
+        numBlocks = int((np.prod(out_dims) + blockSize - 1) // blockSize)
+
+        circ_pad[numBlocks, blockSize, numba_stream](
+            a, all_pads_d, orig_dims_d, c, Y_shape_d, np.prod(out_dims)
+        )
+
+        return 0
+
+    def clone(self):
+        cloned_plugin = CircPadPlugin()
+        cloned_plugin.__dict__.update(self.__dict__)
+        return cloned_plugin
+
+    #
+    # The following defaults take effect since the respective methods are not overriden
+    #
+
+    # def initialize(self):
+    #     pass
+
+    # def get_serialization_size(self):
+    #     return len(to_json({"pads": self.pads}))
+
+    # def get_workspace_size(self, input_desc, output_desc):
+    #     return 0
+
+    # def destroy(self):
+    #     pass
+
+    # def terminate(self):
+    #     pass
+
+
+class CircPadPluginCreator(trt.IPluginCreator):
+    def __init__(self):
+        trt.IPluginCreator.__init__(self)
+        self.name = "CircPadPlugin"
+        self.plugin_namespace = ""
+        self.plugin_version = "1"
+        self.field_names = trt.PluginFieldCollection(
+            [trt.PluginField("pads", np.array([]), trt.PluginFieldType.INT32)]
+        )
+
+    def create_plugin(self, name, fc):
+        return CircPadPlugin(fc)
+
+    def deserialize_plugin(self, name, data):
+        j = dict(from_json(data.decode("utf-8")))
+        deserialized = CircPadPlugin()
+        deserialized.__dict__.update(j)
+        return deserialized
+
+
+if __name__ == "__main__":
+
+    args = parseArgs()
+    precision = np.float32 if args.precision == "fp32" else np.float16
+
+    inp_shape = (10, 3, 32, 32)
+    X = np.random.normal(size=inp_shape).astype(precision)
+
+    pads = (1, 1, 1, 1)
+
+    # Register plugin creator
+    plg_registry = trt.get_plugin_registry()
+    my_plugin_creator = CircPadPluginCreator()
+    plg_registry.register_creator(my_plugin_creator, "")
+
+    # create ONNX model
+    onnx_path = f"test_CircPadPlugin_numba_{args.precision}.onnx"
+    inputA = gs.Variable(name="X", shape=inp_shape, dtype=precision)
+    Y = gs.Variable(name="Y", dtype=precision)
+    myPluginNode = gs.Node(
+        name="CircPadPlugin",
+        op="CircPadPlugin",
+        inputs=[inputA],
+        outputs=[Y],
+        attrs={"pads": pads},
+    )
+    graph = gs.Graph(nodes=[myPluginNode], inputs=[inputA], outputs=[Y], opset=16)
+    onnx.save(gs.export_onnx(graph), onnx_path)
+
+    # build engine
+    build_engine = EngineFromNetwork(
+        NetworkFromOnnxPath(onnx_path, strongly_typed=True), CreateConfig()
+    )
+
+    Y_ref = np.pad(X, [[0, 0], [0, 0], [pads[0], pads[1]], [pads[2], pads[3]]], "wrap")
+    # Run
+    with TrtRunner(build_engine, "trt_runner") as runner:
+        outputs = runner.infer({"X": X})
+        Y = outputs["Y"]
+
+        if np.allclose(Y, Y_ref):
+            print("Inference result correct!")
+        else:
+            print("Inference result incorrect!")

samples/python/python_plugin/circ_pad_plugin_torch.py

@@ -0,0 +1,214 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2025 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 onnx_graphsurgeon as gs
+import numpy as np
+import onnx
+import cupy as cp
+import sys
+import os
+
+import tensorrt as trt
+from polygraphy.backend.trt import (
+    CreateConfig,
+    EngineFromNetwork,
+    NetworkFromOnnxPath,
+    TrtRunner,
+)
+
+from polygraphy.json import to_json, from_json
+import torch
+
+sys.path.insert(1, os.path.join(os.path.dirname(os.path.realpath(__file__)), os.pardir))
+from plugin_utils import volume, parseArgs
+
+
+
+class CircPadPlugin(trt.IPluginV2DynamicExt):
+    def __init__(self, fc=None):
+        trt.IPluginV2DynamicExt.__init__(self)
+        self.pads = []
+        self.X_shape = []
+
+        self.num_outputs = 1
+        self.plugin_namespace = ""
+        self.plugin_type = "CircPadPlugin"
+        self.plugin_version = "1"
+
+        if fc is not None:
+            assert fc[0].name == "pads"
+            self.pads = fc[0].data
+
+    def get_output_datatype(self, index, input_types):
+        return input_types[0]
+
+    def get_output_dimensions(self, output_index, inputs, exprBuilder):
+
+        output_dims = trt.DimsExprs(inputs[0])
+
+        for i in range(np.size(self.pads) // 2):
+            output_dims[len(output_dims) - i - 1] = exprBuilder.operation(
+                trt.DimensionOperation.SUM,
+                inputs[0][len(output_dims) - i - 1],
+                exprBuilder.constant(self.pads[i * 2] + self.pads[i * 2 + 1]),
+            )
+
+        return output_dims
+
+    def serialize(self):
+        return to_json({"pads": self.pads})
+
+    def configure_plugin(self, inp, out):
+        X_dims = inp[0].desc.dims
+        self.X_shape = np.zeros((len(X_dims),))
+        for i in range(len(X_dims)):
+            self.X_shape[i] = X_dims[i]
+
+    def supports_format_combination(self, pos, in_out, num_inputs):
+        assert num_inputs == 1
+        assert pos < len(in_out)
+
+        desc = in_out[pos]
+        if desc.format != trt.TensorFormat.LINEAR:
+            return False
+
+        # first input should be float16 or float32
+        if pos == 0:
+            return desc.type == trt.DataType.FLOAT or desc.type == trt.DataType.HALF
+
+        # output should have the same type as the input
+        if pos == 1:
+            return in_out[0].type == desc.type
+
+        assert False
+
+    def enqueue(self, input_desc, output_desc, inputs, outputs, workspace, stream):
+
+        inp_dtype = trt.nptype(input_desc[0].type)
+
+        a_mem = cp.cuda.UnownedMemory(
+            inputs[0], volume(input_desc[0].dims) * cp.dtype(inp_dtype).itemsize, self
+        )
+        c_mem = cp.cuda.UnownedMemory(
+            outputs[0],
+            volume(output_desc[0].dims) * cp.dtype(inp_dtype).itemsize,
+            self,
+        )
+
+        a_ptr = cp.cuda.MemoryPointer(a_mem, 0)
+        c_ptr = cp.cuda.MemoryPointer(c_mem, 0)
+
+        a_d = cp.ndarray(tuple(input_desc[0].dims), dtype=inp_dtype, memptr=a_ptr)
+        c_d = cp.ndarray((volume(output_desc[0].dims)), dtype=inp_dtype, memptr=c_ptr)
+
+        a_t = torch.as_tensor(a_d, device="cuda")
+
+        # Use PyTorch functional op - no need to write kernel
+        out = torch.nn.functional.pad(a_t, self.pads.tolist(), mode="circular")
+        cp.copyto(c_d, cp.reshape(cp.asarray(out), (-1,)))
+
+        return 0
+
+    def clone(self):
+        cloned_plugin = CircPadPlugin()
+        cloned_plugin.__dict__.update(self.__dict__)
+        return cloned_plugin
+
+    #
+    # The following defaults take effect since the respective methods are not overriden
+    #
+
+    # def initialize(self):
+    #     pass
+
+    # def get_serialization_size(self):
+    #     return len(to_json({"pads": self.pads}))
+
+    # def get_workspace_size(self, input_desc, output_desc):
+    #     return 0
+
+    # def destroy(self):
+    #     pass
+
+    # def terminate(self):
+    #     pass
+
+
+class CircPadPluginCreator(trt.IPluginCreator):
+    def __init__(self):
+        trt.IPluginCreator.__init__(self)
+        self.name = "CircPadPlugin"
+        self.plugin_namespace = ""
+        self.plugin_version = "1"
+        self.field_names = trt.PluginFieldCollection(
+            [trt.PluginField("pads", np.array([]), trt.PluginFieldType.INT32)]
+        )
+
+    def create_plugin(self, name, fc):
+        return CircPadPlugin(fc)
+
+    def deserialize_plugin(self, name, data):
+        j = dict(from_json(data.decode("utf-8")))
+        deserialized = CircPadPlugin()
+        deserialized.__dict__.update(j)
+        return deserialized
+
+
+if __name__ == "__main__":
+
+    args = parseArgs()
+    precision = np.float32 if args.precision == "fp32" else np.float16
+
+    inp_shape = (10, 3, 32, 32)
+    X = np.random.normal(size=inp_shape).astype(precision)
+
+    pads = (1, 1, 1, 1)
+
+    # Register plugin creator
+    plg_registry = trt.get_plugin_registry()
+    my_plugin_creator = CircPadPluginCreator()
+    plg_registry.register_creator(my_plugin_creator, "")
+
+    # create ONNX model
+    onnx_path = f"test_CircPadPlugin_torch_{args.precision}.onnx"
+    inputA = gs.Variable(name="X", shape=inp_shape, dtype=precision)
+    Y = gs.Variable(name="Y", dtype=precision)
+    myPluginNode = gs.Node(
+        name="CircPadPlugin",
+        op="CircPadPlugin",
+        inputs=[inputA],
+        outputs=[Y],
+        attrs={"pads": pads},
+    )
+    graph = gs.Graph(nodes=[myPluginNode], inputs=[inputA], outputs=[Y], opset=16)
+    onnx.save(gs.export_onnx(graph), onnx_path)
+
+    # build engine
+    build_engine = EngineFromNetwork(
+        NetworkFromOnnxPath(onnx_path, strongly_typed=True), CreateConfig()
+    )
+
+    Y_ref = np.pad(X, [[0, 0], [0, 0], [pads[0], pads[1]], [pads[2], pads[3]]], "wrap")
+    # Run
+    with TrtRunner(build_engine, "trt_runner") as runner:
+        outputs = runner.infer({"X": X})
+        Y = outputs["Y"]
+
+        if np.allclose(Y, Y_ref):
+            print("Inference result correct!")
+        else:
+            print("Inference result incorrect!")

samples/python/python_plugin/circ_pad_plugin_triton.py

@@ -0,0 +1,297 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2025 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 onnx_graphsurgeon as gs
+import numpy as np
+import onnx
+import cupy as cp
+import sys
+import os
+
+import triton
+import triton.language as tl
+
+import tensorrt as trt
+from polygraphy.backend.trt import (
+    CreateConfig,
+    EngineFromNetwork,
+    NetworkFromOnnxPath,
+    TrtRunner,
+)
+
+from polygraphy.json import to_json, from_json
+import torch
+
+sys.path.insert(1, os.path.join(os.path.dirname(os.path.realpath(__file__)), os.pardir))
+from plugin_utils import volume, parseArgs
+
+
+
+@triton.jit
+def circ_pad(
+    X,
+    all_pads_0,
+    all_pads_2,
+    all_pads_4,
+    all_pads_6,
+    orig_dims_0,
+    orig_dims_1,
+    orig_dims_2,
+    orig_dims_3,
+    Y,
+    Y_shape_1,
+    Y_shape_2,
+    Y_shape_3,
+    X_len,
+    Y_len,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    i = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+
+    mask_y = i < Y_len
+
+    i3 = i % Y_shape_3
+    i2 = (i // Y_shape_3) % Y_shape_2
+    i1 = (i // Y_shape_3 // Y_shape_2) % Y_shape_1
+    i0 = i // Y_shape_3 // Y_shape_2 // Y_shape_1
+
+    j0 = (i0 - all_pads_0 + orig_dims_0) % orig_dims_0
+    j1 = (i1 - all_pads_2 + orig_dims_1) % orig_dims_1
+    j2 = (i2 - all_pads_4 + orig_dims_2) % orig_dims_2
+    j3 = (i3 - all_pads_6 + orig_dims_3) % orig_dims_3
+
+    load_idx = (
+        orig_dims_3 * orig_dims_2 * orig_dims_1 * j0
+        + orig_dims_3 * orig_dims_2 * j1
+        + orig_dims_3 * j2
+        + j3
+    )
+    mask_x = load_idx < X_len
+
+    x = tl.load(X + load_idx, mask=mask_x)
+
+    tl.store(Y + i, x, mask=mask_y)
+
+
+class CircPadPlugin(trt.IPluginV2DynamicExt):
+    def __init__(self, fc=None):
+        trt.IPluginV2DynamicExt.__init__(self)
+        self.pads = []
+        self.X_shape = []
+
+        self.num_outputs = 1
+        self.plugin_namespace = ""
+        self.plugin_type = "CircPadPlugin"
+        self.plugin_version = "1"
+
+        if fc is not None:
+            assert fc[0].name == "pads"
+            self.pads = fc[0].data
+
+    def get_output_datatype(self, index, input_types):
+        return input_types[0]
+
+    def get_output_dimensions(self, output_index, inputs, exprBuilder):
+
+        output_dims = trt.DimsExprs(inputs[0])
+
+        for i in range(np.size(self.pads) // 2):
+            output_dims[len(output_dims) - i - 1] = exprBuilder.operation(
+                trt.DimensionOperation.SUM,
+                inputs[0][len(output_dims) - i - 1],
+                exprBuilder.constant(self.pads[i * 2] + self.pads[i * 2 + 1]),
+            )
+
+        return output_dims
+
+    def serialize(self):
+        return to_json({"pads": self.pads})
+
+    def configure_plugin(self, inp, out):
+        X_dims = inp[0].desc.dims
+        self.X_shape = np.zeros((len(X_dims),))
+        for i in range(len(X_dims)):
+            self.X_shape[i] = X_dims[i]
+
+    def supports_format_combination(self, pos, in_out, num_inputs):
+        assert num_inputs == 1
+        assert pos < len(in_out)
+
+        desc = in_out[pos]
+        if desc.format != trt.TensorFormat.LINEAR:
+            return False
+
+        # first input should be float16 or float32
+        if pos == 0:
+            return desc.type == trt.DataType.FLOAT or desc.type == trt.DataType.HALF
+
+        # output should have the same type as the input
+        if pos == 1:
+            return in_out[0].type == desc.type
+
+        assert False
+
+    def enqueue(self, input_desc, output_desc, inputs, outputs, workspace, stream):
+
+        inp_dtype = trt.nptype(input_desc[0].type)
+
+        a_mem = cp.cuda.UnownedMemory(
+            inputs[0], volume(input_desc[0].dims) * cp.dtype(inp_dtype).itemsize, self
+        )
+        c_mem = cp.cuda.UnownedMemory(
+            outputs[0],
+            volume(output_desc[0].dims) * cp.dtype(inp_dtype).itemsize,
+            self,
+        )
+
+        a_ptr = cp.cuda.MemoryPointer(a_mem, 0)
+        c_ptr = cp.cuda.MemoryPointer(c_mem, 0)
+
+        a_d = cp.ndarray((volume(input_desc[0].dims)), dtype=inp_dtype, memptr=a_ptr)
+        c_d = cp.ndarray((volume(output_desc[0].dims)), dtype=inp_dtype, memptr=c_ptr)
+
+        a_t = torch.as_tensor(a_d, device="cuda")
+        c_t = torch.as_tensor(c_d, device="cuda")
+
+        N = len(self.X_shape)
+        all_pads = np.zeros((N * 2,), dtype=np.int32)
+        orig_dims = np.array(self.X_shape, dtype=np.int32)
+        out_dims = np.array(self.X_shape, dtype=np.int32)
+
+        for i in range(np.size(pads) // 2):
+            out_dims[N - i - 1] += pads[i * 2] + pads[i * 2 + 1]
+            all_pads[N * 2 - 2 * i - 2] = pads[i * 2]
+            all_pads[N * 2 - 2 * i - 1] = pads[i * 2 + 1]
+
+        all_pads = all_pads.tolist()
+        orig_dims = orig_dims.tolist()
+        out_dims = out_dims.tolist()
+
+        blockSize = 256
+        numBlocks = (int((np.prod(out_dims) + blockSize - 1) // blockSize),)
+
+        circ_pad[numBlocks](
+            a_t,
+            all_pads[0],
+            all_pads[2],
+            all_pads[4],
+            all_pads[6],
+            orig_dims[0],
+            orig_dims[1],
+            orig_dims[2],
+            orig_dims[3],
+            c_t,
+            out_dims[1],
+            out_dims[2],
+            out_dims[3],
+            int(np.prod(orig_dims)),
+            int(np.prod(out_dims)),
+            BLOCK_SIZE=256,
+        )
+
+        return 0
+
+    def clone(self):
+        cloned_plugin = CircPadPlugin()
+        cloned_plugin.__dict__.update(self.__dict__)
+        return cloned_plugin
+
+    #
+    # The following defaults take effect since the respective methods are not overriden
+    #
+
+    # def initialize(self):
+    #     pass
+
+    # def get_serialization_size(self):
+    #     return len(to_json({"pads": self.pads}))
+
+    # def get_workspace_size(self, input_desc, output_desc):
+    #     return 0
+
+    # def destroy(self):
+    #     pass
+
+    # def terminate(self):
+    #     pass
+
+
+class CircPadPluginCreator(trt.IPluginCreator):
+    def __init__(self):
+        trt.IPluginCreator.__init__(self)
+        self.name = "CircPadPlugin"
+        self.plugin_namespace = ""
+        self.plugin_version = "1"
+        self.field_names = trt.PluginFieldCollection(
+            [trt.PluginField("pads", np.array([]), trt.PluginFieldType.INT32)]
+        )
+
+    def create_plugin(self, name, fc):
+        return CircPadPlugin(fc)
+
+    def deserialize_plugin(self, name, data):
+        j = dict(from_json(data.decode("utf-8")))
+        deserialized = CircPadPlugin()
+        deserialized.__dict__.update(j)
+        return deserialized
+
+
+if __name__ == "__main__":
+
+    args = parseArgs()
+    precision = np.float32 if args.precision == "fp32" else np.float16
+
+    inp_shape = (10, 3, 32, 32)
+    X = np.random.normal(size=inp_shape).astype(precision)
+
+    pads = (1, 1, 1, 1)
+
+    # Register plugin creator
+    plg_registry = trt.get_plugin_registry()
+    my_plugin_creator = CircPadPluginCreator()
+    plg_registry.register_creator(my_plugin_creator, "")
+
+    # create ONNX model
+    onnx_path = f"test_CircPadPlugin_triton_{args.precision}.onnx"
+    inputA = gs.Variable(name="X", shape=inp_shape, dtype=precision)
+    Y = gs.Variable(name="Y", dtype=precision)
+    myPluginNode = gs.Node(
+        name="CircPadPlugin",
+        op="CircPadPlugin",
+        inputs=[inputA],
+        outputs=[Y],
+        attrs={"pads": pads},
+    )
+    graph = gs.Graph(nodes=[myPluginNode], inputs=[inputA], outputs=[Y], opset=16)
+    onnx.save(gs.export_onnx(graph), onnx_path)
+
+    # build engine
+    build_engine = EngineFromNetwork(
+        NetworkFromOnnxPath(onnx_path, strongly_typed=True), CreateConfig()
+    )
+
+    Y_ref = np.pad(X, [[0, 0], [0, 0], [pads[0], pads[1]], [pads[2], pads[3]]], "wrap")
+    # Run
+    with TrtRunner(build_engine, "trt_runner") as runner:
+        outputs = runner.infer({"X": X})
+        Y = outputs["Y"]
+
+        if np.allclose(Y, Y_ref):
+            print("Inference result correct!")
+        else:
+            print("Inference result incorrect!")

samples/python/python_plugin/circ_plugin_cpp/circ_pad_plugin.cu

@@ -0,0 +1,387 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2026 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.
+ */
+
+#include "NvInfer.h"
+
+#include <algorithm>
+#include <cstdint>
+#include <cstring>
+#include <string_view>
+#include <iostream>
+#include <memory>
+#include <numeric>
+#include <vector>
+
+#include <cuda.h>
+#include <cuda_fp16.h>
+
+using namespace nvinfer1;
+using namespace std::string_view_literals;
+
+static void caughtError(std::exception const& e)
+{
+    std::cout << e.what() << std::endl;
+}
+
+#define ASSERT(condition)                                                                                              \
+    do                                                                                                                 \
+    {                                                                                                                  \
+        if (!(condition))                                                                                              \
+        {                                                                                                              \
+            std::cout << "Assertion failure: " << #condition << std::endl;                                             \
+            abort();                                                                                                   \
+        }                                                                                                              \
+    } while (0)
+
+template <typename Dtype>
+struct CudaBind
+{
+    size_t mSize;
+    Dtype* mPtr;
+
+    CudaBind(size_t size)
+    {
+        mSize = size;
+        ASSERT(!cudaMalloc((void**) &mPtr, sizeof(Dtype) * mSize));
+    }
+
+    ~CudaBind()
+    {
+        if (mPtr != nullptr)
+        {
+            ASSERT(!cudaFree(mPtr));
+            mPtr = nullptr;
+        }
+    }
+};
+
+static int64_t volume(Dims const& dims)
+{
+    return std::accumulate(dims.d, dims.d + dims.nbDims, int64_t{1}, std::multiplies<int64_t>{});
+}
+
+template <typename T>
+__global__ void circPadKernel(
+    T const* x, int32_t const* allPads, int32_t const* origDims, T* y, int32_t const* yShape, int32_t yLen)
+{
+    int32_t index = blockIdx.x * blockDim.x + threadIdx.x;
+    int32_t stride = blockDim.x * gridDim.x;
+
+    for (int32_t i = index; i < yLen; i += stride)
+    {
+        int32_t i3 = i % yShape[3];
+        int32_t i2 = (i / yShape[3]) % yShape[2];
+        int32_t i1 = (i / yShape[3] / yShape[2]) % yShape[1];
+        int32_t i0 = i / yShape[3] / yShape[2] / yShape[1];
+
+        int32_t j0 = (i0 - allPads[0] + origDims[0]) % origDims[0];
+        int32_t j1 = (i1 - allPads[2] + origDims[1]) % origDims[1];
+        int32_t j2 = (i2 - allPads[4] + origDims[2]) % origDims[2];
+        int32_t j3 = (i3 - allPads[6] + origDims[3]) % origDims[3];
+
+        y[i] = x[origDims[3] * origDims[2] * origDims[1] * j0 + origDims[3] * origDims[2] * j1 + origDims[3] * j2 + j3];
+    }
+}
+
+class CircPadPlugin : public IPluginV3,
+                      public IPluginV3OneCore,
+                      public IPluginV3OneBuild,
+                      public IPluginV3OneRuntime
+{
+public:
+    CircPadPlugin() = default;
+
+    CircPadPlugin(std::vector<int32_t> pads)
+        : mPads(std::move(pads))
+    {
+    }
+
+    CircPadPlugin(CircPadPlugin const& p) = default;
+
+    ~CircPadPlugin() override = default;
+
+    int32_t getNbOutputs() const noexcept override
+    {
+        return 1;
+    }
+
+    bool supportsFormatCombination(
+        int32_t pos, DynamicPluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept override
+    {
+        PluginTensorDesc const& desc = inOut[pos].desc;
+        if (desc.format != TensorFormat::kLINEAR)
+        {
+            return false;
+        }
+
+        // first input should be float16 or float32
+        if (pos == 0)
+        {
+            return (desc.type == DataType::kFLOAT || desc.type == DataType::kHALF);
+        }
+
+        // output should have the same type as the input
+        if (pos == 1)
+        {
+            return (desc.type == inOut[0].desc.type);
+        }
+
+        return false;
+    }
+
+    int32_t enqueue(PluginTensorDesc const* inputDesc, PluginTensorDesc const* outputDesc, void const* const* inputs,
+        void* const* outputs, void* workspace, cudaStream_t stream) noexcept override
+    {
+        auto inpDType = inputDesc[0].type;
+
+        int32_t const blockSize = 256;
+        int32_t const numBlocks = (volume(outputDesc[0].dims) + blockSize - 1) / blockSize;
+
+        ASSERT(inpDType == DataType::kFLOAT || inpDType == DataType::kHALF);
+
+        if (inpDType == DataType::kFLOAT)
+        {
+            circPadKernel<float><<<numBlocks, blockSize, 0, stream>>>(static_cast<float const*>(inputs[0]),
+                mAllPadsPtr->mPtr, mOrigDimsPtr->mPtr, static_cast<float*>(outputs[0]), mOutDimsPtr->mPtr,
+                volume(outputDesc[0].dims));
+        }
+        else if (inpDType == DataType::kHALF)
+        {
+            circPadKernel<half><<<numBlocks, blockSize, 0, stream>>>(static_cast<half const*>(inputs[0]),
+                mAllPadsPtr->mPtr, mOrigDimsPtr->mPtr, static_cast<half*>(outputs[0]), mOutDimsPtr->mPtr,
+                volume(outputDesc[0].dims));
+        }
+        return 0;
+    }
+
+    char const* getPluginName() const noexcept override
+    {
+        return "CircPadPlugin";
+    }
+
+    char const* getPluginVersion() const noexcept override
+    {
+        return "1";
+    }
+
+    IPluginV3* clone() noexcept override
+    {
+        try
+        {
+            auto plugin = std::make_unique<CircPadPlugin>(*this);
+            // Build-time clones do not need GPU memory. Clear shared_ptrs so the
+            // clone does not share GPU allocations with the source.
+            plugin->mAllPadsPtr.reset();
+            plugin->mOrigDimsPtr.reset();
+            plugin->mOutDimsPtr.reset();
+            return plugin.release();
+        }
+        catch (std::exception const& e)
+        {
+            caughtError(e);
+        }
+        return nullptr;
+    }
+
+    void setPluginNamespace(char const* libNamespace) noexcept
+    {
+        mNamespace = libNamespace;
+    }
+
+    char const* getPluginNamespace() const noexcept override
+    {
+        return mNamespace.c_str();
+    }
+
+    int32_t getOutputDataTypes(
+        DataType* outputTypes, int32_t nbOutputs, DataType const* inputTypes, int32_t nbInputs) const noexcept override
+    {
+        outputTypes[0] = inputTypes[0];
+        return 0;
+    }
+
+    int32_t getOutputShapes(DimsExprs const* inputs, int32_t nbInputs, DimsExprs const* shapeInputs,
+        int32_t nbShapeInputs, DimsExprs* outputs, int32_t nbOutputs, IExprBuilder& exprBuilder) noexcept override
+    {
+        outputs[0] = inputs[0];
+        int32_t nbOutDims = inputs[0].nbDims;
+
+        for (int32_t i = 0; i < static_cast<int32_t>(mPads.size()) / 2; ++i)
+        {
+            outputs[0].d[nbOutDims - i - 1] = exprBuilder.operation(DimensionOperation::kSUM,
+                *inputs[0].d[nbOutDims - i - 1], *exprBuilder.constant(mPads[i * 2] + mPads[i * 2 + 1]));
+        }
+
+        return 0;
+    }
+
+    int32_t configurePlugin(DynamicPluginTensorDesc const* in, int32_t nbInputs,
+        DynamicPluginTensorDesc const* out, int32_t nbOutputs) noexcept override
+    {
+        return 0;
+    }
+
+    size_t getWorkspaceSize(DynamicPluginTensorDesc const* inputs, int32_t nbInputs,
+        DynamicPluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept override
+    {
+        return 0;
+    }
+
+    IPluginCapability* getCapabilityInterface(PluginCapabilityType type) noexcept override
+    {
+        if (type == PluginCapabilityType::kBUILD)
+        {
+            return static_cast<IPluginV3OneBuild*>(this);
+        }
+        if (type == PluginCapabilityType::kRUNTIME)
+        {
+            return static_cast<IPluginV3OneRuntime*>(this);
+        }
+        ASSERT(type == PluginCapabilityType::kCORE);
+        return static_cast<IPluginV3OneCore*>(this);
+    }
+
+    int32_t onShapeChange(
+        PluginTensorDesc const* in, int32_t nbInputs, PluginTensorDesc const* out, int32_t nbOutputs) noexcept override
+    {
+        mN = in[0].dims.nbDims;
+
+        std::vector<int32_t> allPads(mN * 2);
+        std::vector<int32_t> origDims(mN);
+        std::vector<int32_t> outDims(mN);
+
+        for (int32_t i = 0; i < mN; ++i)
+        {
+            origDims[i] = in[0].dims.d[i];
+            outDims[i] = in[0].dims.d[i];
+        }
+
+        for (int32_t i = 0; i < static_cast<int32_t>(mPads.size()) / 2; ++i)
+        {
+            outDims[mN - i - 1] += mPads[i * 2] + mPads[i * 2 + 1];
+            allPads[mN * 2 - 2 * i - 2] = mPads[i * 2];
+            allPads[mN * 2 - 2 * i - 1] = mPads[i * 2 + 1];
+        }
+
+        mAllPadsPtr = std::make_shared<CudaBind<int32_t>>(mN * 2);
+        mOrigDimsPtr = std::make_shared<CudaBind<int32_t>>(mN);
+        mOutDimsPtr = std::make_shared<CudaBind<int32_t>>(mN);
+
+        ASSERT(
+            !cudaMemcpy(mAllPadsPtr->mPtr, &allPads.front(), allPads.size() * sizeof(int32_t), cudaMemcpyHostToDevice));
+        ASSERT(!cudaMemcpy(
+            mOrigDimsPtr->mPtr, &origDims.front(), origDims.size() * sizeof(int32_t), cudaMemcpyHostToDevice));
+        ASSERT(
+            !cudaMemcpy(mOutDimsPtr->mPtr, &outDims.front(), outDims.size() * sizeof(int32_t), cudaMemcpyHostToDevice));
+
+        return 0;
+    }
+
+    IPluginV3* attachToContext(IPluginResourceContext* context) noexcept override
+    {
+        return clone();
+    }
+
+    PluginFieldCollection const* getFieldsToSerialize() noexcept override
+    {
+        mDataToSerialize.clear();
+        mDataToSerialize.emplace_back("pads", mPads.data(), PluginFieldType::kINT32, mPads.size());
+        mFCToSerialize.nbFields = mDataToSerialize.size();
+        mFCToSerialize.fields = mDataToSerialize.data();
+        return &mFCToSerialize;
+    }
+
+private:
+    std::vector<int32_t> mPads{};
+    int32_t mN{};
+    std::shared_ptr<CudaBind<int32_t>> mAllPadsPtr{};
+    std::shared_ptr<CudaBind<int32_t>> mOrigDimsPtr{};
+    std::shared_ptr<CudaBind<int32_t>> mOutDimsPtr{};
+    std::string mNamespace;
+
+    std::vector<PluginField> mDataToSerialize;
+    PluginFieldCollection mFCToSerialize;
+};
+
+class CircPadPluginCreator : public IPluginCreatorV3One
+{
+public:
+    CircPadPluginCreator()
+    {
+        mPluginAttributes.clear();
+        mPluginAttributes.emplace_back(PluginField("pads", nullptr, PluginFieldType::kINT32, 1));
+        mFC.nbFields = mPluginAttributes.size();
+        mFC.fields = mPluginAttributes.data();
+    }
+
+    char const* getPluginName() const noexcept override
+    {
+        return "CircPadPlugin";
+    }
+
+    char const* getPluginVersion() const noexcept override
+    {
+        return "1";
+    }
+
+    PluginFieldCollection const* getFieldNames() noexcept override
+    {
+        return &mFC;
+    }
+
+    IPluginV3* createPlugin(char const* name, PluginFieldCollection const* fc, TensorRTPhase phase) noexcept override
+    {
+        try
+        {
+            std::vector<int32_t> pads;
+
+            for (int32_t i = 0; i < fc->nbFields; i++)
+            {
+                if (fc->fields[i].name == "pads"sv)
+                {
+                    pads.resize(fc->fields[i].length);
+                    auto const* padsPtr = static_cast<int32_t const*>(fc->fields[i].data);
+                    std::copy_n(padsPtr, fc->fields[i].length, pads.data());
+                }
+            }
+
+            return new CircPadPlugin(pads);
+        }
+        catch (std::exception const& e)
+        {
+            caughtError(e);
+        }
+        return nullptr;
+    }
+
+    void setPluginNamespace(char const* libNamespace) noexcept
+    {
+        mNamespace = libNamespace;
+    }
+
+    char const* getPluginNamespace() const noexcept override
+    {
+        return mNamespace.c_str();
+    }
+
+private:
+    PluginFieldCollection mFC;
+    std::vector<PluginField> mPluginAttributes;
+    std::string mNamespace;
+};
+
+REGISTER_TENSORRT_PLUGIN(CircPadPluginCreator);

samples/python/python_plugin/requirements.txt

@@ -0,0 +1,17 @@
+cuda-python==12.9.0
+cupy-cuda12x
+numba
+numba-cuda[cu12]
+triton; platform_system != "Windows"
+torch
+--extra-index-url https://pypi.ngc.nvidia.com
+polygraphy
+colored
+numpy==1.26.4
+onnx==1.18.0; platform_system == "Windows"
+--extra-index-url https://pypi.ngc.nvidia.com
+onnx-graphsurgeon
+pywin32; platform_system == "Windows"
+pyyaml==6.0.3
+requests==2.32.4
+tqdm==4.66.4

samples/python/quickly_deployable_plugins/README.md

@@ -0,0 +1,352 @@
+# Quickly Deployable TensorRT Python Plugins [Experimental]
+
+This is a sample to showcase quickly deployable Python-based plugin definitions (QDPs) in TensorRT (TRT). QDPs are able to support a large majority of use cases for adding custom operators to TRT, and will be the recommended option when it becomes a stable feature in 10.9.
+
+This sample contains several mini-samples that demonstrate a few common use cases.
+
+# Contents
+- [Introduction](#introduction)
+- [Setting up the environment](#setting-up-the-environment)
+- [Implementing a Quickly Deployable Python (QDP) Plugin](#implementing-a-quickly-deployable-python-qdp-plugin)
+- [A Simple Plugin: Elementwise-Add](#a-simple-plugin-elementwise-add)
+- [Implementing in-place custom ops with I/O aliasing](#implementing-in-place-custom-ops-with-io-aliasing)
+- [An Op with data-dependent output shapes: Non-zero](#an-op-with-data-dependent-output-shapes-non-zero)
+- [Using multiple tactics and ONNX: Cirular padding](#using-multiple-tactics-and-onnx-cirular-padding)
+- [Providing an Ahead-of-Time (AOT) implementation for Circular padding](#providing-an-ahead-of-time-aot-implementation-for-circular-padding)
+- [Additional Resources](#additional-resources)
+- [License](#license)
+- [Changelog](#changelog)
+- [Known issues](#known-issues)
+
+
+# Introduction
+
+While the regular TRT plugin interfaces are powerful in the flexibility and tunability they provide, for the vast majority of use cases, users will benefit from the simplicity offered by the QDP workflow.
+ - The `tensorrt.plugin` module provides many intuitive APIs that drastically reduces the amount of boilerplate required to implement a plugin
+ - The concept of plugin registration, plugin creators and the plugin registry is abstracted away
+ - The stateless nature of QDPs eliminates the complications of having to comply with a predefined plugin lifecycle
+
+
+# Setting Up The Environment
+
+To build and install the bindings, follow the instructions in `$TRT_OSSPATH/python/README.md`.
+
+Then install the requisite packages
+```bash
+cd $TRT_OSSPATH/samples/python/quickly_deployable_plugins
+pip3 install -r requirements.txt
+```
+
+# Implementing a Quickly Deployable Python (QDP) Plugin
+
+QDP definitions consist of a set of decorated functions that define properties and behaviors of the plugin.
+### `@tensorrt.plugin.register`
+Returns shape and type characteristics of output tensors, and any attributes the plugin needs to function.
+
+### `@tensorrt.plugin.impl`
+Performs the plugin computation. The decorated python function is executed 'just in time', as a python callback during runtime.
+
+### (Optional) `@tensorrt.plugin.aot_impl`
+The decorated function directly returns an 'ahead of time' compiled kernel, along with information required to invoke it at runtime by TRT. This is in contrast with the above `@tensorrt.plugin.impl` in that, the returned kernel is baked into the built TRT engine. This is beneficial, when we need an engine that is fully independent of the python runtime - and hence, can be executed solely in a standard TensorRT C++ runtime (through `trtexec`, for example).
+
+### (Optional) `@tensorrt.plugin.autotune`
+Defines the different data types and formats (tensor layouts) supported by the plugin's IO and any tactics supported by the plugin. Defining this function allows TensorRT to "tune" the plugin during the engine build to find the most performant type/format and tactic combination on the target system.
+
+The specifics of these functions will become clear through the following mini-samples.
+
+# A Simple Plugin: Elementwise-Add
+
+This mini-sample contains an elementwise addition plugin, where the computation is being performed with an OpenAI Triton kernel. Let's first take a look at the `tensorrt.plugin.register` function.
+
+```python
+import tensorrt.plugin as trtp
+
+@trtp.register("sample::elemwise_add_plugin")
+def add_plugin_desc(inp0: trtp.TensorDesc, block_size: int) -> trtp.TensorDesc:
+    return inp0.like()
+```
+
+The argument "sample::elemwise_add_plugin" defines the namespace ("sample") and name ("elemwise_add_plugin") of the plugin. Input arguments to the decorated function (`plugin_desc`) annotated with `trt.plugin.TensorDesc` denote the input tensors; all others are interpreted as plugin attributes (see the [TRT API Reference](https://docs.nvidia.com/deeplearning/tensorrt/api/python_api/infer/tensorrt.plugin/trt_plugin_register.html) for a full list of allowed attribute types). The output signature is a `trt.plugin.TensorDesc` describing the output. `inp0.like()` returns a tensor descriptor with identical shape and type characteristics to `inp0`.
+
+The computation function, decorated with `trt.plugin.impl`, receives `trt.plugin.Tensor`s for each input and output. In contrast to `TensorDesc`s, a `Tensor` references an underlying data buffer, directly accessible through `Tensor.data_ptr`. When working with Torch and OpenAI Triton kernels, it is easier to use `torch.as_tensor()` to zero-copy construct a `torch.Tensor` corresponding to the `trt.plugin.Tensor`.
+
+This sample also showcases the effect of omitting/defining a `trt.plugin.autotune` function, which must return a list of `trt.plugin.AutoTuneCombination`s. In this case, we define a single combination `AutoTuneCombination("FP32|FP16, FP32|FP16")`; this indicates that the input and output must be either both FP32 or both FP16. See the TRT API Reference for a detailed description of the grammar underlying `AutoTuneCombination`s.
+
+## Running the sample
+
+```bash
+python3 qdp_runner.py add [--autotune] [-v]
+```
+
+`--autotune` simulates having defined a `trt.plugin.autotune` function. Enabling verbose logging (`-v`) is recommended to see the effect of autotuning. It can be observed that the `trt.plugin.impl` function is invoked several times during the engine build process when autotune is enabled. With autotuning turned off, `trt.plugin.impl` is invoked only once (when inference is run after building the engine).
+
+```bash
+$ python3 qdp_runner.py add --autotune -v
+...
+Executing for inp0.dtype=DataType.FLOAT and output[0].dtype=DataType.FLOAT
+Executing for inp0.dtype=DataType.FLOAT and output[0].dtype=DataType.FLOAT
+Executing for inp0.dtype=DataType.FLOAT and output[0].dtype=DataType.FLOAT
+Executing for inp0.dtype=DataType.FLOAT and output[0].dtype=DataType.FLOAT
+Executing for inp0.dtype=DataType.HALF and output[0].dtype=DataType.HALF
+Executing for inp0.dtype=DataType.HALF and output[0].dtype=DataType.HALF
+Executing for inp0.dtype=DataType.HALF and output[0].dtype=DataType.HALF
+Executing for inp0.dtype=DataType.HALF and output[0].dtype=DataType.HALF
+[I] Finished engine building in 1.073 seconds
+Executing for inp0.dtype=DataType.HALF and output[0].dtype=DataType.HALF
+```
+
+# Implementing in-place custom ops with I/O aliasing
+
+In-place computations can be accomplished with TRT plugins via aliased I/O. i.e. An input that needs to be modified in-place can be represented by an input-output pair, where the output is aliased to the input. For example, if in-place addition is needed (instead of the out-of-place addition of the above sample), that can be achieved as below:
+```python
+import tensorrt.plugin as trtp
+
+@trtp.register("sample::elemwise_add_plugin_")
+def add_plugin_desc_(inp0: trtp.TensorDesc) -> trtp.TensorDesc:
+    return inp0.aliased()
+```
+
+Note the use of `trt.plugin.TensorDesc.aliased()` to produce an output `TensorDesc` that is aliased to `inp0`.
+
+To appreciate the effect of aliasing better, this sample adds two in-place add plugins chained together.
+
+## Running the sample
+
+Enabling verbose logging (`-v`) is recommended to see the effect of autotuning, which is always enabled.
+
+```bash
+python3 qdp_runner.py inplace_add [--autotune] [-v]
+```
+
+# An Op with data-dependent output shapes: Non-zero
+
+Non-zero is an operation where the indices of the non-zero elements of the input tensor is found -- it has data-dependent output shapes (DDS). As such, typical shape calculations cannot be done with input shapes.
+
+To handle DDS, the extent of each data-dependent output dimension must be expressed in terms of a *_size tensor_*, which is a scalar that communicates to TRT an upper-bound and an autotune value for that dimension, in terms of the input shapes. The TRT engine build may be optimized for the autotune value, but the extent of that dimension may stretch up to the upper-bound at runtime.
+
+In this sample, we consider a 2D input tensor `inp0`; the output will be an $N x 2$ tensor (a set of $N$ 2D indices), where $N$ is the number of non-zero indices. At maximum, all elements could be non-zero, and so the upper-bound could be expressed as `upper_bound = inp0.shape_expr[0] * inp0.shape_expr[1]`. Note that `trt.plugin.TensorDesc.shape_expr` returns symbolic shape expressions for that tensor. Arithmetic operations on shape expressions are supported through standard Python binary operators (see [TRT Python API reference](https://docs.nvidia.com/deeplearning/tensorrt/api/python_api/infer/tensorrt.plugin/Shape/ShapeExpr.html) for full list of supported operations).
+
+On average, we can expect half of the input to be filled with zero, so a size tensor can be constructed with that as the autotune value:
+```python
+st = trtp.size_tensor(opt = upper_bound // 2, upper_bound = upper_bound)
+```
+
+Now we're ready to construct the output shape. `st.expr()` returns a shape expression for the size tensor, so a tensor descriptor for the output shape can be constructed as `trt.plugin.from_shape_expr((st.expr(), 2), dtype=trt.int32)`. TRT requires that any size tensors also be made outputs of the plugin. Putting things together, we arrive at the following:
+
+```python
+import tensorrt.plugin as trtp
+
+@trtp.register("sample::non_zero_plugin")
+def non_zero_plugin_reg(
+    inp0: trtp.TensorDesc,
+) -> Tuple[trtp.TensorDesc, trtp.TensorDesc]:
+    upper_bound = inp0.shape_expr[0] * inp0.shape_expr[1]
+    st = trtp.size_tensor(upper_bound // 2, upper_bound)
+    return trtp.from_shape_expr((st.expr(), 2), dtype=trt.int32), st
+```
+
+## Running the sample
+
+Enabling verbose logging (`-v`) is recommended to see the effect of autotuning, which is always enabled.
+
+```bash
+python3 qdp_runner.py non_zero [-v]
+```
+
+# Using multiple tactics and ONNX: Cirular padding
+
+This sample contains a circular padding plugin, which is useful for ops like circular convolution. It is equivalent to PyTorch's [torch.nn.CircularPad2d](https://pytorch.org/docs/stable/generated/torch.nn.CircularPad2d.html#torch.nn.CircularPad2d).
+
+Refer [this section about circular padding plugin](https://docs.nvidia.com/deeplearning/tensorrt/latest/_static/python-api/pluginGuide.html#example-circular-padding-plugin) in the python plugin guide for more info.
+
+## ONNX model with a plugin
+
+It is often useful to run an ONNX node with a custom op through a TRT plugin that you have written. To allow the TRT ONNX parser to correctly recognize your plugin as being mapped to an ONNX node, ensure that
+ - The `op` property of the node is exactly the same as your plugin name.
+ - The node contains a string attribute called "plugin_namespace" with the namespace of your plugin.
+
+In this sample, we define a plugin with the ID "sample::circ_pad_plugin", so if using ONNX Graphsurgeon, the custom op node can be constructed as follows:
+
+```python
+import onnx_graphsurgeon as gs
+
+var_x = gs.Variable(name="x", shape=inp_shape, dtype=np.float32)
+var_y = gs.Variable(name="y", dtype=np.float32)
+
+circ_pad_node = gs.Node(
+    name="circ_pad_plugin",
+    op="circ_pad_plugin",
+    inputs=[var_x],
+    outputs=[var_y],
+    attrs={"pads": pads, "plugin_namespace": "sample"},
+)
+```
+
+## Multiple tactics
+
+Sometimes, you may have multiple kernels (or backends) that can be used to perform the computation of the plugin -- these are typically called *_tactics_*. If it cannot be predetermined which of these tactics may perform the fastest, it is possible to let TRT time the plugin for each tactic and determine which one is fastest.
+
+Communicating the availability of multiple tactics can simply be done through the `trt.plugin.autotune` function.
+```python
+import tensorrt.plugin as trtp
+from enum import IntEnum
+
+class Tactic(IntEnum):
+    TORCH = 1
+    TRITON = 2
+
+@trt.plugin.autotune("sample::circ_pad_plugin")
+def circ_pad_plugin_autotune(inp0: trtp.TensorDesc, pads: npt.NDArray[np.int32], outputs: Tuple[trtp.TensorDesc]) -> List[trtp.AutoTuneCombination]:
+    c = trtp.AutoTuneCombination()
+    c.pos([0, 1], "FP32|FP16")
+    c.tactics([int(Tactic.TORCH), int(Tactic.TRITON)])
+    return [c]
+```
+
+Note that we're using another way of constructing a `trt.plugin.AutoTuneCombination` here -- namely, through `pos(...)` to populate the type/format information and `tactics(...)` to specify the tactics. In this sample, we use an OpenAI Triton kernel and `torch.nn.functional.pad` as two methods to compute the circular padding.
+
+Refer [this section](https://docs.nvidia.com/deeplearning/tensorrt/latest/_static/python-api/pluginGuide.html#example-plugins-with-multiple-backends-using-custom-tactics) in the Python plugin guide for more info.
+
+## Loading and running a TRT engine containing a plugin
+
+If you have a TRT engine built with a plugin, executing that engine only requires the plugin definitions for `trt.plugin.register` and `trt.plugin.impl` to be available in the module where the engine is being deserialized (note: the `trt.plugin.autotune` definition is not required to be present).
+
+To simulate the loading of an engine, first run this sample with the `--save_engine` flag, followed by `--artifacts_dir [dir]` with a directory in which you wish the engine to be saved. Then run the sample again with `--load engine` and `--artifacts_dir` set to the same directory.
+
+## Running the sample
+
+```bash
+python3 qdp_runner.py circ_pad [--multi_tactic] [--save_engine] [--load_engine] --mode {onnx,inetdef} [--artifacts_dir ARTIFACTS_DIR]  [-v]
+
+options:
+  --multi_tactic        Enable multiple tactics.
+  --save_engine         Save engine to the artifacts_dir.
+  --load_engine         Load engine from the artifacts_dir. Ignores all other options.
+  --artifacts_dir ARTIFACTS_DIR
+                        Whether to store (or retrieve) artifacts.
+  --mode {onnx,inetdef} Whether to use ONNX parser or INetworkDefinition APIs to construct the network.
+  -v, --verbose         Enable verbose log output.
+```
+
+# Providing an Ahead-of-Time (AOT) implementation for Circular padding
+
+Let's extend the [above sample](#using-multiple-tactics-and-onnx-cirular-padding) by providing an AOT implementation for the same circular padding operation.
+Instead of specifying the OpenAI Triton Kernel callback to TRT through `@trt.plugin.impl`, we can directly
+compile the kernel ahead of time, and provide that to TRT under `@trt.plugin.aot_impl`.
+
+Refer [this section](https://docs.nvidia.com/deeplearning/tensorrt/latest/_static/python-api/pluginGuide.html#providing-an-ahead-of-time-aot-implementation) in the Python plugin guide for more info.
+
+## ONNX model with an AOT plugin
+
+The same rules apply as mentioned in the above [ONNX model with a plugin](#onnx-model-with-a-plugin) section.
+In addition to that, if the plugin has an AOT implementation that we'd like to use, we can modify the ONNX node to communicate that to the TRT ONNX parser.
+This should be done by adding a bool attribute called "aot" to the ONNX node, and setting it to True.
+Note that, this is on top of making sure that the ONNX node has the appropriate `op` property and "plugin_namespace" attribute as mentioned [previously](#onnx-model-with-a-plugin).
+
+Therefore, using ONNX Graphsurgeon, the custom op node that uses the AOT implementation of "sample::circ_pad_plugin" can be constructed similarly:
+
+```python
+import onnx_graphsurgeon as gs
+
+var_x = gs.Variable(name="x", shape=inp_shape, dtype=np.float32)
+var_y = gs.Variable(name="y", dtype=np.float32)
+
+circ_pad_aot_node = gs.Node(
+    name="circ_pad_plugin_aot",
+    op="circ_pad_plugin",
+    inputs=[var_x],
+    outputs=[var_y],
+    attrs={"pads": pads, "plugin_namespace": "sample", "aot": True},
+)
+```
+
+## Loading and running a TRT engine containing an AOT plugin
+
+If you have a TRT engine built with an AOT plugin, the plugin computation is already part of the engine. Therefore, it does not require any Python modules or definitions to be present at runtime. This means that the engine can be executed on the standard TRT runtime, as part of any tool that is capable of deserializing and running the engine (like [trtexec](../../trtexec/README.md).
+
+To simulate the loading of an engine, first run this sample with the `--save_engine` flag, followed by `--artifacts_dir [dir]` with a directory in which you wish the engine to be saved. Then run the sample again with `--load engine` and `--artifacts_dir` set to the same directory.
+
+## Running the sample
+
+```bash
+python3 qdp_runner.py circ_pad [--multi_tactic] [--aot] [--save_engine] [--load_engine] --mode {onnx,inetdef} [--artifacts_dir ARTIFACTS_DIR]  [-v]
+
+options:
+  --multi_tactic        Enable multiple tactics. Combined with --aot, the advertised tactics are AOT-compiled.
+  --save_engine         Save engine to the artifacts_dir.
+  --load_engine         Load engine from the artifacts_dir. Ignores all other options.
+  --artifacts_dir ARTIFACTS_DIR
+                        Whether to store (or retrieve) artifacts.
+  --mode {onnx,inetdef} Whether to use ONNX parser or INetworkDefinition APIs to construct the network.
+  --aot                 Use the AOT implementation of the plugin.
+  -v, --verbose         Enable verbose log output.
+```
+
+## Combining AOT with multiple tactics
+
+The AOT path can advertise more than one tactic. Each tactic must be a precompiled GPU kernel, so a Torch-dispatched tactic (like the `Tactic.TORCH` variant used in the JIT multi-tactic example above) cannot participate. In this sample we instead expose two variants of the same Triton kernel that differ only in `BLOCK_SIZE`. TRT times both PTX blobs at build time and bakes the winner into the engine.
+
+The autotune declaration lists both tactic IDs:
+
+```python
+class Tactic(IntEnum):
+    BLOCK_256 = 1
+    BLOCK_1024 = 2
+
+@trt.plugin.autotune("sample::circ_pad_plugin")
+def circ_pad_plugin_autotune(inp0, outputs):
+    c = trtp.AutoTuneCombination()
+    c.pos([0, 1], "FP32|FP16")
+    c.tactics([int(Tactic.BLOCK_256), int(Tactic.BLOCK_1024)])
+    return [c]
+```
+
+The AOT impl switches on the `tactic` argument, compiles the Triton kernel with the matching `BLOCK_SIZE` constexpr, and returns the corresponding PTX plus a `KernelLaunchParams` whose `grid_x` reflects that block size:
+
+```python
+@trt.plugin.aot_impl("sample::circ_pad_plugin")
+def circ_pad_plugin_aot_impl(inp0, pads, outputs, tactic):
+    block_size = {1: 256, 2: 1024}[tactic]
+    src = triton.compiler.ASTSource(fn=circ_pad_kernel, signature=..., constexprs={"BLOCK_SIZE": block_size})
+    compiled = triton.compile(src)
+    launch_params = trtp.KernelLaunchParams()
+    launch_params.grid_x = trtp.cdiv(outputs[0].shape_expr.numel(), block_size)
+    launch_params.block_x = compiled.metadata.num_warps * 32
+    launch_params.shared_mem = compiled.metadata.shared
+    extra_args = trtp.SymIntExprs.from_tuple(...)  # symbolic int32 kernel args. see source for more details.
+    return compiled.metadata.name.encode(), compiled.asm["ptx"].encode(), launch_params, extra_args
+```
+
+Run it with both flags:
+
+```bash
+python3 qdp_runner.py circ_pad --mode inetdef --multi_tactic --aot -v
+```
+
+Verbose logs show TRT timing both tactics during engine build, then a single winner is serialized into the engine.
+
+# Additional Resources
+
+**Python Plugin Guide**
+- [pluginGuide.md](../../../documentation/python/pluginGuide.md)
+
+**`tensorrt.plugin` API reference**
+- [`tensorrt.plugin` module API reference](https://docs.nvidia.com/deeplearning/tensorrt/api/python_api/infer/tensorrt.plugin/index.html)
+
+**Developer Guide**
+- [Extending TensorRT with Custom Layers](https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#extending)
+
+# License
+
+For terms and conditions for use, reproduction, and distribution, see the [TensorRT Software License Agreement](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sla/index.html) documentation.
+
+# Changelog
+- May 2026: Added multi-tactic AOT subsection for circular padding.
+- October 2025: Migrate to strongly typed APIs.
+- August 2025: Removed support for Python versions < 3.10.
+- December 2024: Added section on AOT Plugins, added contents section
+- October 2024: Initial release of this sample
+
+# Known issues
+
+There are no known issues in this sample

samples/python/quickly_deployable_plugins/oait_kernels.py

@@ -0,0 +1,78 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025 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 triton
+import triton.language as tl
+
+@triton.jit
+def add_kernel(x_ptr, y_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
+    pid = tl.program_id(0)
+    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < n_elements
+    x = tl.load(x_ptr + offsets, mask=mask)
+    tl.store(y_ptr + offsets, x + 1, mask=mask)
+
+
+@triton.jit
+def circ_pad_kernel(
+    # input tensor
+    X,
+    # extra scalar args in between input and output tensors
+    # for kernel signature to be compatible with AOT plugin impl
+    all_pads_0,
+    all_pads_2,
+    all_pads_4,
+    all_pads_6,
+    orig_dims_0,
+    orig_dims_1,
+    orig_dims_2,
+    orig_dims_3,
+    Y_shape_1,
+    Y_shape_2,
+    Y_shape_3,
+    X_len,
+    Y_len,
+    # output tensor
+    Y,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    i = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+
+    mask_y = i < Y_len
+
+    i3 = i % Y_shape_3
+    i2 = (i // Y_shape_3) % Y_shape_2
+    i1 = (i // Y_shape_3 // Y_shape_2) % Y_shape_1
+    i0 = i // Y_shape_3 // Y_shape_2 // Y_shape_1
+
+    j0 = (i0 - all_pads_0 + orig_dims_0) % orig_dims_0
+    j1 = (i1 - all_pads_2 + orig_dims_1) % orig_dims_1
+    j2 = (i2 - all_pads_4 + orig_dims_2) % orig_dims_2
+    j3 = (i3 - all_pads_6 + orig_dims_3) % orig_dims_3
+
+    load_idx = (
+        orig_dims_3 * orig_dims_2 * orig_dims_1 * j0
+        + orig_dims_3 * orig_dims_2 * j1
+        + orig_dims_3 * j2
+        + j3
+    )
+    mask_x = load_idx < X_len
+
+    x = tl.load(X + load_idx, mask=mask_x)
+
+    tl.store(Y + i, x, mask=mask_y)

samples/python/quickly_deployable_plugins/qdp_defs.py

@@ -0,0 +1,386 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 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 tensorrt as trt
+import torch
+import numpy as np
+
+from typing import Tuple, List, Union
+
+import tensorrt.plugin as trtp
+import numpy.typing as npt
+
+import logging
+
+logging.basicConfig(level=logging.INFO)
+logging.getLogger("QuicklyDeployablePlugins").setLevel(logging.INFO)
+
+########## Elemwise-add plugin definition ##########
+
+
+@trtp.register("sample::elemwise_add_plugin")
+def add_plugin_desc(inp0: trtp.TensorDesc, block_size: int) -> trtp.TensorDesc:
+    return inp0.like()
+
+
+# Helper to simulate defining/omitting an autotune definition for the plugin
+def register_autotune():
+    # Type annotations can be omitted for autotune and impl definitions, but will be checked for consistency if added
+    @trtp.autotune("sample::elemwise_add_plugin")
+    def add_plugin_autotune(
+        inp0: trtp.TensorDesc, outputs: Tuple[trtp.TensorDesc]
+    ) -> List[trtp.AutoTuneCombination]:
+        return [trtp.AutoTuneCombination("FP32|FP16, FP32|FP16")]
+
+
+@trtp.impl("sample::elemwise_add_plugin")
+def add_plugin_impl(
+    inp0: trtp.Tensor, block_size: int, outputs: Tuple[trtp.Tensor], stream: int
+) -> None:
+
+    log = logging.getLogger("QuicklyDeployablePlugins")
+    log.debug(
+        f"Executing for inp0: dtype={inp0.dtype},format={inp0.format} and output[0]: dtype={outputs[0].dtype},format={outputs[0].format}"
+    )
+
+    n = inp0.numel()
+
+    with torch.cuda.stream(torch.cuda.ExternalStream(stream)):
+        inp0_t = torch.as_tensor(inp0, device="cuda")
+        out_t = torch.as_tensor(outputs[0], device="cuda")
+
+        import triton
+        from oait_kernels import add_kernel
+
+        add_kernel[(triton.cdiv(n, block_size),)](inp0_t, out_t, n, BLOCK_SIZE=block_size)
+
+
+########## In-place elemwise-add plugin definition ##########
+
+
+@trtp.register("sample::elemwise_add_plugin_")
+def add_plugin_desc_(inp0: trtp.TensorDesc, delta: int) -> trtp.TensorDesc:
+    return inp0.aliased()
+
+
+@trtp.autotune("sample::elemwise_add_plugin_")
+def add_plugin_autotune_(inp0, outputs) -> List[trtp.AutoTuneCombination]:
+    return [
+        trtp.AutoTuneCombination("FP32, FP32", "LINEAR*HWC"),
+        trtp.AutoTuneCombination("FP32|FP16, FP32|FP16", "LINEAR"),
+    ]
+
+
+@trtp.impl("sample::elemwise_add_plugin_")
+def add_plugin_impl_(inp0, delta: int, outputs, stream) -> None:
+
+    log = logging.getLogger("QuicklyDeployablePlugins")
+    log.debug(
+        f"Executing for inp0: dtype={inp0.dtype},format={inp0.format} and output[0]: dtype={outputs[0].dtype},format={outputs[0].format}"
+    )
+
+    with torch.cuda.stream(torch.cuda.ExternalStream(stream)):
+        inp0_t = torch.as_tensor(inp0, device="cuda")
+        inp0_t.add_(delta)
+
+
+########## Non-zero plugin (DDS) ##########
+
+
+@trtp.register("sample::non_zero_plugin")
+def non_zero_plugin_reg(
+    inp0: trtp.TensorDesc,
+) -> Tuple[trtp.TensorDesc, trtp.TensorDesc]:
+    upper_bound = inp0.shape_expr[0] * inp0.shape_expr[1]
+    st = trtp.size_tensor(upper_bound // 2, upper_bound)
+    st.dtype = trt.int64
+    return trtp.from_shape_expr((st.expr(), 2), dtype=trt.int32), st
+
+
+@trtp.autotune("sample::non_zero_plugin")
+def non_zero_plugin_autotune(inp0, outputs) -> List[trtp.AutoTuneCombination]:
+    return [trtp.AutoTuneCombination("FP32|FP16, INT32, INT64")]
+
+
+@trtp.impl("sample::non_zero_plugin")
+def non_zero_plugin_impl(inp0, outputs, stream) -> None:
+
+    log = logging.getLogger("QuicklyDeployablePlugins")
+    log.debug(
+        f"Executing for inp0: dtype={inp0.dtype},format={inp0.format} and output[0]: dtype={outputs[0].dtype},format={outputs[0].format}"
+    )
+
+    with torch.cuda.stream(torch.cuda.ExternalStream(stream)):
+        inp0_t = torch.as_tensor(inp0, device="cuda")
+        out_1 = torch.as_tensor(outputs[1], device="cuda").reshape((-1,))
+
+        out = torch.nonzero(inp0_t)
+
+        out0 = torch.as_tensor(outputs[0].aliased(out.shape), device="cuda")
+        out0.copy_(out)
+        out_1.copy_(torch.Tensor([out.shape[0]]))
+
+
+########## Circular padding plugin ########
+
+
+@trtp.register("sample::circ_pad_plugin")
+def circ_pad_plugin_desc(
+    inp0: trtp.TensorDesc, pads: npt.NDArray[np.int32]
+) -> trtp.TensorDesc:
+    ndim = inp0.ndim
+    out_desc = inp0.like()
+
+    for i in range(np.size(pads) // 2):
+        out_desc.shape_expr[ndim - i - 1] += int(pads[i * 2] + pads[i * 2 + 1])
+
+    return out_desc
+
+
+# Helper to define a multi-tactic implementation of the plugin
+def enable_multi_tactic_circ_pad():
+
+    from enum import IntEnum
+
+    class Tactic(IntEnum):
+        TORCH = 1
+        TRITON = 2
+
+    @trtp.autotune("sample::circ_pad_plugin")
+    def circ_pad_plugin_autotune(
+        inp0: trtp.TensorDesc,
+        outputs: Tuple[trtp.TensorDesc],
+    ) -> List[trtp.AutoTuneCombination]:
+        c = trtp.AutoTuneCombination()
+        c.pos([0, 1], "FP32|FP16")
+        c.tactics([int(Tactic.TORCH), int(Tactic.TRITON)])
+        return [c]
+
+    @trtp.impl("sample::circ_pad_plugin")
+    def circ_pad_plugin_impl(
+        inp0: trtp.Tensor,
+        pads: npt.NDArray[np.int32],
+        outputs: Tuple[trtp.Tensor],
+        stream: int,
+        tactic: int,
+    ) -> None:
+
+        log = logging.getLogger("QuicklyDeployablePlugins")
+        log.debug(
+            f"Executing for inp0: dtype={inp0.dtype},format={inp0.format} and output[0]: dtype={outputs[0].dtype},format={outputs[0].format}"
+        )
+
+        with torch.cuda.stream(torch.cuda.ExternalStream(stream)):
+            inp_t = torch.as_tensor(inp0, device="cuda")
+            out_t = torch.as_tensor(outputs[0], device="cuda")
+
+            if tactic == Tactic.TORCH:
+                out = torch.nn.functional.pad(inp_t, pads.tolist(), mode="circular")
+                out_t.copy_(out)
+            elif tactic == Tactic.TRITON:
+                N = inp0.ndim
+                all_pads = np.zeros((N * 2,), dtype=np.int32)
+                out_dims = trtp.Shape(tuple(inp0.shape))
+
+                for i in range(np.size(pads) // 2):
+                    out_dims[N - i - 1] += pads[i * 2] + pads[i * 2 + 1]
+                    all_pads[N * 2 - 2 * i - 2] = pads[i * 2]
+                    all_pads[N * 2 - 2 * i - 1] = pads[i * 2 + 1]
+
+                all_pads = all_pads.tolist()
+
+                block_size = 256
+                num_blocks = tuple(
+                    [int((np.prod(out_dims) + block_size - 1) // block_size)]
+                )
+
+                from oait_kernels import circ_pad
+
+                circ_pad[num_blocks](
+                    inp_t,
+                    all_pads[0],
+                    all_pads[2],
+                    all_pads[4],
+                    all_pads[6],
+                    inp0.shape[0],
+                    inp0.shape[1],
+                    inp0.shape[2],
+                    inp0.shape[3],
+                    int(out_dims[1]),
+                    int(out_dims[2]),
+                    int(out_dims[3]),
+                    inp0.numel(),
+                    out_dims.numel(),
+                    out_t,
+                    BLOCK_SIZE=block_size,
+                )
+
+
+# Shared AOT compilation body for the circ_pad plugin: build SymInt args,
+# compile the Triton kernel with the given BLOCK_SIZE, and pack launch params.
+def _compile_circ_pad_aot(
+    inp0: trtp.TensorDesc,
+    pads: npt.NDArray[np.int32],
+    outputs: Tuple[trtp.TensorDesc],
+    block_size: int,
+) -> Tuple[Union[str, bytes], Union[str, bytes], trtp.KernelLaunchParams, trtp.SymExprs]:
+
+    N = inp0.ndim
+    all_pads = np.zeros((N * 2,), dtype=np.int32)
+    inp_dims = inp0.shape_expr
+    out_dims = outputs[0].shape_expr
+
+    for i in range(np.size(pads) // 2):
+        all_pads[N * 2 - 2 * i - 2] = pads[i * 2]
+        all_pads[N * 2 - 2 * i - 1] = pads[i * 2 + 1]
+
+    all_pads = all_pads.tolist()
+
+    # Representing all int32-scalar-kernel-inputs as symbolic expressions.
+    # These inputs are either constants or derivatives of input/output shapes (that may be dynamic).
+    # The symbolic expressions are resolved after the full shape context becomes available at runtime.
+    extra_args = trtp.SymIntExprs.from_tuple(
+        [
+            trtp.SymInt32(e)
+            for e in [
+                all_pads[0],
+                all_pads[2],
+                all_pads[4],
+                all_pads[6],
+                inp_dims[0],
+                inp_dims[1],
+                inp_dims[2],
+                inp_dims[3],
+                out_dims[1],
+                out_dims[2],
+                out_dims[3],
+                inp_dims.numel(),
+                out_dims.numel(),
+            ]
+        ]
+    )
+
+    type_str = "fp32" if inp0.dtype == trt.float32 else "fp16"
+
+    from oait_kernels import circ_pad_kernel
+    import triton
+
+    src = triton.compiler.ASTSource(
+        fn=circ_pad_kernel,
+        signature={
+            "X": f"*{type_str}",
+            "all_pads_0": "i32",
+            "all_pads_2": "i32",
+            "all_pads_4": "i32",
+            "all_pads_6": "i32",
+            "orig_dims_0": "i32",
+            "orig_dims_1": "i32",
+            "orig_dims_2": "i32",
+            "orig_dims_3": "i32",
+            "Y_shape_1": "i32",
+            "Y_shape_2": "i32",
+            "Y_shape_3": "i32",
+            "X_len": "i32",
+            "Y_len": "i32",
+            "Y": f"*{type_str}",
+        },
+        constexprs={"BLOCK_SIZE": block_size},
+    )
+
+    compiled_kernel = triton.compile(src)
+    launch_params = trtp.KernelLaunchParams()
+    launch_params.grid_x = trtp.cdiv(out_dims.numel(), block_size)
+    launch_params.block_x = compiled_kernel.metadata.num_warps * 32
+    launch_params.shared_mem = compiled_kernel.metadata.shared
+
+    return (
+        compiled_kernel.metadata.name.encode(),
+        compiled_kernel.asm["ptx"].encode(),
+        launch_params,
+        extra_args,
+    )
+
+
+# Helper to define a single tactic implementation of the plugin
+def enable_single_tactic_circ_pad():
+    @trtp.autotune("sample::circ_pad_plugin")
+    def circ_pad_plugin_autotune(
+        inp0: trtp.TensorDesc,
+        outputs: Tuple[trtp.TensorDesc],
+    ) -> List[trtp.AutoTuneCombination]:
+
+        return [trtp.AutoTuneCombination("FP32|FP16, FP32|FP16")]
+
+    @trtp.impl("sample::circ_pad_plugin")
+    def circ_pad_plugin_impl(
+        inp0: trtp.Tensor,
+        pads: npt.NDArray[np.int32],
+        outputs: Tuple[trtp.Tensor],
+        stream: int,
+    ) -> None:
+        with torch.cuda.stream(torch.cuda.ExternalStream(stream)):
+            inp_t = torch.as_tensor(inp0, device="cuda")
+            out_t = torch.as_tensor(outputs[0], device="cuda")
+
+            out = torch.nn.functional.pad(inp_t, pads.tolist(), mode="circular")
+            out_t.copy_(out)
+
+    @trtp.aot_impl("sample::circ_pad_plugin")
+    def circ_pad_plugin_aot_impl(
+        inp0: trtp.TensorDesc, pads: npt.NDArray[np.int32], outputs: Tuple[trtp.TensorDesc], tactic: int
+    ) -> Tuple[Union[str, bytes], Union[str, bytes], trtp.KernelLaunchParams, trtp.SymExprs]:
+        return _compile_circ_pad_aot(inp0, pads, outputs, block_size=256)
+
+
+# Helper to define a multi-tactic AOT implementation of the plugin.
+# Each tactic precompiles the same Triton kernel with a different BLOCK_SIZE,
+# so TRT times two PTX variants at build time and bakes the winner into the engine.
+def enable_multi_tactic_aot_circ_pad():
+
+    from enum import IntEnum
+
+    class Tactic(IntEnum):
+        BLOCK_256 = 1
+        BLOCK_1024 = 2
+
+    block_size_by_tactic = {
+        int(Tactic.BLOCK_256): 256,
+        int(Tactic.BLOCK_1024): 1024,
+    }
+
+    @trtp.autotune("sample::circ_pad_plugin")
+    def circ_pad_plugin_autotune(
+        inp0: trtp.TensorDesc,
+        outputs: Tuple[trtp.TensorDesc],
+    ) -> List[trtp.AutoTuneCombination]:
+        c = trtp.AutoTuneCombination()
+        c.pos([0, 1], "FP32|FP16")
+        c.tactics([int(Tactic.BLOCK_256), int(Tactic.BLOCK_1024)])
+        return [c]
+
+    @trtp.aot_impl("sample::circ_pad_plugin")
+    def circ_pad_plugin_aot_impl(
+        inp0: trtp.TensorDesc,
+        pads: npt.NDArray[np.int32],
+        outputs: Tuple[trtp.TensorDesc],
+        tactic: int,
+    ) -> Tuple[Union[str, bytes], Union[str, bytes], trtp.KernelLaunchParams, trtp.SymExprs]:
+        block_size = block_size_by_tactic[tactic]
+        logging.getLogger("QuicklyDeployablePlugins").debug(
+            f"aot_impl invoked: tactic={tactic} -> BLOCK_SIZE={block_size}"
+        )
+        return _compile_circ_pad_aot(inp0, pads, outputs, block_size=block_size)

samples/python/quickly_deployable_plugins/qdp_runner.py

@@ -0,0 +1,363 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 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 tensorrt as trt
+import torch
+import numpy as np
+
+from polygraphy.backend.trt import (
+    CreateConfig,
+    TrtRunner,
+    create_network,
+    engine_from_network,
+    network_from_onnx_path,
+    bytes_from_engine,
+    engine_from_bytes,
+)
+
+from polygraphy.backend.common import bytes_from_path
+from polygraphy import cuda
+
+import onnx_graphsurgeon as gs
+import onnx
+import os
+import argparse
+
+import tensorrt.plugin as trtp
+
+import qdp_defs
+import logging
+
+def run_add(enable_autotune=False):
+
+    if enable_autotune:
+        qdp_defs.register_autotune()
+
+    BLOCK_SIZE = 256
+
+    builder, network = create_network(strongly_typed=True)
+    x = torch.randint(10, (10, 3, 32, 32), dtype=torch.float32, device="cuda")
+
+    # Populate network
+    i_x = network.add_input(name="x", dtype=trt.DataType.FLOAT, shape=x.shape)
+
+    out = network.add_plugin(
+        trtp.op.sample.elemwise_add_plugin(i_x, block_size=BLOCK_SIZE)
+    )
+    out.get_output(0).name = "y"
+    network.mark_output(tensor=out.get_output(0))
+
+    builder.create_builder_config()
+
+    engine = engine_from_network(
+        (builder, network),
+        CreateConfig(),
+    )
+
+    with TrtRunner(engine, "trt_runner") as runner:
+        outputs = runner.infer(
+            {
+                "x": x,
+            },
+            copy_outputs_to_host=False,
+        )
+
+    if torch.allclose(x + 1, outputs["y"]):
+        print("Inference result is correct!")
+    else:
+        print("Inference result is incorrect!")
+
+
+def run_inplace_add():
+    builder, network = create_network(strongly_typed=True)
+    x = torch.ones((10, 3, 32, 32), dtype=torch.float32, device="cuda")
+
+    x_clone = x.clone()
+
+    i_x = network.add_input(name="x", dtype=trt.DataType.FLOAT, shape=x.shape)
+
+    # Amounts to elementwise-add in the first and second plugins
+    deltas = (2, 4)
+
+    out0 = network.add_plugin(trtp.op.sample.elemwise_add_plugin_(i_x, delta=deltas[0]))
+    out1 = network.add_plugin(
+        trtp.op.sample.elemwise_add_plugin_(out0.get_output(0), delta=deltas[1])
+    )
+    out1.get_output(0).name = "y"
+    network.mark_output(tensor=out1.get_output(0))
+
+    builder.create_builder_config()
+
+    # Enable preview feature for aliasing plugin I/O
+    config = CreateConfig(
+        preview_features=[trt.PreviewFeature.ALIASED_PLUGIN_IO_10_03]
+    )
+
+    engine = engine_from_network(
+        (builder, network),
+        config,
+    )
+
+    context = engine.create_execution_context()
+
+    stream = cuda.Stream()
+
+    context.set_tensor_address("x", x.data_ptr())
+    context.set_tensor_address("y", x.data_ptr())
+    context.execute_async_v3(stream.ptr)
+    stream.synchronize()
+
+    if torch.allclose(x, x_clone + sum(deltas), atol=1e-2):
+        print("Inference result is correct!")
+    else:
+        print("Inference result is incorrect!")
+        print(x[0][0][0][:10])
+        print(x_clone[0][0][0][:10])
+
+
+def run_non_zero():
+    builder, network = create_network(strongly_typed=True)
+    inp_shape = (128, 128)
+
+    X = np.random.normal(size=inp_shape).astype(trt.nptype(trt.DataType.FLOAT))
+
+    # Zero out some random indices
+    indices = np.random.choice(
+        np.prod(inp_shape),
+        replace=False,
+        size=np.random.randint(0, np.prod(inp_shape) + 1),
+    )
+    X[np.unravel_index(indices, inp_shape)] = 0
+
+    # Populate network
+    i_x = network.add_input(name="X", dtype=trt.DataType.FLOAT, shape=inp_shape)
+
+    out = network.add_plugin(trtp.op.sample.non_zero_plugin(i_x))
+    out.get_output(0).name = "Y"
+    network.mark_output(tensor=out.get_output(0))
+
+    builder.create_builder_config()
+
+    engine = engine_from_network(
+        (builder, network),
+        config=CreateConfig(),
+    )
+
+    Y_ref = np.transpose(np.nonzero(X))
+
+    with TrtRunner(engine, "trt_runner") as runner:
+        outputs = runner.infer({"X": X})
+        Y = outputs["Y"]
+        Y = Y[np.lexsort(np.fliplr(Y).T)]
+
+    if np.allclose(Y, Y_ref, atol=1e-3):
+        print("Inference result is correct!")
+    else:
+        print("Inference result is incorrect!")
+
+
+def check_artifacts_dir_exists(artifacts_dir):
+    if not os.path.exists(artifacts_dir):
+        raise ValueError(f"artifacts_dir '{artifacts_dir}' does not exist")
+
+
+def run_circ_pad(
+    enable_multi_tactic=False, mode="onnx", artifacts_dir=None, save_or_load_engine=None, aot=False
+):
+
+    if enable_multi_tactic and aot:
+        qdp_defs.enable_multi_tactic_aot_circ_pad()
+    elif enable_multi_tactic:
+        qdp_defs.enable_multi_tactic_circ_pad()
+    else:
+        qdp_defs.enable_single_tactic_circ_pad()
+
+    inp_shape = (10, 3, 32, 32)
+    x = np.random.normal(size=inp_shape).astype(trt.nptype(trt.DataType.FLOAT))
+
+    pads = np.array((1, 1, 1, 1), dtype=np.int32)
+
+    if save_or_load_engine is not None and save_or_load_engine is False:
+        check_artifacts_dir_exists(artifacts_dir)
+        engine_path = os.path.join(artifacts_dir, "circ_pad.engine")
+        engine = engine_from_bytes(bytes_from_path(engine_path))
+    else:
+        if mode == "inetdef":
+            builder, network = create_network(strongly_typed=True)
+            i_x = network.add_input(name="x", dtype=trt.DataType.FLOAT, shape=x.shape)
+            out = network.add_plugin(trtp.op.sample.circ_pad_plugin(i_x, pads=pads), aot = aot)
+            out.get_output(0).name = "y"
+            network.mark_output(tensor=out.get_output(0))
+
+            engine = engine_from_network(
+                (builder, network),
+                CreateConfig(),
+            )
+        elif mode == "onnx":
+            if artifacts_dir is None:
+                raise ValueError("'artifacts_dir' must be specified in onnx mode")
+
+            check_artifacts_dir_exists(artifacts_dir)
+
+            onnx_path = os.path.join(artifacts_dir, "circ_pad.onnx")
+            var_x = gs.Variable(name="x", shape=inp_shape, dtype=np.float32)
+            var_y = gs.Variable(name="y", dtype=np.float32)
+            circ_pad_node = gs.Node(
+                name="circ_pad_plugin 0",
+                op="circ_pad_plugin",
+                inputs=[var_x],
+                outputs=[var_y],
+                attrs={"pads": pads, "plugin_namespace": "sample", "aot": aot},
+            )
+            graph = gs.Graph(
+                nodes=[circ_pad_node], inputs=[var_x], outputs=[var_y], opset=16
+            )
+            onnx.save(gs.export_onnx(graph), onnx_path)
+
+            engine = engine_from_network(
+                network_from_onnx_path(onnx_path, strongly_typed=True), CreateConfig()
+            )
+        else:
+            raise ValueError(f"Unknown mode {mode}")
+
+        if save_or_load_engine is not None and save_or_load_engine is True:
+            check_artifacts_dir_exists(artifacts_dir)
+            engine_path = os.path.join(artifacts_dir, "circ_pad.engine")
+            with open(engine_path, "wb") as f:
+                f.write(bytes_from_engine(engine))
+
+    Y_ref = np.pad(x, [[0, 0], [0, 0], [pads[0], pads[1]], [pads[2], pads[3]]], "wrap")
+
+    with TrtRunner(engine, "trt_runner") as runner:
+        outputs = runner.infer({"x": x})
+        Y = outputs["y"]
+
+        if np.allclose(Y, Y_ref, atol=1e-2):
+            print("Inference result is correct!")
+        else:
+            print("Inference result is incorrect!")
+
+
+def setup_add_sample(subparsers):
+    subparser = subparsers.add_parser("add", help="'add' sample help")
+    subparser.add_argument("--autotune", action="store_true", help="Enable autotuning")
+    subparser.add_argument("--aot", action="store_true", help="Use the AOT implementation of the plugin")
+    subparser.add_argument(
+        "-v", "--verbose", action="store_true", help="Enable more verbose log output"
+    )
+
+
+def setup_inplace_add_sample(subparsers):
+    subparser = subparsers.add_parser("inplace_add", help="inplace_add sample help")
+    subparser.add_argument(
+        "-v", "--verbose", action="store_true", help="Enable more verbose log output"
+    )
+
+
+def setup_non_zero_sample(subparsers):
+    subparser = subparsers.add_parser("non_zero", help="non_zero sample help")
+    subparser.add_argument(
+        "-v", "--verbose", action="store_true", help="Enable more verbose log output"
+    )
+
+
+def setup_circ_pad_sample(subparsers):
+    subparser = subparsers.add_parser("circ_pad", help="circ_pad sample help.")
+    subparser.add_argument(
+        "--multi_tactic", action="store_true", help="Enable multiple tactics."
+    )
+    subparser.add_argument(
+        "--save_engine", action="store_true", help="Save engine to the artifacts_dir."
+    )
+    subparser.add_argument(
+        "--load_engine",
+        action="store_true",
+        help="Load engine from the artifacts_dir. Ignores all other options.",
+    )
+    subparser.add_argument(
+        "--artifacts_dir",
+        type=str,
+        help="Whether to store (or retrieve) artifacts.",
+    )
+    subparser.add_argument(
+        "--mode",
+        type=str,
+        choices=["onnx", "inetdef"],
+        help="Whether to use ONNX parser or INetworkDefinition APIs to construct the network.",
+    )
+    subparser.add_argument("--aot", action="store_true", help="Use the AOT implementation of the plugin.")
+    subparser.add_argument(
+        "-v", "--verbose", action="store_true", help="Enable verbose log output."
+    )
+
+    return subparser
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+
+    parser = argparse.ArgumentParser(description="Main script help")
+    subparsers = parser.add_subparsers(dest="sample", help="Mode help", required=True)
+
+    setup_add_sample(subparsers)
+    setup_inplace_add_sample(subparsers)
+    circ_pad_subparser = setup_circ_pad_sample(subparsers)
+    setup_non_zero_sample(subparsers)
+
+    args = parser.parse_args()
+
+    if args.verbose:
+        logging.getLogger("QuicklyDeployablePlugins").setLevel(logging.DEBUG)
+
+    if args.sample == "add":
+        run_add(args.autotune)
+    if args.sample == "inplace_add":
+        run_inplace_add()
+    if args.sample == "non_zero":
+        run_non_zero()
+    if args.sample == "circ_pad":
+        if args.mode == "onnx":
+            if args.artifacts_dir is None:
+                parser.error(
+                    "circ_pad: argument --mode: When mode is 'onnx', artifacts_dir is required"
+                )
+
+        save_or_load_engine = None
+
+        if args.load_engine is True:
+            if args.save_engine is True:
+                parser.error(
+                    "circ_pad: save_engine and load_engine cannot be specified at the same time. First save_engine and load_engine separately."
+                )
+            else:
+                if args.multi_tactic is True or args.mode is not None:
+                    print(
+                        "warning circ_pad: when load_engine is specified, all other options except 'artifacts_dir' is ignored."
+                    )
+
+            save_or_load_engine = False
+        else:
+            if args.mode is None:
+                circ_pad_subparser.print_help()
+                parser.error(
+                    "circ_pad: '--mode' option is required."
+                )
+
+        if args.save_engine is True:
+            save_or_load_engine = True
+
+        run_circ_pad(args.multi_tactic, args.mode, args.artifacts_dir, save_or_load_engine, args.aot)

samples/python/quickly_deployable_plugins/requirements.txt

@@ -0,0 +1,12 @@
+triton==3.5.0; (platform_system != "Windows")
+torch
+--extra-index-url https://pypi.ngc.nvidia.com
+polygraphy>=0.50.1
+colored
+numpy==1.26.4
+onnx==1.18.0; platform_system == "Windows"
+--extra-index-url https://pypi.ngc.nvidia.com
+onnx-graphsurgeon
+pyyaml==6.0.3
+requests==2.32.4
+tqdm==4.66.4

samples/python/refactored/1_run_onnx_with_tensorrt/README.md

@@ -0,0 +1,66 @@
+# Run ONNX with TensorRT
+This sample demonstrates:
+
+- Converting a pre-trained [EfficientNet](https://arxiv.org/abs/1905.11946)-B0 ONNX model to a `TensorRT` engine
+- Performing inference with `TensorRT` using Python APIs
+- Comparing inference performance between `ONNX Runtime` and `TensorRT`
+- Proper memory management and resource cleanup in both Python implementations
+
+## Key features demonstrated:
+
+- `TensorRT`'s ONNX parser + ONNX model -> `TensorRT` engine
+- Engine building and serialization
+- Input/output tensor handling
+- Performance profiling
+- Editable timing cache for deterministic engine builds
+- Memory pool optimization with workspace configuration
+
+## Implementation Details
+
+### Memory Management
+- Configures workspace memory pool for running under limited hardware
+
+### Engine Building
+- Supports editable timing cache for deterministic builds
+- Serialization and deserialization of TensorRT engines
+
+### Inference Pipeline
+- Efficient image preprocessing with `PIL` and `NumPy`
+- Supports batch inference
+- Implements proper error handling and resource cleanup
+- Provides performance comparison between `ONNX Runtime` and `TensorRT`
+- Performs inference on a real-world image
+
+## CLI Tools
+Users can run their onnx model and generate the engine with similar functionality using `trtexec`:
+
+```bash
+# Basic conversion with performance profiling
+trtexec --onnx=efficientnet-b0.onnx \
+        --saveEngine=efficientnet-b0_trtexec.plan \
+        --dumpProfile \
+        --iterations=100 \
+        --avgRuns=100 \
+        --workspace=1024 \
+        --batch=1
+```
+
+Key options explained:
+- `--onnx`: Input ONNX model
+- `--saveEngine`: Output TensorRT engine
+- `--dumpProfile`: Performance profiling
+- `--iterations`: Number of inference iterations
+- `--avgRuns`: Number of runs to average for timing
+- `--workspace`: Workspace size in MB (1024MB = 1GB)
+- `--batch`: Batch size for inference
+
+## Additional Resources
+
+- [TensorRT Documentation](https://docs.nvidia.com/deeplearning/tensorrt/latest/index.html)
+- [ONNX Documentation](https://onnx.ai/)
+- [EfficientNet Paper](https://arxiv.org/abs/1905.11946)
+
+# Changelog
+
+August 2025
+Removed support for Python versions < 3.10.

samples/python/refactored/1_run_onnx_with_tensorrt/main.ipynb

@@ -0,0 +1,589 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "\n",
+    "SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.SPDX-License-Identifier: Apache-2.0\n",
+    "\n",
+    "Licensed under the Apache License, Version 2.0 (the \"License\"); you may not use\n",
+    "\n",
+    "this file except in compliance with the License. You may obtain a copy of the License at\n",
+    "\n",
+    "\n",
+    "\n",
+    "http://www.apache.org/licenses/LICENSE-2.0\n",
+    "\n",
+    "\n",
+    "\n",
+    "Unless required by applicable law or agreed to in writing, software\n",
+    "\n",
+    "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+    "\n",
+    "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+    "\n",
+    "See the License for the specific language governing permissions and\n",
+    "\n",
+    "limitations under the License.\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Getting Started with TensorRT: Accelerate Your Deep Learning Inference\n",
+    "\n",
+    "Welcome to your first TensorRT tutorial! In this notebook, you'll learn how to:\n",
+    "1. Load a pre-trained EfficientNet model in ONNX format\n",
+    "2. Convert it to a TensorRT engine for faster inference\n",
+    "3. Run inference and see the speedup firsthand\n",
+    "4. Make predictions on real images\n",
+    "\n",
+    "## Understanding ONNX: The Universal Model Format\n",
+    "\n",
+    "ONNX (Open Neural Network Exchange) is a standard format for representing deep learning models. Think of it as a universal language that different deep learning frameworks can understand. Here's why it's important:\n",
+    "\n",
+    "- **Framework Independence**: Models trained in PyTorch, TensorFlow, or other frameworks can be exported to ONNX\n",
+    "- **Interoperability**: ONNX models can be imported into various inference engines and frameworks\n",
+    "- **Production Ready**: ONNX is widely used in production environments for model deployment\n",
+    "\n",
+    "### The ONNX to TensorRT Workflow\n",
+    "\n",
+    "TensorRT is NVIDIA's deep learning inference optimizer that can import models from ONNX. This makes it a powerful tool in your deployment pipeline:\n",
+    "\n",
+    "```\n",
+    "Your Framework (PyTorch/TF/etc.) → ONNX → TensorRT ===> Optimized Inference\n",
+    "```\n",
+    "\n",
+    "This workflow is particularly powerful because:\n",
+    "1. You can train your model in any framework you prefer\n",
+    "2. Export it to ONNX (a one-time conversion)\n",
+    "3. Use TensorRT to optimize it for NVIDIA GPUs\n",
+    "4. Get significant speedup in production\n",
+    "\n",
+    "## Prerequisites\n",
+    "\n",
+    "Before we start, make sure you have:\n",
+    "- NVIDIA GPU with CUDA support\n",
+    "- Python 3.10+ installed\n",
+    "- Basic understanding of deep learning and inference\n",
+    "\n",
+    "Let's begin by installing and importing the required packages:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%pip install tensorrt cuda-python pillow onnxruntime\n",
+    "import tensorrt as trt\n",
+    "from cuda.bindings import runtime as cudart\n",
+    "from PIL import Image\n",
+    "import numpy as np\n",
+    "from pathlib import Path\n",
+    "import time\n",
+    "from typing import Optional, Union, Tuple"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "root = Path.cwd()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# define a function to download files\n",
+    "\n",
+    "import requests\n",
+    "from requests.adapters import HTTPAdapter\n",
+    "from urllib3.util.retry import Retry\n",
+    "\n",
+    "def download_file(url: str, output_path: Union[str, Path]):\n",
+    "    \"\"\"Download a file with retry mechanism.\"\"\"\n",
+    "    session = requests.Session()\n",
+    "    retry = Retry(total=10, backoff_factor=1)\n",
+    "    adapter = HTTPAdapter(max_retries=retry)\n",
+    "    session.mount('http://', adapter)\n",
+    "    session.mount('https://', adapter)\n",
+    "\n",
+    "    response = session.get(url, verify=False, timeout=30)\n",
+    "    output_path.parent.mkdir(parents=True, exist_ok=True)\n",
+    "    with open(output_path, 'wb') as f:\n",
+    "        f.write(response.content)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Step 1: Download a Pre-trained Model\n",
+    "\n",
+    "We'll use EfficientNet-B0, a popular and efficient image classification model, as an example for this sample. \n",
+    "\n",
+    "### Understanding ONNX Model Structure\n",
+    "\n",
+    "An ONNX model contains:\n",
+    "- Model architecture (layers, connections)\n",
+    "- Weights and biases\n",
+    "- Input/output specifications\n",
+    "- Metadata about the model\n",
+    "just like any other model representations. \n",
+    "\n",
+    "This standardized format makes it easy to move models between different frameworks and inference engines."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "download_file(\"https://github.com/onnx/models/raw/refs/heads/main/Computer_Vision/efficientnet_b0_Opset17_timm/efficientnet_b0_Opset17.onnx\", root / \"efficientnet-b0.onnx\")\n",
+    "assert (root / \"efficientnet-b0.onnx\").exists(), \"Model file not found. Please check if the download was successful.\""
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Step 2: Convert ONNX to TensorRT Engine\n",
+    "\n",
+    "This is where the magic happens! We'll convert our ONNX model into a TensorRT engine. The engine is optimized for your specific GPU and will run much faster than the original model.\n",
+    "\n",
+    "### The Conversion Process\n",
+    "\n",
+    "1. **Load ONNX Model**: TensorRT reads the ONNX file and understands the model structure\n",
+    "2. **Optimize**: TensorRT performs several optimizations:\n",
+    "   - Layer fusion\n",
+    "   - Memory optimization\n",
+    "   - Precision calibration\n",
+    "3. **Generate Engine**: Creates a highly optimized inference engine\n",
+    "\n",
+    "The resulting engine is specific to your GPU and will run much faster than the original ONNX model."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "logger = trt.Logger(trt.Logger.WARNING)\n",
+    "builder = trt.Builder(logger)\n",
+    "network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED))\n",
+    "\n",
+    "# Bind the TensorRT network to the parser so that the parser can update the network later accordingly\n",
+    "parser = trt.OnnxParser(network, logger)\n",
+    "\n",
+    "onnx_path = root / \"efficientnet-b0.onnx\"\n",
+    "print(f'Parsing ONNX model at {onnx_path}...')\n",
+    "with open(onnx_path, \"rb\") as model:\n",
+    "    parser.parse(model.read())\n",
+    "print('Parsing ONNX model... done')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Now that we have the TensorRT `INetworkDefinition`, we can start building the engine"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "config = builder.create_builder_config()\n",
+    "\n",
+    "# TensorRT needs memory for layer operations and intermediate activations during inference\n",
+    "# Setting a memory limit helps control resource usage and prevents out-of-memory errors\n",
+    "config.set_memory_pool_limit(\n",
+    "        trt.MemoryPoolType.WORKSPACE, 1 << 30\n",
+    ") # 1GB\n",
+    "\n",
+    "print('Starting to build engine. This might take several minutes depending on the hardware...')\n",
+    "engine = builder.build_serialized_network(network, config)\n",
+    "assert engine is not None, 'Engine build failed'\n",
+    "\n",
+    "engine_path = root / \"efficientnet-b0.plan\"\n",
+    "with open(engine_path, 'wb') as f:\n",
+    "    f.write(engine)\n",
+    "\n",
+    "print(\"TensorRT engine created successfully!\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Optional: Using Editable Timing Cache\n",
+    "\n",
+    "TensorRT engines may vary between builds because kernel selection is based on runtime performance measurements. The hardware state (GPU utilization, temperature, system load) affects which kernels are chosen since kernels might outperform each other under different scenarios. \n",
+    "\n",
+    "To ensure consistent builds, TensorRT provides an editable timing cache that:\n",
+    "- Stores intermediate optimization results\n",
+    "- Enables deterministic engine builds\n",
+    "- Speeds up subsequent builds since they don't need to measure kernel execution time for each op again"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def build_engine_with_cache(onnx_path: Union[str, Path], timing_cache: Optional[trt.ITimingCache]):\n",
+    "    builder = trt.Builder(logger)\n",
+    "    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED))\n",
+    "    parser = trt.OnnxParser(network, logger)\n",
+    "    with open(onnx_path, 'rb') as model:\n",
+    "        parser.parse(model.read())\n",
+    "    config = builder.create_builder_config()\n",
+    "    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)\n",
+    "\n",
+    "    # Enable editable timing cache\n",
+    "    config.set_flag(trt.BuilderFlag.EDITABLE_TIMING_CACHE)\n",
+    "\n",
+    "    # Create timing cache if not provided\n",
+    "    if not timing_cache:\n",
+    "        timing_cache = config.create_timing_cache(bytes())\n",
+    "    config.set_timing_cache(timing_cache, True)\n",
+    "\n",
+    "    # Build engine\n",
+    "    print('Start building engine...')\n",
+    "    tik = time.time()\n",
+    "    engine = builder.build_serialized_network(network, config)\n",
+    "    tok = time.time()\n",
+    "\n",
+    "    print(f'Engine build cost {tok - tik}ms')\n",
+    "    return engine, timing_cache\n",
+    "\n",
+    "# First build (creates cache)\n",
+    "engine1, timing_cache = build_engine_with_cache(onnx_path, None)\n",
+    "print(\"First build completed with cache creation\")\n",
+    "\n",
+    "# Second build (uses cache)\n",
+    "engine2, timing_cache = build_engine_with_cache(onnx_path, timing_cache)\n",
+    "print(\"Second build completed with cache creation\")\n",
+    "\n",
+    "is_identical = np.array_equal(\n",
+    "    np.frombuffer(engine1, dtype=np.uint8),\n",
+    "    np.frombuffer(engine2, dtype=np.uint8))\n",
+    "print(f'Is engine identical: {is_identical}')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Step 3: Run Inference and Compare Performance\n",
+    "\n",
+    "Now let's see the real power of TensorRT! We'll:\n",
+    "1. Run inference with both ONNX and TensorRT\n",
+    "2. Compare their performance\n",
+    "3. See the speedup TensorRT provides\n",
+    "\n",
+    "### Understanding the Performance Difference\n",
+    "\n",
+    "The speedup comes from several optimizations:\n",
+    "- Layer fusion: Combining multiple operations into one\n",
+    "- Memory optimization: Better memory access patterns\n",
+    "- Precision optimization: Using optimal precision for each layer\n",
+    "- CUDA optimization: Direct GPU execution without framework overhead"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def load_and_preprocess_image(image_path: Union[str, Path], input_size: Tuple[int, int] = (224, 224)):\n",
+    "    img = Image.open(image_path)\n",
+    "    img = img.resize(input_size)\n",
+    "    img = np.array(img).astype(np.float32)\n",
+    "    img = img / 255.0  # Normalize from [0, 255] to [0, 1]\n",
+    "    img = np.transpose(img, (2, 0, 1))  # HWC to CHW\n",
+    "    img = np.expand_dims(img, axis=0)  # Add batch dimension\n",
+    "    return img\n",
+    "\n",
+    "def check_cuda_error(error):\n",
+    "    if isinstance(error, tuple):\n",
+    "        error = error[0]\n",
+    "    if error != cudart.cudaError_t.cudaSuccess:\n",
+    "        error_name = cudart.cudaGetErrorName(error)[1]\n",
+    "        error_string = cudart.cudaGetErrorString(error)[1]\n",
+    "        raise RuntimeError(f\"CUDA Error: {error_name} ({error_string})\")\n",
+    "\n",
+    "def run_inference_trt(engine: trt.ICudaEngine, input_data: np.ndarray):\n",
+    "    # Create execution context - this stores the device memory allocations\n",
+    "    # and bindings needed for inference\n",
+    "    context = engine.create_execution_context()\n",
+    "\n",
+    "    # Initialize lists to store input/output information and GPU memory allocations\n",
+    "    inputs = []\n",
+    "    outputs = []\n",
+    "    allocations = []\n",
+    "\n",
+    "    # Iterate through all input/output tensors to set up memory and bindings\n",
+    "    for i in range(engine.num_io_tensors):\n",
+    "        name = engine.get_tensor_name(i)\n",
+    "        # Check if this tensor is an input or output\n",
+    "        is_input = engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT\n",
+    "        # Get tensor datatype and shape information\n",
+    "        dtype = engine.get_tensor_dtype(name)\n",
+    "        shape = engine.get_tensor_shape(name)\n",
+    "\n",
+    "        # Calculate required memory size for this tensor\n",
+    "        size = np.dtype(trt.nptype(dtype)).itemsize\n",
+    "        for s in shape:\n",
+    "            size *= s\n",
+    "\n",
+    "        # Allocate GPU memory for this tensor\n",
+    "        err, allocation = cudart.cudaMalloc(size)\n",
+    "        check_cuda_error(err)\n",
+    "\n",
+    "        # Store tensor information in a dictionary for easy access\n",
+    "        binding = {\n",
+    "            \"index\": i,\n",
+    "            \"name\": name,\n",
+    "            \"dtype\": np.dtype(trt.nptype(dtype)),\n",
+    "            \"shape\": list(shape),\n",
+    "            \"allocation\": allocation,\n",
+    "            \"size\": size,\n",
+    "        }\n",
+    "\n",
+    "        # Keep track of all allocations and sort tensors into inputs/outputs\n",
+    "        allocations.append(allocation)\n",
+    "        if is_input:\n",
+    "            inputs.append(binding)\n",
+    "        else:\n",
+    "            outputs.append(binding)\n",
+    "\n",
+    "    # Ensure input data is contiguous in memory for efficient GPU transfer\n",
+    "    input_data = np.ascontiguousarray(input_data)\n",
+    "\n",
+    "    # Copy input data from host (CPU) to device (GPU)\n",
+    "    err = cudart.cudaMemcpy(\n",
+    "        inputs[0][\"allocation\"],\n",
+    "        input_data.ctypes.data,\n",
+    "        inputs[0][\"size\"],\n",
+    "        cudart.cudaMemcpyKind.cudaMemcpyHostToDevice,\n",
+    "    )\n",
+    "    check_cuda_error(err)\n",
+    "\n",
+    "    # Set tensor addresses for all tensors\n",
+    "    for i in range(engine.num_io_tensors):\n",
+    "        context.set_tensor_address(engine.get_tensor_name(i), allocations[i])\n",
+    "\n",
+    "    # Create a CUDA stream for asynchronous execution\n",
+    "    err, stream = cudart.cudaStreamCreate()\n",
+    "    check_cuda_error(err)\n",
+    "\n",
+    "    # Run inference using the TensorRT engine\n",
+    "    context.execute_async_v3(stream_handle=stream)\n",
+    "    err = cudart.cudaStreamSynchronize(stream)\n",
+    "    check_cuda_error(err)\n",
+    "\n",
+    "    # Prepare numpy array for output and copy results from GPU to CPU\n",
+    "    output_shape = outputs[0][\"shape\"]\n",
+    "    output = np.empty(output_shape, dtype=outputs[0][\"dtype\"])\n",
+    "\n",
+    "    err = cudart.cudaMemcpy(\n",
+    "        output.ctypes.data,\n",
+    "        outputs[0][\"allocation\"],\n",
+    "        outputs[0][\"size\"],\n",
+    "        cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost,\n",
+    "    )\n",
+    "    check_cuda_error(err)\n",
+    "\n",
+    "    # Free all GPU memory allocations\n",
+    "    for allocation in allocations:\n",
+    "        err = cudart.cudaFree(allocation)\n",
+    "        check_cuda_error(err)\n",
+    "\n",
+    "    # Destroy the CUDA stream\n",
+    "    err = cudart.cudaStreamDestroy(stream)\n",
+    "    check_cuda_error(err)\n",
+    "\n",
+    "    return output\n",
+    "\n",
+    "import onnxruntime as ort\n",
+    "def run_inference_onnx(session, input_data: np.ndarray):\n",
+    "    output = session.run(None, {'x': input_data})[0]\n",
+    "    return output"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Let's Compare Performance!\n",
+    "\n",
+    "We'll run both models multiple times to get an accurate comparison of their performance. This will show you the baseline speedup that TensorRT provides. \n",
+    "\n",
+    "Refer to https://docs.nvidia.com/deeplearning/tensorrt/latest/index.html for more information about how to further optimize your engine"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Create a sample input\n",
+    "sample_input = np.random.randn(1, 3, 224, 224).astype(np.float32)\n",
+    "\n",
+    "# Benchmark ONNX Runtime\n",
+    "session = ort.InferenceSession(onnx_path)\n",
+    "onnx_times = []\n",
+    "for _ in range(100):\n",
+    "    start_time = time.time()\n",
+    "    _ = run_inference_onnx(session, sample_input)\n",
+    "    onnx_times.append(time.time() - start_time)\n",
+    "\n",
+    "# Benchmark TensorRT\n",
+    "with open(engine_path, \"rb\") as f, trt.Runtime(logger) as runtime:\n",
+    "    engine = runtime.deserialize_cuda_engine(f.read())\n",
+    "trt_times = []\n",
+    "for _ in range(100):\n",
+    "    start_time = time.time()\n",
+    "    _ = run_inference_trt(engine, sample_input)\n",
+    "    trt_times.append(time.time() - start_time)\n",
+    "\n",
+    "print(f\"ONNX Runtime Average Time: {np.mean(onnx_times)*1000:.2f} ms\")\n",
+    "print(f\"TensorRT Average Time: {np.mean(trt_times)*1000:.2f} ms\")\n",
+    "print(f\"Speedup: {np.mean(onnx_times)/np.mean(trt_times):.2f}x\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Step 4: Run Inference on a Real Image\n",
+    "\n",
+    "Now let's try our optimized model on a real image! We'll:\n",
+    "1. Download a sample image\n",
+    "2. Load the ImageNet class labels\n",
+    "3. Make predictions and show the results\n",
+    "\n",
+    "This will demonstrate how the optimized TensorRT engine performs in a real-world scenario."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Download a sample image\n",
+    "download_file(\"https://raw.githubusercontent.com/pytorch/hub/master/images/dog.jpg\", root / \"test_image.jpg\")\n",
+    "\n",
+    "from PIL import Image\n",
+    "from IPython.display import display\n",
+    "\n",
+    "# Open and display the image\n",
+    "img = Image.open(root/\"test_image.jpg\")\n",
+    "display(img)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def load_imagenet_labels():\n",
+    "    # Download ImageNet labels if not exists\n",
+    "    if not (root / \"imagenet_classes.txt\").is_file():\n",
+    "        download_file(\"https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt\", root / \"imagenet_classes.txt\")\n",
+    "    # Read the labels\n",
+    "    with open(root / \"imagenet_classes.txt\") as f:\n",
+    "        categories = [s.strip() for s in f.readlines()]\n",
+    "    return categories\n",
+    "\n",
+    "# Load ImageNet labels\n",
+    "categories = load_imagenet_labels()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Load and preprocess a test image\n",
+    "test_image_path = root / \"test_image.jpg\"\n",
+    "input_data = load_and_preprocess_image(test_image_path)\n",
+    "\n",
+    "# Run inference\n",
+    "output = run_inference_trt(engine, input_data)\n",
+    "\n",
+    "# Get top 5 predictions\n",
+    "top5_idx = np.argsort(output[0])[-5:][::-1]\n",
+    "print(\"Top 5 predictions:\")\n",
+    "for idx in top5_idx:\n",
+    "    print(f\"{categories[idx]}: {output[0][idx]:.2f}%\")\n",
+    "assert categories[top5_idx[0]] == \"Samoyed\", 'Incorrect prediction'\n",
+    "print('Correctly recognized!')\n",
+    "print('Notebook executed successfully')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Congratulations! 🎉\n",
+    "\n",
+    "### You've successfully:\n",
+    "1. Loaded a pre-trained EfficientNet model in ONNX format\n",
+    "2. Converted it to a TensorRT engine\n",
+    "3. Achieved significant speedup in inference\n",
+    "4. Made predictions on real images\n",
+    "5. Learned how to use timing cache to speed up engine building and ensure engine build determinism. \n",
+    "\n",
+    "### What's Next?\n",
+    "\n",
+    "Now that you understand the ONNX to TensorRT workflow, you can:\n",
+    "- Export your own models from PyTorch/TensorFlow to ONNX\n",
+    "- Try different optimization settings in TensorRT\n",
+    "- Apply this workflow to your production models and get instant performance boost with NVIDIA GPUs!\n"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

samples/python/refactored/2_construct_network_with_layer_apis/README.md

@@ -0,0 +1,21 @@
+# Sample 2: Constructing a Network with TensorRT Layer APIs
+
+This sample demonstrates how to build a TensorRT network definition from scratch using the TensorRT Layer APIs, focusing on constructing a recurrent neural network (LSTM) and utilizing advanced builder features.
+
+## Description
+
+This sample constructs a simple, single-layer Long Short-Term Memory (LSTM) network using the TensorRT Layer APIs. The primary goal is to illustrate how to:
+
+1.  Define individual network layers and their connections programmatically using Python (**TensorRT Layer API**). This includes layers like constants, matrix multiply, element-wise operations, activations, and slicing.
+2.  Implement recurrent logic by building an LSTM cell and using TensorRT's `add_loop` construct to create a recurrent LSTM layer (**Recurrent Network Construction**).
+3.  Monitor the potentially lengthy engine build process by implementing `IProgressMonitor` for real-time feedback (**Build Progress Monitoring**).
+4.  Configure the builder for engine portability using `BuilderFlag.VERSION_COMPATIBLE` to create more portable engines (**Version-Compatible Engines**).
+5.  Run inference and verify the custom network's correctness by utilizing Polygraphy's `TrtRunner` for simplified engine loading/execution (**Inference with Polygraphy**) and comparing the TensorRT engine's output against a reference NumPy implementation (**NumPy Verification**).
+
+## Additional Resources
+*   [tensorrtx repo](https://github.com/wang-xinyu/tensorrtx): Offers real-world examples of constructing complex networks using the TensorRT Layer APIs.
+
+# Changelog
+
+August 2025
+Removed support for Python versions < 3.10.

samples/python/refactored/2_construct_network_with_layer_apis/main.ipynb

@@ -0,0 +1,694 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.SPDX-License-Identifier: Apache-2.0\n",
+    "\n",
+    "Licensed under the Apache License, Version 2.0 (the \"License\"); you may not use\n",
+    "\n",
+    "this file except in compliance with the License. You may obtain a copy of the License at\n",
+    "\n",
+    "\n",
+    "\n",
+    "http://www.apache.org/licenses/LICENSE-2.0\n",
+    "\n",
+    "\n",
+    "\n",
+    "Unless required by applicable law or agreed to in writing, software\n",
+    "\n",
+    "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+    "\n",
+    "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+    "\n",
+    "See the License for the specific language governing permissions and\n",
+    "\n",
+    "limitations under the License."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 2. Constructing a Network with TensorRT Layer APIs\n",
+    "\n",
+    "In this notebook, you'll learn how to move beyond pre-built model formats and directly construct a neural network using TensorRT's versatile Layer APIs. This approach offers fine-grained control over your network architecture and optimizations.\n",
+    "\n",
+    "Specifically, we will cover:\n",
+    "\n",
+    "1.  **Building a Recurrent Network (LSTM) from Scratch:** Understand how to define each layer of a Long Short-Term Memory (LSTM) cell and then use these components to construct an entire recurrent LSTM layer. This involves using various Layer API functionalities like `add_constant`, `add_matrix_multiply`, `add_elementwise`, `add_slice`, and `add_activation`.\n",
+    "2.  **Implementing Loops for Recurrence:** Utilize TensorRT's `add_loop` functionality to efficiently handle the recurrent nature of the LSTM, processing an input sequence step-by-step.\n",
+    "3.  **Monitoring Build Progress:** Implement an `IProgressMonitor` to track the engine creation process in real-time, providing visibility into potentially long build times.\n",
+    "4.  **Creating Version-Compatible Engines:** Learn to save TensorRT engines with the `BuilderFlag.VERSION_COMPATIBLE` flag, enhancing their portability across different TensorRT patch versions and compatible hardware.\n",
+    "\n",
+    "This sample uses a small, single-layer LSTM to keep the focus on these core TensorRT API features. We'll also verify its output against an equivalent NumPy implementation to ensure correctness."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Introduction\n",
+    "\n",
+    "While importing models via ONNX offers convenience, constructing networks directly with TensorRT APIs provides fine-grained control over the network definition. The **[TensorRT Layer API](https://docs.nvidia.com/deeplearning/tensorrt/latest/python_api/infer/Graph/Layers.html)** enables users to define each layer explicitly, offering flexibility and optimization opportunities.\n",
+    "\n",
+    "To facilitate understanding and verification, this demonstration employs small tensors, allowing for direct comparison with an equivalent NumPy implementation.\n",
+    "\n",
+    "> **Note: This sample assumes familiarity with the basic concepts of Long Short-Term Memory (LSTM) networks. If you're new to LSTMs, you might find it helpful to review their structure and operation before proceeding.**"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Step 0: Prerequisites"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%pip install numpy tensorrt polygraphy --extra-index-url https://pypi.ngc.nvidia.com\n",
+    "import tensorrt as trt\n",
+    "import numpy as np\n",
+    "from typing import Tuple"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "To simplify, network parameters and weight initializations use small, illustrative values and dimensions."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# === Network Parameters & Weights Initialization ===\n",
+    "batch_size = 1\n",
+    "seq_len = 5      # Length of the sequence\n",
+    "input_size = 1   # Dimension of input vector at each time step\n",
+    "hidden_size = 2  # Dimension of hidden state and cell state\n",
+    "num_units = 1\n",
+    "\n",
+    "# --- Create Fixed Dummy Weights and Biases (NumPy arrays with dummy values) ---\n",
+    "# These will be used by both the TensorRT build and the NumPy verification\n",
+    "w_val, u_val, b_val = 0.01, 0.05, 0.3\n",
+    "initial_h_val = 0.1\n",
+    "initial_c_val = 0.2\n",
+    "\n",
+    "# Define shapes\n",
+    "w_shape = (input_size, 4 * hidden_size) # e.g., [1, 8] for layer 0\n",
+    "u_shape = (hidden_size, 4 * hidden_size)       # e.g., [2, 8]\n",
+    "b_shape = (4 * hidden_size,)                   # e.g., [8]\n",
+    "initial_h_shape = (batch_size, hidden_size)\n",
+    "initial_c_shape = (batch_size, hidden_size)\n",
+    "\n",
+    "# Create NumPy arrays\n",
+    "np_weight_W = np.full(w_shape, w_val, dtype=np.float32)\n",
+    "np_weight_U = np.full(u_shape, u_val, dtype=np.float32)\n",
+    "np_bias = np.full(b_shape, b_val, dtype=np.float32)\n",
+    "np_initial_h = np.full(initial_h_shape, initial_h_val, dtype=np.float32)\n",
+    "np_initial_c = np.full(initial_c_shape, initial_c_val, dtype=np.float32)\n",
+    "\n",
+    "# Create inputs for the network\n",
+    "np_inputs = np.ones((seq_len, batch_size, input_size), dtype=np.float32)\n",
+    "\n",
+    "print(\"NumPy Weights Initialized:\")\n",
+    "print(f\"  W shape : {np_weight_W.shape}\")\n",
+    "print(f\"  U shape : {np_weight_U.shape}\")\n",
+    "print(f\"  Bias shape : {np_bias.shape}\")\n",
+    "print(f\"  Initial H shape : {np_initial_h.shape}\")\n",
+    "print(f\"  Initial C shape : {np_initial_c.shape}\")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Step 1: Defining LSTM Operations with the Layer API\n",
+    "\n",
+    "This step involves defining the LSTM operations by adding layers to the TensorRT `INetworkDefinition`.\n",
+    "\n",
+    "### Typical Usage Pattern for TensorRT Layer APIs\n",
+    "\n",
+    "When adding layers to a TensorRT network using the Layer API, the common pattern is:\n",
+    "\n",
+    "1.  **Add the layer:** Use a `network.add_*` method (e.g., `network.add_matrix_multiply`) to add the desired layer. This method takes input tensors and layer-specific parameters, returning an `ILayer` object representing the newly added layer.\n",
+    "2.  **Configure the layer:** Access the returned `ILayer` object to configure its properties. This is optional but useful for naming layer's name and output tensors name for easier debugging and more helpful logs. \n",
+    "\n",
+    "```python\n",
+    "# Example: Adding and configuring a generic layer\n",
+    "\n",
+    "# 1. Add the layer (replace with a specific layer like add_matrix_multiply)\n",
+    "layer = network.add_some_layer(input_tensor, ...)\n",
+    "\n",
+    "# 2. Configure the layer (optional)\n",
+    "output_tensor = layer.get_output(0)\n",
+    "output_tensor.name = 'my_layer_output'  # Name the output\n",
+    "# ... other configurations ...\n",
+    "```\n",
+    "\n",
+    "For a comprehensive list of available layer types and their specific methods and properties, consult the official [TensorRT Layer API documentation](https://docs.nvidia.com/deeplearning/tensorrt/latest/_static/python-api/infer/Graph/Layers.html)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "TRT_LOGGER = trt.Logger(trt.Logger.INFO)\n",
+    "\n",
+    "def add_lstm_unit(network: trt.INetworkDefinition,\n",
+    "                  input_x: trt.ITensor,      # Shape: [batch_size, input_size]\n",
+    "                  prev_h: trt.ITensor,       # Shape: [batch_size, hidden_size]\n",
+    "                  prev_c: trt.ITensor,       # Shape: [batch_size, hidden_size]\n",
+    "                  W: np.ndarray,            # Shape: [input_size, 4 * hidden_size]\n",
+    "                  U: np.ndarray,            # Shape: [hidden_size, 4 * hidden_size]\n",
+    "                  bias: np.ndarray,         # Shape: [4 * hidden_size]\n",
+    "                  hidden_size: int,\n",
+    "                  input_size: int\n",
+    "                  ) -> Tuple[trt.ITensor, trt.ITensor]:\n",
+    "    \"\"\"\n",
+    "    Adds the computations for a single LSTM time step.\n",
+    "    Assumes input tensors have a leading batch dimension.\n",
+    "    \"\"\"\n",
+    "    batch_size = input_x.shape[0] # Get batch size from input\n",
+    "\n",
+    "    # Create constant layers for weights and biases\n",
+    "    W_layer = network.add_constant(W.shape, trt.Weights(W))\n",
+    "    W_layer.get_output(0).name = \"W_const\"\n",
+    "    U_layer = network.add_constant(U.shape, U)\n",
+    "    U_layer.get_output(0).name = \"U_const\"\n",
+    "    # Reshape bias for broadcasting: [4*hidden] -> [1, 4*hidden]\n",
+    "    bias_reshaped_np = np.expand_dims(bias.copy(), axis=0)\n",
+    "    bias_layer = network.add_constant(bias_reshaped_np.shape, bias_reshaped_np)\n",
+    "    bias_layer.get_output(0).name = \"Bias_const\"\n",
+    "\n",
+    "\n",
+    "    # Linear transformations: Wx = input_x * W ; Uh = prev_h * U\n",
+    "    # Wx = [batch, input] * [input, 4*hidden] = [batch, 4*hidden]\n",
+    "    mm_wx = network.add_matrix_multiply(input_x, trt.MatrixOperation.NONE,\n",
+    "                                        W_layer.get_output(0), trt.MatrixOperation.NONE)\n",
+    "    mm_wx.get_output(0).name = \"Wx\"\n",
+    "\n",
+    "    # Uh = [batch, hidden] * [hidden, 4*hidden] = [batch, 4*hidden]\n",
+    "    mm_uh = network.add_matrix_multiply(prev_h, trt.MatrixOperation.NONE,\n",
+    "                                        U_layer.get_output(0), trt.MatrixOperation.NONE)\n",
+    "    mm_uh.get_output(0).name = \"Uh\"\n",
+    "\n",
+    "\n",
+    "    # Combined gates = Wx + Uh + Bias\n",
+    "    gates_wx_uh = network.add_elementwise(mm_wx.get_output(0), mm_uh.get_output(0),\n",
+    "                                         trt.ElementWiseOperation.SUM)\n",
+    "    gates_wx_uh.get_output(0).name = \"Wx_plus_Uh\"\n",
+    "\n",
+    "    gates = network.add_elementwise(gates_wx_uh.get_output(0), bias_layer.get_output(0),\n",
+    "                                    trt.ElementWiseOperation.SUM)\n",
+    "\n",
+    "    gates_output = gates.get_output(0) # Shape [batch, 4*hidden]\n",
+    "    gates_output.name = \"Gates_Combined\"\n",
+    "\n",
+    "    # Split the combined gates tensor [batch, 4*hidden] -> four [batch, hidden] gate tensors (Input, Forget, Candidate, Output)\n",
+    "    def add_gate_slice(index):\n",
+    "        gate_slice_layer = network.add_slice(input=gates_output,\n",
+    "                                       start=(0, index * hidden_size), # Start [batch_idx=0, col_idx]\n",
+    "                                       shape=(batch_size, hidden_size), # Slice shape\n",
+    "                                       stride=(1, 1))                   # Stride\n",
+    "        return gate_slice_layer.get_output(0)\n",
+    "\n",
+    "    slice_i = add_gate_slice(0)\n",
+    "    slice_i.name = \"Slice_I\"\n",
+    "    slice_f = add_gate_slice(1)\n",
+    "    slice_f.name = \"Slice_F\"\n",
+    "    slice_c = add_gate_slice(2)\n",
+    "    slice_c.name = \"Slice_C_candidate\" # Cell candidate\n",
+    "    slice_o = add_gate_slice(3)\n",
+    "    slice_o.name = \"Slice_O\"\n",
+    "\n",
+    "    # Apply activations\n",
+    "    act_i_layer = network.add_activation(slice_i, trt.ActivationType.SIGMOID)\n",
+    "    act_i = act_i_layer.get_output(0)\n",
+    "    act_i.name = \"Gate_I\"\n",
+    "    act_f_layer = network.add_activation(slice_f, trt.ActivationType.SIGMOID)\n",
+    "    act_f = act_f_layer.get_output(0)\n",
+    "    act_f.name = \"Gate_F\"\n",
+    "    act_c_layer = network.add_activation(slice_c, trt.ActivationType.TANH)\n",
+    "    act_c = act_c_layer.get_output(0)\n",
+    "    act_c.name = \"Gate_C_candidate\"\n",
+    "    act_o_layer = network.add_activation(slice_o, trt.ActivationType.SIGMOID)\n",
+    "    act_o = act_o_layer.get_output(0)\n",
+    "    act_o.name = \"Gate_O\"\n",
+    "\n",
+    "    # Cell state update: c_t = f_t * c_{t-1} + i_t * g_t\n",
+    "    term1_c = network.add_elementwise(act_f, prev_c, trt.ElementWiseOperation.PROD)\n",
+    "    term2_c = network.add_elementwise(act_i, act_c, trt.ElementWiseOperation.PROD)\n",
+    "    next_c_layer = network.add_elementwise(term1_c.get_output(0), term2_c.get_output(0), trt.ElementWiseOperation.SUM)\n",
+    "\n",
+    "    next_c = next_c_layer.get_output(0)\n",
+    "    next_c.name = \"next_c\" # Shape [batch, hidden]\n",
+    "\n",
+    "    # Hidden state update: h_t = o_t * tanh(c_t)\n",
+    "    tanh_c_layer = network.add_activation(next_c, trt.ActivationType.TANH)\n",
+    "    tanh_c = tanh_c_layer.get_output(0)\n",
+    "    next_h_layer = network.add_elementwise(act_o, tanh_c, trt.ElementWiseOperation.PROD)\n",
+    "\n",
+    "    next_h = next_h_layer.get_output(0)\n",
+    "    next_h.name = \"next_h\" # Shape [batch, hidden]\n",
+    "\n",
+    "    return next_h, next_c\n",
+    "\n",
+    "\n",
+    "def add_lstm_layer(network: trt.INetworkDefinition,\n",
+    "                   input_sequence: trt.ITensor, # Shape: [seq_len, batch_size, input_size]\n",
+    "                   hidden_size: int,\n",
+    "                   seq_len: int,\n",
+    "                   weight_W: np.ndarray, # [input_size, 4*hidden] or [hidden, 4*hidden]\n",
+    "                   weight_U: np.ndarray, # [hidden, 4*hidden]\n",
+    "                   bias: np.ndarray    # [4*hidden]\n",
+    "                   ) -> trt.ITensor:\n",
+    "    \"\"\"\n",
+    "    Adds a LSTM to the network by adding one lstm_unit, and run multiple times with loops.\n",
+    "    \"\"\"\n",
+    "    # Infer batch_size and input_size from the input tensor shape\n",
+    "    assert len(input_sequence.shape) == 3, f\"Input sequence tensor must have 3 dimensions [seq, batch, input]. Got shape {input_sequence.shape}\"\n",
+    "    input_size = input_sequence.shape[2]\n",
+    "\n",
+    "    # Shape: [batch_size, hidden_size]\n",
+    "    initial_h = network.add_constant(np_initial_h.shape, np_initial_h).get_output(0)\n",
+    "    initial_h.name = \"Initial_H\"\n",
+    "    initial_c = network.add_constant(np_initial_c.shape, np_initial_c).get_output(0)\n",
+    "    initial_c.name = \"Initial_C\"\n",
+    "\n",
+    "    loop = network.add_loop()\n",
+    "    loop.name = \"Time_Loop_Layer\"\n",
+    "\n",
+    "    # add_trip_limit determines when the loop should stop. For here we want the loop to run seq_len times.\n",
+    "    trip_limit = network.add_constant((), np.array([seq_len], dtype=np.int32)).get_output(0)\n",
+    "    loop.add_trip_limit(trip_limit, trt.TripLimit.COUNT)\n",
+    "\n",
+    "    # Recurrences for hidden and cell states\n",
+    "    h_recurrence = loop.add_recurrence(initial_h)\n",
+    "    c_recurrence = loop.add_recurrence(initial_c)\n",
+    "    prev_h_tensor = h_recurrence.get_output(0)\n",
+    "    prev_h_tensor.name = \"Prev_H\"\n",
+    "    prev_c_tensor = c_recurrence.get_output(0)\n",
+    "    prev_c_tensor.name = \"Prev_C\"\n",
+    "\n",
+    "    # add_iterator iterates through slices of the input sequence along the specified axis, providing one slice per iteration.\n",
+    "    x_t_iterator = loop.add_iterator(input_sequence, axis=0)\n",
+    "    x_t = x_t_iterator.get_output(0)\n",
+    "    x_t.name = \"x_t\"\n",
+    "\n",
+    "\n",
+    "    # Call the LSTM unit function\n",
+    "    next_h, next_c = add_lstm_unit(network=network,\n",
+    "                                    input_x=x_t,\n",
+    "                                    prev_h=prev_h_tensor,\n",
+    "                                    prev_c=prev_c_tensor,\n",
+    "                                    W=weight_W,\n",
+    "                                    U=weight_U,\n",
+    "                                    bias=bias,\n",
+    "                                    hidden_size=hidden_size,\n",
+    "                                    input_size=input_size)\n",
+    "\n",
+    "    # Feed the computed states back into the recurrence inputs\n",
+    "    h_recurrence.set_input(1, next_h)\n",
+    "    c_recurrence.set_input(1, next_c)\n",
+    "\n",
+    "    # add_loop_output() collects the values in the loop and outputs them. For this example, we concatenate the values along the first axis.\n",
+    "    loop_output_h = loop.add_loop_output(next_h, trt.LoopOutput.CONCATENATE, axis=0)\n",
+    "\n",
+    "    # when using CONCATENATE, the second input must be the trip limit.\n",
+    "    loop_output_h.set_input(1, trip_limit)\n",
+    "    loop_output_h.get_output(0).name = \"Hidden_Sequence\"\n",
+    "\n",
+    "    # --- End of time step loop definition ---\n",
+    "\n",
+    "    layer_output_sequence = loop_output_h.get_output(0)\n",
+    "\n",
+    "    # The final output sequence is the sequence from the last layer\n",
+    "    if layer_output_sequence is None:\n",
+    "         raise RuntimeError(\"LSTM layer output was not generated (num_layers may be 0)\")\n",
+    "    layer_output_sequence.name = \"Final_LSTM_Output_Sequence\"\n",
+    "    return layer_output_sequence"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Step 2: Build the Network\n",
+    "\n",
+    "Now that we have the LSTM layer implementation (`add_lstm_layer`), let's proceed to build the TensorRT `INetworkDefinition`.\n",
+    "This involves defining the network structure by:\n",
+    "1. Adding the input tensor using `network.add_input`.\n",
+    "2. Adding the LSTM layer using our custom `add_lstm_layer` function.\n",
+    "3. Marking the LSTM layer's output tensor as the network's final output."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "builder = trt.Builder(TRT_LOGGER)\n",
+    "network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED))\n",
+    "\n",
+    "# === Network Definition ===\n",
+    "# Shape: [seq_len, batch_size, input_size] -> e.g., [5, 1, 1]\n",
+    "input_tensor = network.add_input(name='input', dtype=trt.float32, shape=(seq_len, batch_size, input_size))\n",
+    "\n",
+    "# --- Add SINGLE LSTM Layer ---\n",
+    "lstm_output = add_lstm_layer(network=network,\n",
+    "                                input_sequence=input_tensor,\n",
+    "                                hidden_size=hidden_size,\n",
+    "                                seq_len=seq_len,\n",
+    "                                weight_W=np_weight_W, \n",
+    "                                weight_U=np_weight_U, \n",
+    "                                bias=np_bias) \n",
+    "# lstm_output shape: [seq_len, batch_size, hidden_size] -> e.g., [5, 1, 2]\n",
+    "\n",
+    "# --- Mark Output ---\n",
+    "lstm_output.name = 'hidden_state_sequence'\n",
+    "network.mark_output(lstm_output)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Step 3: Build the Engine\n",
+    "\n",
+    "Now that we have defined the network (`INetworkDefinition`), the next step is to build the optimized TensorRT engine. This process involves using the `trt.Builder` along with an `trt.BuilderConfig` object to specify how the engine should be built.\n",
+    "\n",
+    "The `IBuilderConfig` allows you to control various aspects of the build process, such as:\n",
+    "*   Setting memory constraints (e.g., workspace size using `set_memory_pool_limit`).\n",
+    "*   Setting builder flags to control optimization strategies and compatibility.\n",
+    "\n",
+    "Once the network and configuration are ready, the `builder.build_serialized_network(network, config)` method is called to produce the serialized engine, which can then be saved to a file or used directly.\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## (Optional) Defining a Progress Monitor\n",
+    "Building a TensorRT engine can sometimes take a while, especially for complex models. Don't worry if the build seems long! TensorRT offers a helpful tool called `IProgressMonitor`. This interface lets you track the build process step-by-step, making it easier to monitor progress and even debug if needed. \n",
+    "\n",
+    "### Implementing `IProgressMonitor`\n",
+    "\n",
+    "To use the progress monitor, inherit from `trt.IProgressMonitor` and override its key methods:\n",
+    "\n",
+    "*   `phase_start(self, phase_name, parent_phase, num_steps)`: TensorRT calls this method when it begins a significant phase of the build process (e.g., \"Parsing ONNX Model\", \"Building Engine\"). \n",
+    "    *   `phase_name`: Name of the phase starting.\n",
+    "    *   `parent_phase`: Name of the parent phase, if this is a sub-phase (can be `None`).\n",
+    "    *   `num_steps`: The total number of steps expected for this phase.\n",
+    "*   `step_complete(self, phase_name, step)`: Called after each incremental step within a phase is completed.\n",
+    "    *   `phase_name`: Name of the current phase.\n",
+    "    *   `step`: The index of the step that just finished (0-based).\n",
+    "    *   *Your implementation* usually updates the corresponding progress indicator.\n",
+    "    *   **Crucially, this method must return `True` to allow the build to continue.** Returning `False` or `None` will signal TensorRT to cancel the build.\n",
+    "*   `phase_finish(self, phase_name)`: Called when a phase (and all its steps) is completed.\n",
+    "    *   `phase_name`: Name of the phase that finished.\n",
+    "    *   *Your implementation* typically finalizes and removes the progress indicator for this phase.\n",
+    "\n",
+    "After that, hook it with `IBuilderConfig` by setting `config.progress_monitor = MyProgressMonitor()`"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class SimpleProgressMonitor(trt.IProgressMonitor):\n",
+    "    def __init__(self):\n",
+    "        trt.IProgressMonitor.__init__(self)\n",
+    "        self._active_phases = 0\n",
+    "\n",
+    "    def phase_start(self, phase_name, parent_phase, num_steps):\n",
+    "        print(f\"[ProgressMonitor] Phase Start: {phase_name} ({num_steps} steps)\")\n",
+    "        self._active_phases += 1\n",
+    "\n",
+    "    def phase_finish(self, phase_name):\n",
+    "        print(f\"[ProgressMonitor] Phase Finish: {phase_name}\")\n",
+    "        self._active_phases -= 1\n",
+    "\n",
+    "    def step_complete(self, phase_name, step):\n",
+    "        print(f\"[ProgressMonitor] Step Complete: {phase_name}, Step {step}\")\n",
+    "        return True\n",
+    "\n",
+    "    @property\n",
+    "    def active_phases(self):\n",
+    "        return self._active_phases"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## (Optional) Version Compatible Engine\n",
+    "TensorRT engines are typically optimized for the specific GPU and TensorRT version they are built on. This maximizes performance but can cause incompatibility if the deployment environment differs.\n",
+    "\n",
+    "The `trt.BuilderFlag.VERSION_COMPATIBLE` flag addresses this by creating a more portable engine. This engine is less sensitive to minor variations in TensorRT versions or GPU models (within a compatible family), potentially at the cost of some performance compared to a non-compatible engine optimized for the exact target. It also reduces the need to rebuild the engine for every minor TensorRT update. Version compatibility is supported from TensorRT 8.6 onwards; the plan must be built with a version at least 8.6 or higher, and the runtime must also be 8.6 or higher.\n",
+    "\n",
+    "### Use Cases\n",
+    "*   Deploying across diverse hardware fleets with compatible GPUs/TRT versions.\n",
+    "*   Distributing applications where end-user system configurations vary.\n",
+    "*   Simplifying maintenance by avoiding frequent rebuilds for minor updates.\n",
+    "\n",
+    "### How it Works\n",
+    "Enabling `trt.BuilderFlag.VERSION_COMPATIBLE` instructs TensorRT to use more generic optimizations. By default, this flag also causes a copy of a \"lean runtime\" (a version-specific, stripped-down runtime component) to be packaged within the engine plan file. When you deserialize this engine plan on a compatible system, TensorRT recognizes the embedded lean runtime, loads it, and uses this runtime to deserialize and execute the rest of the plan. \n",
+    "\n",
+    "Because this process involves loading and executing code (the lean runtime) directly from the engine plan file, you must explicitly indicate that you trust the origin and integrity of the plan. This is done by setting `runtime.engine_host_code_allowed = True` on your `trt.Runtime` instance before attempting to deserialize the engine.\n",
+    "\n",
+    "> **Considerations for Multiple Version-Compatible Engines:**\n",
+    "If deploying many version-compatible engines, the embedded lean runtime in each plan can lead to large overall application sizes. An alternative is to exclude the runtime from the engine plan (using `trt.BuilderFlag.EXCLUDE_LEAN_RUNTIME`) and load it manually. This approach can significantly reduce the total deployment footprint. For detailed instructions, refer to the NVIDIA TensorRT documentation on [Manually Loading the Runtime](https://docs.nvidia.com/deeplearning/tensorrt/latest/inference-library/advanced.html#manually-loading-the-runtime)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ENGINE_FILE_PATH = './lstm_network.plan'\n",
+    "\n",
+    "config = builder.create_builder_config()\n",
+    "config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 28) # 256MB\n",
+    "config.progress_monitor = SimpleProgressMonitor()\n",
+    "config.set_flag(trt.BuilderFlag.VERSION_COMPATIBLE)\n",
+    "\n",
+    "print(\"Building engine...\")\n",
+    "serialized_engine = builder.build_serialized_network(network, config)\n",
+    "\n",
+    "print(\"Engine build completed.\")\n",
+    "with open(ENGINE_FILE_PATH, 'wb') as f:\n",
+    "    f.write(serialized_engine)\n",
+    "print(f\"Engine saved to {ENGINE_FILE_PATH}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Inference\n",
+    "\n",
+    "Once the TensorRT engine is built, the next step is typically to run inference to verify its functionality and performance. The standard process involves creating a runtime and execution context, managing GPU memory for inputs and outputs, transferring data between host and device, and executing the engine etc. While this process provides fine-grained control, it involves boilerplate code. This standard procedure was demonstrated in detail in Sample 1.\n",
+    "\n",
+    "In this sample, we'll simplify the inference process by using **[Polygraphy](https://github.com/NVIDIA/TensorRT/tree/main/tools/Polygraphy)**, a versatile toolkit included with TensorRT that automates many underlying details, such as:\n",
+    "*   Context creation\n",
+    "*   Buffer management\n",
+    "*   Data transfers\n",
+    "\n",
+    "> **Important Note:** While Polygraphy is excellent for debugging and testing due to its ease of use, it may introduce overhead.\n",
+    "> For optimal performance in deployment scenarios, consider handcrafting the inference code as demonstrated in the `1_run_onnx_with_tensorrt` sample.\n",
+    "\n",
+    "For more examples, please refer to [polygraphy examples](https://github.com/NVIDIA/TensorRT/tree/main/tools/Polygraphy/examples)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from polygraphy.backend.common import BytesFromPath\n",
+    "from polygraphy.backend.trt import EngineFromBytes, TrtRunner\n",
+    "\n",
+    "def run_inference_with_polygraphy(h_input: np.ndarray) -> np.ndarray:\n",
+    "    input_name = 'input'\n",
+    "    output_name = 'hidden_state_sequence'\n",
+    "\n",
+    "    # Prepare the feed dictionary required by Polygraphy\n",
+    "    # Ensure input is contiguous C-style array, which Polygraphy prefers.\n",
+    "    h_input_contiguous = np.ascontiguousarray(h_input)\n",
+    "    feed_dict = {input_name: h_input_contiguous}\n",
+    "\n",
+    "    print(f\"Loading engine from: {ENGINE_FILE_PATH}\")\n",
+    "    outputs = None\n",
+    "    load_engine = EngineFromBytes(BytesFromPath(ENGINE_FILE_PATH))\n",
+    "    with TrtRunner(load_engine) as runner:\n",
+    "        outputs = runner.infer(feed_dict=feed_dict)\n",
+    "        # Polygraphy automatically synchronizes, so no explicit stream sync needed here\n",
+    "\n",
+    "    output_sequence = outputs[output_name]\n",
+    "    print(f\"Output '{output_name}' shape: {output_sequence.shape}, dtype: {output_sequence.dtype}\")\n",
+    "    return output_sequence\n",
+    "\n",
+    "output_sequence = run_inference_with_polygraphy(np_inputs) \n",
+    "\n",
+    "if output_sequence is not None:\n",
+    "    print(f\"\\nInput Sequence (shape {np_inputs.shape}):\\n{np_inputs}\")\n",
+    "    print(f\"\\nOutput Hidden State Sequence (shape {output_sequence.shape}):\\n{output_sequence}\")\n",
+    "else:\n",
+    "    print(\"Inference failed.\")\n",
+    "    "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Verifying the Output (Comparison with Equivalent Operations in NumPy)\n",
+    "\n",
+    "To ensure our TensorRT LSTM implementation is correct, we'll compare its output with a reference implementation in NumPy. This is a common practice to validate custom layer logic.\n",
+    "\n",
+    "The NumPy version will mimic the same LSTM cell computations and unroll the loop over the time sequence."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def sigmoid_np(x):\n",
+    "    x_clipped = np.clip(x, -500, 500)  # avoid overflow\n",
+    "    return 1.0 / (1.0 + np.exp(-x_clipped))\n",
+    "\n",
+    "\n",
+    "def tanh_np(x):\n",
+    "    x_clipped = np.clip(x, -100, 100)  # avoid overflow\n",
+    "    return np.tanh(x_clipped)\n",
+    "\n",
+    "\n",
+    "def lstm_step_numpy(x_t, prev_h, prev_c, W, U, bias):\n",
+    "    # W: shape [input_size, 4*hidden_size]\n",
+    "    # U: shape [hidden_size, 4*hidden_size]\n",
+    "    # bias: shape [4*hidden_size]\n",
+    "    # x_t: shape [batch_size, input_size]\n",
+    "    # prev_h, prev_c: shape [batch_size, hidden_size]\n",
+    "\n",
+    "    hidden_size_ = prev_h.shape[1]\n",
+    "\n",
+    "    Wx = x_t @ W  # Shape [batch_size, 4*hidden_size]\n",
+    "    Uh = prev_h @ U  # Shape [batch_size, 4*hidden_size]\n",
+    "    gates = Wx + Uh + bias\n",
+    "\n",
+    "    # Split gates\n",
+    "    i = gates[:, 0 * hidden_size_ : 1 * hidden_size_]\n",
+    "    f = gates[:, 1 * hidden_size_ : 2 * hidden_size_]\n",
+    "    c = gates[:, 2 * hidden_size_ : 3 * hidden_size_]  # Cell candidate\n",
+    "    o = gates[:, 3 * hidden_size_ : 4 * hidden_size_]\n",
+    "\n",
+    "    i_act = sigmoid_np(i)\n",
+    "    f_act = sigmoid_np(f)\n",
+    "    c_act = tanh_np(c)\n",
+    "    o_act = sigmoid_np(o)\n",
+    "\n",
+    "    next_c = f_act * prev_c + i_act * c_act\n",
+    "    next_h = o_act * tanh_np(next_c)\n",
+    "\n",
+    "    return next_h, next_c\n",
+    "\n",
+    "\n",
+    "def lstm_layer_numpy(input_sequence_np, np_W, np_U, np_bias):\n",
+    "    seq_len_ = input_sequence_np.shape[0]\n",
+    "    final_output_sequence_np = None\n",
+    "    h = np_initial_h.copy()\n",
+    "    c = np_initial_c.copy()\n",
+    "\n",
+    "    layer_output_sequence_list = []\n",
+    "\n",
+    "    for t in range(seq_len_):\n",
+    "        # Slice the input sequence for this time step\n",
+    "        x_t = input_sequence_np[t, :, :]\n",
+    "\n",
+    "        h, c = lstm_step_numpy(x_t, h, c, np_W, np_U, np_bias)\n",
+    "        layer_output_sequence_list.append(h)\n",
+    "        layer_output_sequence_np = np.stack(layer_output_sequence_list, axis=0)\n",
+    "        final_output_sequence_np = layer_output_sequence_np\n",
+    "\n",
+    "    return final_output_sequence_np\n",
+    "\n",
+    "\n",
+    "numpy_output_sequence = lstm_layer_numpy(np_inputs, np_weight_W, np_weight_U, np_bias)\n",
+    "print(\"\\n--- NumPy LSTM Calculation Results ---\")\n",
+    "print(f\"Input Sequence (all ones, shape {np_inputs.shape}):\\n{np_inputs}\")\n",
+    "print(f\"\\nNumPy Output Hidden State Sequence (shape {numpy_output_sequence.shape}):\\n{numpy_output_sequence}\")\n",
+    "print(\"\\n--- Comparison ---\")\n",
+    "diff = np.abs(output_sequence - numpy_output_sequence)\n",
+    "max_diff = np.max(diff) if diff.size > 0 else 0.0\n",
+    "print(f\"Max absolute difference: {max_diff}\")\n",
+    "assert np.allclose(\n",
+    "    output_sequence, numpy_output_sequence, atol=1e-5\n",
+    "), f\"Output sequence mismatch between TensorRT and NumPy, max diff: {max_diff}\"\n",
+    "print(\"Notebook executed successfully\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Conclusion and Next Steps\n",
+    "\n",
+    "Congratulations! You have successfully:\n",
+    "- Defined an LSTM cell and layer using TensorRT's Layer APIs.\n",
+    "- Implemented a recurrent loop with `add_loop`.\n",
+    "- Monitored the engine build process using `IProgressMonitor`.\n",
+    "- Built a version-compatible TensorRT engine.\n",
+    "- Performed inference using the built engine via Polygraphy.\n",
+    "- Verified the results against a NumPy implementation.\n",
+    "\n",
+    "This sample demonstrates the fundamental building blocks for creating custom network architectures in TensorRT. From here, you can explore:\n",
+    "- More complex network structures.\n",
+    "- Different types of layers available in the TensorRT API.\n",
+    "- Advanced loop constructs and conditional logic.\n",
+    "- Further optimization techniques if performance is critical (though for this sample, we focused on API usage).\n",
+    "\n",
+    "By mastering the Layer API, you gain the power to optimize virtually any deep learning model for inference on NVIDIA GPUs."
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

samples/python/refactored/README.md

@@ -0,0 +1,30 @@
+# TensorRT Refactored Samples
+
+This directory contains refactored and improved versions of TensorRT samples, demonstrating best practices and modern implementations.
+
+## Available Samples
+
+| Sample Name | Description | Format |
+|-------------|-------------|---------|
+| [1_run_onnx_with_tensorrt](./1_run_onnx_with_tensorrt) | Demonstrates ONNX model conversion to TensorRT and inference comparison | `ipynb` |
+| [2_construct_network_with_layer_apis](./2_construct_network_with_layer_apis) | Constructing a Network with TensorRT Layer APIs | `ipynb` |
+
+
+
+## Launch Instructions
+
+1.  Navigate to the desired sample directory and start the Jupyter server:
+    ```bash
+    pip install notebook
+    cd 1_run_onnx_with_tensorrt # or any other sample
+    jupyter notebook
+    ```
+2.  Then, open the `main.ipynb` file in the Jupyter Notebook interface that opens in your web browser.
+
+# Changelog
+
+October 2025
+Migrate to strongly typed APIs.
+
+August 2025
+Removed support for Python versions < 3.10.

samples/python/requirements.txt

@@ -0,0 +1,4 @@
+pyyaml==6.0.3
+requests==2.32.4
+tqdm==4.66.4
+numpy==1.26.4

samples/python/sample_weight_stripping/README.md

@@ -0,0 +1,91 @@
+# Introduction To Building and Refitting Weight-stripped Engines from ONNX Models
+
+
+**Table Of Contents**
+- [Description](#description)
+- [How does this sample work?](#how-does-this-sample-work)
+- [Prerequisites](#prerequisites)
+- [Running the sample](#running-the-sample)
+	* [Sample `--help` options](#sample-help-options)
+- [Additional resources](#additional-resources)
+- [License](#license)
+- [Changelog](#changelog)
+- [Known issues](#known-issues)
+
+## Description
+
+This sample, sample_weight_stripping, is a Python sample which uses TensorRT to build a weight-stripped engine and later refit to a full engine for inference.
+
+## How does this sample work?
+
+This sample demonstrates how to build a weight-stripped engine from an ONNX model file using TensorRT Python API which can reduce the saved engine size. Later, the weight-stripped engine is refitted by parser refitter with the original ONNX model as input. The refitted full engine is used for inference and guarantees no performance and accuracy loss. In this sample, we use ResNet50 to showcase our features.
+
+## Prerequisites
+
+1. Install the dependencies for Python.
+
+	```bash
+	pip3 install -r requirements.txt
+	```
+
+2. Preparing sample data
+See [Preparing sample data](../../README.md#preparing-sample-data) in the main samples README.
+
+## Running the sample
+
+1.  Build and save both normal engine and weight-stripped engine:
+
+	```
+	python3 build_engines.py --output_stripped_engine=stripped_engine.trt --output_normal_engine=normal_engine.trt
+	```
+
+	After running this step, you can see two saved TensorRT engines. `stripped_engine.trt` contains a stripped engine (~2.3MB) and `normal_engine.trt` contains a normal engine with all weights included (~51MB). By using stripped engine build, we can greatly reduce the size of the saved engine file.
+
+
+	**Note:** If the TensorRT sample data is not installed in the default location, for example `/usr/src/tensorrt/data/`, the model directory must be specified. For example: `--stripped_onnx=/path/to/my/data/` sets the model path for building weight-stripped engine and `--original_onnx=/path/to/my/data/` sets the model path for building normal engine. In most of the cases, they can use the same ONNX model.
+
+2.  Refit the weight-stripped engine and perform inference with the weight-stripped engine and the normal engine:
+	```
+	python3 refit_engine_and_infer.py --stripped_engine=stripped_engine.trt -–normal_engine=normal_engine.trt
+	```
+3.  Verify that the sample ran successfully. If the sample runs successfully you should see output similar to the following. The prediction results of the refitted stripped engine is the same as the normal engine. There is no performance loss.
+	```
+	Normal engine inference time on 100 cases: 0.1066 seconds
+	Refitted stripped engine inference time on 100 cases: 0.0606 seconds
+	Normal engine correctly recognized data/samples/resnet50/tabby_tiger_cat.jpg as tiger cat
+	Refitted stripped engine correctly recognized data/samples/resnet50/tabby_tiger_cat.jpg as tiger cat
+	```
+### Sample --help options
+
+To see the full list of available options and their descriptions, use the `-h` or `--help` command line option.
+
+# Additional resources
+
+The following resources provide a deeper understanding about importing a model into TensorRT using Python:
+
+**ResNet-50**
+- [Deep Residual Learning for Image Recognition](https://arxiv.org/pdf/1512.03385.pdf)
+
+**Documentation**
+- [Introduction To NVIDIA’s TensorRT Samples](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sample-support-guide/index.html#samples)
+- [Working With TensorRT Using The Python API](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#python_topics)
+- [NVIDIA’s TensorRT Documentation Library](https://docs.nvidia.com/deeplearning/sdk/tensorrt-archived/index.html)
+
+# License
+
+For terms and conditions for use, reproduction, and distribution, see the [TensorRT Software License Agreement](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sla/index.html) documentation.
+
+# Changelog
+
+October 2025
+Migrate to strongly typed APIs.
+
+August 2025
+Removed support for Python versions < 3.10.
+
+February 2024
+Initial release of this sample.
+
+# Known issues
+
+There are no known issues in this sample.

samples/python/sample_weight_stripping/build_engines.py

@@ -0,0 +1,115 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2025 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 sys
+import argparse
+import math
+import time
+import datetime
+
+import tensorrt as trt
+
+sys.path.insert(1, os.path.join(sys.path[0], os.path.pardir))
+import common
+
+# You can set the logger severity higher to suppress messages (or lower to display more messages).
+TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
+
+
+def convert_size(size_bytes):
+    if size_bytes == 0:
+        return "0B"
+    size_name = ("B", "KB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB")
+    i = int(math.floor(math.log(size_bytes, 1024)))
+    p = math.pow(1024, i)
+    s = round(size_bytes / p, 2)
+    return "%s %s" % (s, size_name[i])
+
+def main(args):
+
+    with trt.Builder(TRT_LOGGER) as builder, builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)) as network, trt.OnnxParser(network, TRT_LOGGER) as parser:
+        with open(args.original_onnx, 'rb') as onnx_model:
+            parser.parse(onnx_model.read())
+
+        with builder.create_builder_config() as config:
+
+            config.set_flag(trt.BuilderFlag.STRIP_PLAN)
+
+            cache = config.create_timing_cache(b"")
+            config.set_timing_cache(cache, ignore_mismatch = False)
+
+            profile = builder.create_optimization_profile()
+            profile.set_shape("gpu_0/data_0", min=[1, 3, 224, 224], opt=[1, 3, 224, 224], max=[1, 3, 224, 224])
+            config.add_optimization_profile(profile)
+
+            def build_and_save_engine(builder, network, config, output):
+                start_time = time.time()
+                engine_bytes = builder.build_serialized_network(network, config)
+                assert engine_bytes is not None
+                with open(output, 'wb') as f:
+                    f.write(engine_bytes)
+                total_time = time.time() - start_time
+                print("built and saved {} in time {}".format(output, str(datetime.timedelta(seconds=int(total_time)))))
+
+            # build weight-stripped engine and generate timing cache.
+            build_and_save_engine(builder, network, config, args.output_stripped_engine)
+
+            # build normal engine with the same timing cache.
+            config.flags &= ~(1 << int(trt.BuilderFlag.STRIP_PLAN))
+            build_and_save_engine(builder, network, config, args.output_normal_engine)
+
+def get_default_model_file():
+    # Set the data path to the directory that contains the ONNX model.
+    _, data_files = common.find_sample_data(
+        description="Runs a ResNet50 network with a TensorRT inference engine.",
+        subfolder="resnet50",
+        find_files=["ResNet50.onnx"],
+    )
+    onnx_model_file = data_files[0]
+    return onnx_model_file
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("--stripped_onnx", default=None, type=str,
+                        help="The ONNX model file to load for building stripped engine.")
+    parser.add_argument("--original_onnx", default=None, type=str,
+                        help="The ONNX model file to load for building normal engine.")
+    parser.add_argument("--output_stripped_engine", default='stripped_engine.trt', type=str,
+                        help="The output path for the weight-stripped TRT engine.")
+    parser.add_argument("--output_normal_engine", default='normal_engine.trt', type=str,
+                        help="The output path for the full TRT engine.")
+    args, _ = parser.parse_known_args()
+
+    onnx_model_file = get_default_model_file()
+    if args.stripped_onnx is None:
+        args.stripped_onnx = onnx_model_file
+    if args.original_onnx is None:
+        args.original_onnx = onnx_model_file
+
+    if not os.path.exists(args.stripped_onnx):
+        parser.print_help()
+        print(f"--stripped_onnx {args.stripped_onnx} does not exist.")
+        sys.exit(1)
+    if not os.path.exists(args.original_onnx):
+        parser.print_help()
+        print(f"--original_onnx {args.original_onnx} does not exist.")
+        sys.exit(1)
+
+    main(args)

samples/python/sample_weight_stripping/refit_engine_and_infer.py

@@ -0,0 +1,173 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2026 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 argparse
+import random
+import sys
+import time
+import datetime
+import numpy as np
+
+import tensorrt as trt
+from PIL import Image
+
+sys.path.insert(1, os.path.join(sys.path[0], os.path.pardir))
+import common
+
+
+# You can set the logger severity higher to suppress messages (or lower to display more messages).
+TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
+
+class ModelData(object):
+    MODEL_PATH = "ResNet50.onnx"
+    INPUT_SHAPE = (3, 224, 224)
+    # We can convert TensorRT data types to numpy types with trt.nptype()
+    DTYPE = trt.float32
+
+def load_stripped_engine_and_refit(input_file, onnx_model_path):
+    runtime = trt.Runtime(TRT_LOGGER)
+
+    with open(input_file, 'rb') as engine_file:
+        engine = runtime.deserialize_cuda_engine(engine_file.read())
+        refitter = trt.Refitter(engine, TRT_LOGGER)
+        parser_refitter = trt.OnnxParserRefitter(refitter, TRT_LOGGER)
+        assert parser_refitter.refit_from_file(onnx_model_path)
+        assert refitter.refit_cuda_engine()
+
+        return engine
+
+def load_normal_engine(input_file):
+    runtime = trt.Runtime(TRT_LOGGER)
+    with open(input_file, 'rb') as engine_file:
+        engine = runtime.deserialize_cuda_engine(engine_file.read())
+
+        return engine
+
+
+def load_normalized_test_case(test_image, pagelocked_buffer):
+    # Converts the input image to a CHW Numpy array
+    def normalize_image(image):
+        # Resize, antialias and transpose the image to CHW.
+        c, h, w = ModelData.INPUT_SHAPE
+        image_arr = (
+            np.asarray(image.resize((w, h), Image.LANCZOS))
+            .transpose([2, 0, 1])
+            .astype(trt.nptype(ModelData.DTYPE))
+            .ravel()
+        )
+        # This particular ResNet50 model requires some preprocessing, specifically, mean normalization.
+        return (image_arr / 255.0 - 0.45) / 0.225
+
+    # Normalize the image and copy to pagelocked memory.
+    np.copyto(pagelocked_buffer, normalize_image(Image.open(test_image)))
+    return test_image
+
+def main(args):
+    # Set the data path to the directory that contains the trained models and test images for inference.
+    _, data_files = common.find_sample_data(
+        description="Runs a ResNet50 network with a TensorRT inference engine.",
+        subfolder="resnet50",
+        find_files=[
+            "binoculars.jpeg",
+            "reflex_camera.jpeg",
+            "tabby_tiger_cat.jpg",
+            ModelData.MODEL_PATH,
+            "class_labels.txt",
+        ],
+    )
+    # Get test images, models and labels.
+    test_images = data_files[0:3]
+    onnx_model_file, labels_file = data_files[3:]
+
+    labels = open(labels_file, "r").read().split("\n")
+
+    # Load a TensorRT engine.
+    engine = load_normal_engine(args.normal_engine)
+    refitted_engine = load_stripped_engine_and_refit(args.stripped_engine, onnx_model_file)
+
+    # Allocate buffers
+    inputs, outputs, bindings = common.allocate_buffers(engine)
+    inputs_1, outputs_1, bindings_1 = common.allocate_buffers(refitted_engine)
+
+    # Contexts are used to perform inference.
+    context = engine.create_execution_context()
+    context_1 = refitted_engine.create_execution_context()
+
+    # Load a normalized test case into the host input page-locked buffer.
+    test_image = random.choice(test_images)
+    test_case = load_normalized_test_case(test_image, inputs[0].host)
+    test_case_1 = load_normalized_test_case(test_image, inputs_1[0].host)
+
+    # Run the engine. The output will be a 1D tensor of length 1000, where each value represents the
+    # probability that the image corresponds to that label
+    
+    # Use context manager for proper stream lifecycle management - Normal engine
+    trt_outputs = []
+    with common.CudaStreamContext() as stream:
+        start_time = time.time()
+        for i in range(100): # count time for 100 times of inference
+            trt_outputs = common.do_inference(context, engine=engine, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream)
+        total_time = time.time() - start_time
+        print("Normal engine inference time on 100 cases: {:.4f} seconds".format(total_time))
+
+    # Use context manager for proper stream lifecycle management - Refitted engine
+    trt_outputs_refitted = []
+    with common.CudaStreamContext() as stream_1:
+        start_time = time.time()
+        for i in range(100):
+            trt_outputs_refitted = common.do_inference(context_1, engine=refitted_engine, bindings=bindings_1, inputs=inputs_1, outputs=outputs_1, stream=stream_1)
+        total_time = time.time() - start_time
+        print("Refitted stripped engine inference time on 100 cases: {:.4f} seconds".format(total_time))
+
+    # We use the highest probability as our prediction. Its index corresponds to the predicted label.
+    pred = labels[np.argmax(trt_outputs[0])]
+    if "_".join(pred.split()) in os.path.splitext(os.path.basename(test_case))[0]:
+        print("Normal engine correctly recognized " + test_case + " as " + pred)
+    else:
+        print("Normal engine incorrectly recognized " + test_case + " as " + pred)
+        exit(1)
+
+    pred_refitted = labels[np.argmax(trt_outputs_refitted[0])]
+    if "_".join(pred_refitted.split()) in os.path.splitext(os.path.basename(test_case_1))[0]:
+        print("Refitted stripped engine correctly recognized " + test_case + " as " + pred_refitted)
+    else:
+        print("Refitted stripped engine incorrectly recognized " + test_case + " as " + pred_refitted)
+        exit(1)
+
+    return trt_outputs, trt_outputs_refitted
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--stripped_engine", default='stripped_engine.trt', type=str,
+                        help="The stripped engine file to load.")
+    parser.add_argument("--normal_engine", default='normal_engine.trt', type=str,
+                        help="The normal engine file to load.")
+
+    args, _ = parser.parse_known_args()
+    if not os.path.exists(args.stripped_engine):
+        parser.print_help()
+        print(f"--stripped_engine {args.stripped_engine} does not exist.")
+        sys.exit(1)
+    if not os.path.exists(args.normal_engine):
+        parser.print_help()
+        print(f"--normal_engine {args.normal_engine} does not exist.")
+        sys.exit(1)
+
+    trt_outputs, trt_outputs_refitted = main(args)
+    print("The MSE of the final layer output is", np.square(np.subtract(trt_outputs, trt_outputs_refitted)).mean())

samples/python/sample_weight_stripping/requirements.txt

@@ -0,0 +1,7 @@
+Pillow==11.3.0
+cuda-python==12.9.0
+pywin32; platform_system == "Windows"
+pyyaml==6.0.3
+requests==2.32.4
+tqdm==4.66.4
+numpy==1.26.4

samples/python/simple_progress_monitor/README.md

@@ -0,0 +1,103 @@
+# Introduction To IProgressMonitor Callbacks Using Python
+
+
+**Table Of Contents**
+- [Description](#description)
+- [How does this sample work?](#how-does-this-sample-work)
+	* [simple_progress_monitor](#simple_progress_monitor)
+- [Prerequisites](#prerequisites)
+- [Running the sample](#running-the-sample)
+	* [Sample `--help` options](#sample-help-options)
+- [Additional resources](#additional-resources)
+- [License](#license)
+- [Changelog](#changelog)
+- [Known issues](#known-issues)
+
+## Description
+
+This sample, simple_progress_monitor, is a Python sample which uses TensorRT and its included ONNX parser, to perform inference with ResNet-50 models saved in ONNX format. It displays animated progress bars while TensorRT builds the engine.
+
+## How does this sample work?
+
+### simple_progress_monitor
+
+This sample demonstrates how to build an engine from an ONNX model file using the open-source ONNX parser and then run inference. The ONNX parser can be used with any framework that supports the ONNX format (typically `.onnx` files). An `IProgressMonitor` object receives updates on the progress of the build, and displays them as ASCII progress bars on stdout.
+
+## Prerequisites
+
+1. Install the dependencies for Python.
+
+```bash
+pip3 install -r requirements.txt
+```
+2. Preparing sample data
+See [Preparing sample data](../../README.md#preparing-sample-data) in the main samples README.
+
+## Running the sample
+
+1.  Run the sample from a terminal to create a TensorRT inference engine and run inference:
+	`python3 simple_progress_monitor.py`
+
+	**Note:** If the TensorRT sample data is not installed in the default location, the `data` directory must be specified. For example: `python3 simple_progress_monitor.py -d $TRT_DATADIR`
+
+	**Note:** Do not redirect the output of this script to a file or pipe.
+
+2.  Verify that the sample ran successfully. If the sample runs successfully you should see output similar to the following:
+	`Correctly recognized data/samples/resnet50/reflex_camera.jpeg as reflex camera`
+
+### Sample --help options
+
+To see the full list of available options and their descriptions, use the `-h` or `--help` command line option. For example:
+```
+usage: simple_progress_monitor.py [-h] [-d DATADIR]
+
+Runs a ResNet50 network with a TensorRT inference engine. Displays intermediate build progress.
+
+optional arguments:
+ -h, --help            show this help message and exit
+ -d DATADIR, --datadir DATADIR
+                       Location of the TensorRT sample data directory.
+                       (default: /usr/src/tensorrt/data)
+```
+
+# Additional resources
+
+The following resources provide a deeper understanding about importing a model into TensorRT using Python:
+
+**ResNet-50**
+- [Deep Residual Learning for Image Recognition](https://arxiv.org/pdf/1512.03385.pdf)
+
+**Parsers**
+- [ONNX Parser](https://docs.nvidia.com/deeplearning/sdk/tensorrt-api/python_api/parsers/Onnx/pyOnnx.html)
+
+**Documentation**
+- [Introduction To NVIDIA’s TensorRT Samples](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sample-support-guide/index.html#samples)
+- [Working With TensorRT Using The Python API](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#python_topics)
+- [Importing A Model Using A Parser In Python](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#import_model_python)
+- [NVIDIA’s TensorRT Documentation Library](https://docs.nvidia.com/deeplearning/sdk/tensorrt-archived/index.html)
+
+**Terminal Escape Sequences**
+- Linux: [XTerm Control Sequences](https://invisible-island.net/xterm/ctlseqs/ctlseqs.html)
+- Windows: [Console Virtual Terminal Sequences](https://learn.microsoft.com/en-us/windows/console/console-virtual-terminal-sequences)
+
+# License
+
+For terms and conditions for use, reproduction, and distribution, see the [TensorRT Software License Agreement](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sla/index.html) documentation.
+
+# Changelog
+
+October 2025
+Migrate to strongly typed APIs.
+
+August 2025
+Removed support for Python versions < 3.10.
+
+August 2023
+Removed support for Python versions < 3.8.
+
+June 2023
+This `README.md` file was created and reviewed.
+
+# Known issues
+
+There are no known issues in this sample

samples/python/simple_progress_monitor/requirements.txt

@@ -0,0 +1,7 @@
+Pillow==11.3.0
+cuda-python==12.9.0
+pywin32; platform_system == "Windows"
+pyyaml==6.0.3
+requests==2.32.4
+tqdm==4.66.4
+numpy==1.26.4

samples/python/simple_progress_monitor/simple_progress_monitor.py

@@ -0,0 +1,222 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2025 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
+
+# This sample demonstrates incremental progress reporting while it uses an ONNX ResNet50 Model to create a TensorRT Inference Engine.
+import random
+import sys
+
+import numpy as np
+
+import tensorrt as trt
+from PIL import Image
+
+sys.path.insert(1, os.path.join(sys.path[0], os.path.pardir))
+import common
+
+
+class ModelData(object):
+    MODEL_PATH = "ResNet50.onnx"
+    INPUT_SHAPE = (3, 224, 224)
+    # We can convert TensorRT data types to numpy types with trt.nptype().
+    DTYPE = trt.float32
+
+
+# This is a simple ASCII-art progress monitor comparable to the C++ version in sample_progress_monitor.
+class SimpleProgressMonitor(trt.IProgressMonitor):
+    def __init__(self):
+        trt.IProgressMonitor.__init__(self)
+        self._active_phases = {}
+        self._step_result = True
+
+    def phase_start(self, phase_name, parent_phase, num_steps):
+        try:
+            if parent_phase is not None:
+                nbIndents = 1 + self._active_phases[parent_phase]["nbIndents"]
+            else:
+                nbIndents = 0
+            self._active_phases[phase_name] = {
+                "title": phase_name,
+                "steps": 0,
+                "num_steps": num_steps,
+                "nbIndents": nbIndents,
+            }
+            self._redraw()
+        except KeyboardInterrupt:
+            # The phase_start callback cannot directly cancel the build, so request the cancellation from within step_complete.
+            _step_result = False
+
+    def phase_finish(self, phase_name):
+        try:
+            del self._active_phases[phase_name]
+            self._redraw(blank_lines=1)  # Clear the removed phase.
+        except KeyboardInterrupt:
+            _step_result = False
+
+    def step_complete(self, phase_name, step):
+        try:
+            self._active_phases[phase_name]["steps"] = step
+            self._redraw()
+            return self._step_result
+        except KeyboardInterrupt:
+            # There is no need to propagate this exception to TensorRT. We can simply cancel the build.
+            return False
+
+    def _redraw(self, *, blank_lines=0):
+        # The Python curses module is not widely available on Windows platforms.
+        # Instead, this function uses raw terminal escape sequences. See the sample documentation for references.
+        def clear_line():
+            print("\x1B[2K", end="")
+
+        def move_to_start_of_line():
+            print("\x1B[0G", end="")
+
+        def move_cursor_up(lines):
+            print("\x1B[{}A".format(lines), end="")
+
+        def progress_bar(steps, num_steps):
+            INNER_WIDTH = 10
+            completed_bar_chars = int(INNER_WIDTH * steps / float(num_steps))
+            return "[{}{}]".format(
+                "=" * completed_bar_chars, "-" * (INNER_WIDTH - completed_bar_chars)
+            )
+
+        # Set max_cols to a default of 200 if not run in interactive mode.
+        max_cols = os.get_terminal_size().columns if sys.stdout.isatty() else 200
+
+        move_to_start_of_line()
+        for phase in self._active_phases.values():
+            phase_prefix = "{indent}{bar} {title}".format(
+                indent=" " * phase["nbIndents"],
+                bar=progress_bar(phase["steps"], phase["num_steps"]),
+                title=phase["title"],
+            )
+            phase_suffix = "{steps}/{num_steps}".format(**phase)
+            allowable_prefix_chars = max_cols - len(phase_suffix) - 2
+            if allowable_prefix_chars < len(phase_prefix):
+                phase_prefix = phase_prefix[0 : allowable_prefix_chars - 3] + "..."
+            clear_line()
+            print(phase_prefix, phase_suffix)
+        for line in range(blank_lines):
+            clear_line()
+            print()
+        move_cursor_up(len(self._active_phases) + blank_lines)
+        sys.stdout.flush()
+
+
+# You can set the logger severity higher to suppress messages (or lower to display more messages).
+TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
+
+
+# The Onnx path is used for Onnx models.
+def build_engine_onnx(model_file):
+    builder = trt.Builder(TRT_LOGGER)
+    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED))
+    config = builder.create_builder_config()
+    if not sys.stdout.isatty():
+        print(
+            "Warning: This sample should be run from an interactive terminal in order to showcase the progress monitor correctly."
+        )
+    config.progress_monitor = SimpleProgressMonitor()
+    parser = trt.OnnxParser(network, TRT_LOGGER)
+
+    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, common.GiB(1))
+    # Load the Onnx model and parse it in order to populate the TensorRT network.
+    with open(model_file, "rb") as model:
+        if not parser.parse(model.read()):
+            print("ERROR: Failed to parse the ONNX file.")
+            for error in range(parser.num_errors):
+                print(parser.get_error(error))
+            return None
+
+    engine_bytes = builder.build_serialized_network(network, config)
+    runtime = trt.Runtime(TRT_LOGGER)
+    return runtime.deserialize_cuda_engine(engine_bytes)
+
+
+def load_normalized_test_case(test_image, pagelocked_buffer):
+    # Converts the input image to a CHW Numpy array.
+    def normalize_image(image):
+        # Resize, antialias and transpose the image to CHW.
+        c, h, w = ModelData.INPUT_SHAPE
+        image_arr = (
+            np.asarray(image.resize((w, h), Image.LANCZOS))
+            .transpose([2, 0, 1])
+            .astype(trt.nptype(ModelData.DTYPE))
+            .ravel()
+        )
+        # This particular ResNet50 model requires some preprocessing, specifically, mean normalization.
+        return (image_arr / 255.0 - 0.45) / 0.225
+
+    # Normalize the image and copy to pagelocked memory.
+    np.copyto(pagelocked_buffer, normalize_image(Image.open(test_image)))
+    return test_image
+
+
+def main():
+    # Set the data path to the directory that contains the trained models and test images for inference.
+    _, data_files = common.find_sample_data(
+        description="Runs a ResNet50 network with a TensorRT inference engine. Displays intermediate build progress.",
+        subfolder="resnet50",
+        find_files=[
+            "binoculars.jpeg",
+            "reflex_camera.jpeg",
+            "tabby_tiger_cat.jpg",
+            ModelData.MODEL_PATH,
+            "class_labels.txt",
+        ],
+    )
+    # Get test images, models and labels.
+    test_images = data_files[0:3]
+    onnx_model_file, labels_file = data_files[3:]
+    labels = open(labels_file, "r").read().split("\n")
+
+    # Build a TensorRT engine.
+    engine = build_engine_onnx(onnx_model_file)
+    # Inference is the same regardless of which parser is used to build the engine, since the model architecture is the same.
+    # Allocate buffers
+    inputs, outputs, bindings = common.allocate_buffers(engine)
+    # Contexts are used to perform inference.
+    context = engine.create_execution_context()
+
+    # Use context manager for proper stream lifecycle management
+    with common.CudaStreamContext() as stream:
+        # Load a normalized test case into the host input page-locked buffer.
+        test_image = random.choice(test_images)
+        test_case = load_normalized_test_case(test_image, inputs[0].host)
+        # Run the engine. The output will be a 1D tensor of length 1000, where each value represents the
+        # probability that the image corresponds to that label
+        trt_outputs = common.do_inference(
+            context,
+            engine=engine,
+            bindings=bindings,
+            inputs=inputs,
+            outputs=outputs,
+            stream=stream,
+        )
+        # We use the highest probability as our prediction. Its index corresponds to the predicted label.
+        pred = labels[np.argmax(trt_outputs[0])]
+        common.free_buffers(inputs, outputs)
+    if "_".join(pred.split()) in os.path.splitext(os.path.basename(test_case))[0]:
+        print("Correctly recognized " + test_case + " as " + pred)
+    else:
+        print("Incorrectly recognized " + test_case + " as " + pred)
+
+
+if __name__ == "__main__":
+    main()

samples/python/stream_writer/README.md

@@ -0,0 +1,72 @@
+# TensorRT Python Sample: Stream Writer
+
+This sample demonstrates how to use the TensorRT Python API to serialize an engine directly to a custom stream using the `IStreamWriter` interface, rather than to a file or in-memory buffer. This is useful for advanced scenarios where you want to control how and where the engine bytes are written (e.g., to a network socket, custom buffer, or in-memory stream).
+
+## What does this sample do?
+
+- Builds a simple TensorRT network with two convolutional layers and ReLU activations.
+- Implements a custom `StreamWriter` class inheriting from `trt.IStreamWriter` to collect serialized engine bytes.
+- Serializes the engine using `builder.build_serialized_network_to_stream()` and writes the bytes to the custom stream.
+- Deserializes the engine from the collected bytes to verify correctness.
+
+## File Structure
+
+- `build.py`: Main script containing the sample code.
+- `README.md`: This document.
+
+## How to Run
+
+1. **Install Requirements**
+
+   Make sure you have the following Python packages installed:
+   - `tensorrt`
+   - `numpy`
+   - `polygraphy`
+
+   You can install Polygraphy via pip:
+   ```
+   pip install polygraphy
+   ```
+
+   The `tensorrt` Python package is typically provided by NVIDIA as a wheel file.
+
+2. **Run the Sample**
+
+   ```
+   python3 build.py
+   ```
+
+   You should see output indicating the network is constructed, the engine is built and serialized to the stream, and then deserialized successfully.
+
+## Key Concepts
+
+- **IStreamWriter**: An interface in TensorRT that allows you to define custom logic for writing serialized engine bytes. You must implement the `write(self, data)` method.
+- **build_serialized_network_to_stream**: A method that serializes the network and writes the bytes to the provided `IStreamWriter` instance.
+
+## Example Output
+
+```
+Constructing network...
+[I] TF32 is disabled by default. Turn on TF32 for better performance with minor accuracy differences.
+[I] Configuring with profiles:[
+        Profile 0:
+            {input [min=[1, 3, 224, 224], opt=[1, 3, 224, 224], max=[1, 3, 224, 224]]}
+    ]
+Building engine and serializing to stream...
+The total bytes written to stream is  267836
+Deserializing engine from stream...
+Engine deserialized successfully
+```
+
+# License
+
+For terms and conditions for use, reproduction, and distribution, see the [TensorRT Software License Agreement](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sla/index.html) documentation.
+
+# Changelog
+
+September 2025
+Initial release of this sample.
+
+# Known issues
+
+There are no known issues in this sample.

samples/python/stream_writer/build.py

@@ -0,0 +1,78 @@
+#!/usr/bin/env python3
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2026 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 tensorrt as trt
+import numpy as np
+from polygraphy.logger import G_LOGGER
+from polygraphy.backend.trt import (
+    CreateNetwork,
+    CreateConfig,
+    engine_bytes_from_network,
+    get_trt_logger
+)
+DEBUG_LOG = False  # Turn on to print TRT verbose logs
+
+if DEBUG_LOG:
+    verbose = G_LOGGER.verbosity(G_LOGGER.SUPER_VERBOSE)
+    verbose.__enter__()
+else:
+    verbose = None
+
+def build_network():
+    builder, network = CreateNetwork()()
+    # A simple network with internal tensors
+    input_tensor = network.add_input(name="input", dtype=trt.float32, shape=(1, 3, 224, 224))
+    conv1_w = np.random.randn(16, 3, 3, 3).astype(np.float32)
+    conv1_b = np.random.randn(16).astype(np.float32)
+    conv1 = network.add_convolution_nd(input=input_tensor, num_output_maps=16, kernel_shape=(3, 3), kernel=conv1_w, bias=conv1_b)
+    relu1 = network.add_activation(input=conv1.get_output(0), type=trt.ActivationType.RELU)
+    conv2_w = np.random.randn(32, 16, 3, 3).astype(np.float32)
+    conv2_b = np.random.randn(32).astype(np.float32)
+    conv2 = network.add_convolution_nd(input=relu1.get_output(0), num_output_maps=32, kernel_shape=(3, 3), kernel=conv2_w, bias=conv2_b)
+    relu2 = network.add_activation(input=conv2.get_output(0), type=trt.ActivationType.RELU)
+    network.mark_output(tensor=relu2.get_output(0))
+    return builder, network
+
+class StreamWriter(trt.IStreamWriter):
+    def __init__(self):
+        trt.IStreamWriter.__init__(self)
+        self.bytes = bytes()
+
+    def write(self, data):
+        self.bytes += data
+        return len(data)
+
+def build_engine():
+    print("Constructing network...")
+    builder, network = build_network()
+    config = CreateConfig()(builder, network)
+    stream_writer = StreamWriter()
+    print("Building engine and serializing to stream...")
+    engine_bytes = builder.build_serialized_network_to_stream(network, config, stream_writer)
+    print("The total bytes written to stream is: ", len(stream_writer.bytes))
+    runtime = trt.Runtime(get_trt_logger())
+    print("Deserializing engine from stream...")
+    engine = runtime.deserialize_cuda_engine(stream_writer.bytes)
+    assert engine is not None, "Engine deserialization failed"
+    print("Engine deserialized successfully")
+
+if __name__ == "__main__":
+    build_engine()
+    if verbose is not None:
+        verbose.__exit__(None, None, None)
+

samples/python/strongly_type_autocast/README.md

@@ -0,0 +1,121 @@
+# Convert FP32 ONNX models to mixed FP32-FP16 precision for TensorRT strong typing usage
+
+**Table Of Contents**
+
+- [Description](#description)
+- [How does this sample work?](#how-does-this-sample-work)
+    * [Data preparation and the original ONNX model verification](#data-preparation-and-the-original-onnx-model-verification)
+    * [Model conversion and the converted ONNX model verification](#model-conversion-and-the-converted-onnx-model-verification)
+    * [Build and verify TensorRT engine from the converted ONNX model](#build-and-verify-tensorrt-engine-from-the-converted-onnx-model)
+- [Prerequisites](#prerequisites)
+- [Running the sample](#running-the-sample)
+    * [Sample `--help` options](#sample-help-options)
+- [Additional resources](#additional-resources)
+- [License](#license)
+- [Changelog](#changelog)
+- [Known issues](#known-issues)
+
+## Description
+
+This sample, `strongly_type_autocast`, uses ModelOpt's AutoCast tool to convert a FP32 ONNX model to mixed FP32-FP16 precision, and builds engine / runs inference with TensorRT's strong typing mode.
+
+[AutoCast](https://nvidia.github.io/TensorRT-Model-Optimizer/guides/8_autocast.html) is a tool for converting FP32 ONNX models to mixed precision FP32-FP16 or FP32-BF16 models. AutoCast intelligently selects nodes to keep in FP32 precision to maintain model accuracy while benefiting from reduced precision on the rest of the nodes. AutoCast automatically injects cast operations around the selected nodes.
+
+[Strong Typing vs Weak Typing](https://docs.nvidia.com/deeplearning/tensorrt/latest/architecture/capabilities.html#strong-vs-weak-typing)
+For strong typing, TensorRT adheres strictly to the type semantics in ONNX frameworks. For weak typing, TensorRT may substitute different precisions for tensors if it increases performance. Weak typing has been deprecated in 10.12. We recommend using AutoCast tool to convert the FP32 ONNX model to mixed precision before doing TensorRT strong typing optimization.
+
+## How does this sample work?
+
+This sample consists of three stages:
+- [Data preparation and the original ONNX model verification](#data-preparation-and-the-original-onnx-model-verification)
+- [Model conversion and the converted ONNX model verification](#model-conversion-and-the-converted-onnx-model-verification)
+- [Build and verify TensorRT engine from the converted ONNX model](#build-and-verify-tensorrt-engine-from-the-converted-onnx-model)
+
+### Data preparation and the original ONNX model verification
+
+The original input data is in pgm format, including ten pictures of numbers from 0 to 9. The input data need to be transformed to npz format.
+
+The original ONNX model is in fp32 precision. To verify the original model, use ONNX Runtime to run inference, printing the predicted numbers and saving all model outputs as gold references.
+
+### Model conversion and the converted ONNX model verification
+
+The original fp32 ONNX model can be converted to fp32-fp16 mixed precision using ModelOpt's AutoCast python package:
+
+    ```
+    from modelopt.onnx.autocast import convert_to_mixed_precision
+
+    converted_model = convert_to_mixed_precision(
+        onnx_path="model.onnx",
+        low_precision_type="fp16",            # or "bf16"
+        nodes_to_exclude=None,                # optional list of node name patterns to keep in FP32
+        op_types_to_exclude=None,             # optional list of op types to keep in FP32
+        data_max=512,                         # maximum absolute I/O values for nodes to convert
+        init_max=65504,                       # maximum absolute I/O values for initializers to convert
+        keep_io_types=False,                  # whether to preserve input/output types
+        calibration_data=None,                # optional path to input data file
+    )
+    ```  
+
+After the fp32-fp16 mixed precision model generated, use ONNX Runtime to run inference on the converted ONNX model, printing the predicted numbers and comparing the model outputs with gold references.
+
+### Build and verify TensorRT engine from the converted ONNX model
+
+Build TensorRT engine from the converted ONNX model, enabling strong type through `NetworkDefinitionCreationFlag.STRONGLY_TYPED`:
+
+    ```
+    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED))
+    ```
+
+Run inference on the generated trt engine, printing the predicted numbers and comparing the model outputs with gold references.
+
+## Prerequisites
+
+Dependencies required for this sample
+
+1. Install the dependencies for Python:
+    ```bash
+    pip3 install --upgrade pip
+    pip3 install -r requirements.txt
+    ```
+
+2. TensorRT
+
+3. The MNIST dataset can be found under the data directory (usually `/usr/src/tensorrt/data/mnist`) if using the TensorRT containers. It is also bundled along with the [TensorRT tarball](https://developer.nvidia.com/nvidia-tensorrt-download).
+
+## Running the sample
+
+1.  Run the sample:
+    `python3 sample.py [--mnist_dir] [--working_dir]`
+
+2.  Verify that the sample ran successfully. If the sample runs successfully you should see the following message:
+    ```
+    Sample finished successfully.
+    ```
+
+### Sample --help options
+
+To see the full list of available options and their descriptions, use the `-h` or `--help` command line option.
+
+# Additional resources
+
+The following resources provide a deeper understanding about AutoCast and TensorRT strong typing:
+
+**Documentation**
+- [Guide of TensorRT-Model-Optimizer Autocast](https://nvidia.github.io/TensorRT-Model-Optimizer/guides/8_autocast.html)
+- [TensorRT Strong Typing vs Weak Typing](https://docs.nvidia.com/deeplearning/tensorrt/latest/architecture/capabilities.html#strong-vs-weak-typing)
+- [Introduction To NVIDIA’s TensorRT Samples](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sample-support-guide/index.html#samples)
+- [Working With TensorRT Using The Python API](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#python_topics)
+- [NVIDIA’s TensorRT Documentation Library](https://docs.nvidia.com/deeplearning/sdk/tensorrt-archived/index.html)
+
+# License
+
+For terms and conditions for use, reproduction, and distribution, see the [TensorRT Software License Agreement](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sla/index.html) documentation.
+
+# Changelog
+
+September 2025
+This is the first version of this `README.md` file.
+
+# Known issues
+
+There are no known issues in this sample.

samples/python/strongly_type_autocast/requirements.txt

@@ -0,0 +1,9 @@
+--index-url https://download.pytorch.org/whl/cpu
+--extra-index-url https://pypi.org/simple
+Pillow==11.3.0
+torch==2.8.0
+nvidia-modelopt[onnx]==0.35.0
+pyyaml==6.0.3
+requests==2.32.4
+tqdm==4.66.4
+numpy==1.26.4

samples/python/strongly_type_autocast/sample.py

@@ -0,0 +1,193 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2026 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 sys
+import onnx
+import argparse
+import numpy as np
+from PIL import Image
+
+import onnxruntime
+import tensorrt as trt
+from modelopt.onnx.autocast import convert_to_mixed_precision
+from polygraphy.backend.onnxrt import OnnxrtRunner
+from polygraphy.backend.trt import TrtRunner
+
+class Sample:
+    def __init__(self, mnist_dir, working_dir):
+        if not os.path.exists(mnist_dir):
+            print(f"ERROR: {mnist_dir} is not existed.")
+            sys.exit(1)            
+        if not os.path.exists(working_dir):
+            print(f"ERROR: {working_dir} is not existed.")
+            sys.exit(1)   
+        self.mnist_dir = mnist_dir
+        self.working_dir = working_dir
+        self.origin_onnx_path = os.path.join(self.mnist_dir, "mnist-12.onnx")
+        self.converted_onnx_path = os.path.join(self.working_dir, "mnist_converted.onnx")
+        self.trt_engine_path = os.path.join(self.working_dir, "mnist_converted.engine")
+        self.origin_output_path = os.path.join(self.working_dir, "origin_output.npz")
+        self.converted_output_path = os.path.join(self.working_dir, "converted_output.npz")
+        self.trt_output_path = os.path.join(self.working_dir, "trt_output.npz")
+
+    # Prepare input data in npz format.
+    def prepare_input_npz(self):
+        print("Start prepare input data in npz format.")
+        for i in range(10):
+            image = Image.open(os.path.join(self.mnist_dir, str(i) + ".pgm"))
+            image_np = np.array(image)
+            image_np = 1.0 - (image_np / 255.0)
+            image_np_reshape = image_np.reshape(1, 1, 28, 28).astype(np.float32)
+            file_path = os.path.join(self.working_dir, str(i) + ".npz")
+            np.savez(file_path, Input3=image_np_reshape)
+            print(f"Saved input number {i} to {file_path}.")
+
+    # Convert fp32 onnx model to fp32-fp16 mixed precision
+    def convert_model(self):
+        fp32_nodes = ["Plus214"]
+        fp32_op_types = ["MatMul"]
+        print(f"Start autocast on {self.origin_onnx_path}.")
+        print("Use 8.npz as input data file for reference runner.")
+        calibration_path = os.path.join(self.working_dir, "8.npz")
+        converted_model = convert_to_mixed_precision(
+            onnx_path=self.origin_onnx_path,                        # Path to the input ONNX model.
+            low_precision_type="fp16",                              # Target precision to reduce to ('fp16' or 'bf16').
+            nodes_to_exclude=fp32_nodes,                            # List of regex patterns to match node names that should remain in FP32.
+            op_types_to_exclude=fp32_op_types,                      # List of operation types that should remain in FP32.
+            data_max=4.0,                                           # Maximum absolute value for node input and output values.
+            init_max=4.0,                                           # Maximum absolute value for initializers.
+            keep_io_types=True,                                     # Whether to preserve input/output types.
+            calibration_data=calibration_path,                      # Path to input data file for reference runner.
+        )
+        onnx.save(converted_model, self.converted_onnx_path)
+        print(f"Saved converted model to {self.converted_onnx_path}, autocast stage finished.")
+
+    # Infer with runner (onnxrt or trtrt)
+    def infer_with_runner(self, runner, output_file_name):
+        all_output = []
+        all_result = True
+        with runner:
+            for i in range(10):
+                input = np.load(os.path.join(self.working_dir, str(i) + ".npz"))
+                output = runner.infer(feed_dict=input)
+                all_output.append(output["Plus214_Output_0"].copy())
+                max_index = np.argmax(output["Plus214_Output_0"])
+                if max_index != i:
+                    all_result = False
+                result = "passed" if max_index == i else "failed"      
+                print(f"Input number is {i}, inferred number is {max_index}, {result}.")
+        if all_result:
+            print("Number infer test passed.")
+        else:
+            print("ERROR: Number infer test failed.")
+            sys.exit(1)
+        output_path = os.path.join(self.working_dir, output_file_name)
+        output_data = np.concatenate(all_output, axis=0)
+        np.savez(output_path, Plus214_Output_0=output_data)
+        print(f"Saved all outputs to {output_path}, infer stage finished.")    
+
+    # Check if outputs meet tolerances
+    def check_accuracy(self, output1_path, output2_path, r_tol=5e-3, a_tol=5e-3):
+        print(f" Comparing '{output1_path}' and '{output2_path}' with rtol({r_tol})+a_tol({a_tol})...")
+        output1 = np.load(output1_path)
+        output2 = np.load(output2_path)
+        result = np.allclose(output1["Plus214_Output_0"], output2["Plus214_Output_0"], rtol=r_tol, atol=a_tol)
+        if result:
+            print("Outputs and references matches.")
+        else:
+            print("ERROR: Outputs and references mismatches.")
+            sys.exit(1)
+
+    # Build trt engine with strongly_type
+    def build_trt_engine(self):
+        print(f"Start build trt engine on {self.converted_onnx_path}.")
+        TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
+        builder = trt.Builder(TRT_LOGGER)
+        network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED))
+        config = builder.create_builder_config()
+
+        parser = trt.OnnxParser(network, TRT_LOGGER)
+        with open(self.converted_onnx_path, "rb") as model:
+            if not parser.parse(model.read()):
+                for i in range(parser.num_errors):
+                    print(parser.get_error(i))
+                sys.exit(1)
+
+        engine = builder.build_engine_with_config(network, config)
+        with open(self.trt_engine_path, "wb") as f:
+            f.write(engine.serialize())
+            print(f"Saved trt engine to {self.trt_engine_path}.")
+
+    # Run onnxruntime infer stage
+    def infer_with_onnxrt(self, onnx_path, output_file_name):
+        print(f"Start onnxruntime infer on {onnx_path}.")
+        sess = onnxruntime.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
+        runner = OnnxrtRunner(sess)
+        self.infer_with_runner(runner, output_file_name)
+
+    # Run trtruntime infer stage
+    def infer_with_trt(self, output_file_name):
+        print(f"Start trtruntime infer on {self.trt_engine_path}.")
+        TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
+        runtime = trt.Runtime(TRT_LOGGER)
+        with open(self.trt_engine_path, 'rb') as f:
+            engine_data = f.read()
+            engine = runtime.deserialize_cuda_engine(engine_data)
+            runner = TrtRunner(engine)
+            self.infer_with_runner(runner, output_file_name)
+
+    def run(self):
+        print("STAGE 1: Prepare input data, do inference with onnx runtime on original fp32 onnx model.")
+        self.prepare_input_npz()
+        self.infer_with_onnxrt(self.origin_onnx_path, self.origin_output_path)
+
+        print("STAGE 2: Convert original fp32 onnx model to fp32-fp16 mixed onnx model, do inference with onnx runtime on converted onnx model, and check accuracy.")        
+        self.convert_model()
+        self.infer_with_onnxrt(self.converted_onnx_path, self.converted_output_path)
+        self.check_accuracy(self.origin_output_path, self.converted_output_path)
+
+        print("STAGE 3: Build trt engine on the converted onnx model with strong typing mode, do inference with trt runtime on the engine, and check accuracy.") 
+        self.build_trt_engine()
+        self.infer_with_trt(self.trt_output_path)
+        self.check_accuracy(self.origin_output_path, self.trt_output_path)
+        
+        print("Sample finished successfully.")
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Sample to convert FP32 ONNX model to mixed FP32-FP16 precision for TensorRT strong typing usage.",
+    )
+    parser.add_argument(
+        "--mnist_dir",
+        type=str,
+        required=True,
+        help="Path to the MNIST directory.",
+    )
+    parser.add_argument(
+        "--working_dir",
+        type=str,
+        default=os.path.abspath("."),
+        help="Path to the working directory. Defaults to current directory.",
+    )    
+    args = parser.parse_args()
+
+    sample = Sample(args.mnist_dir, args.working_dir)
+    sample.run()
+
+if __name__ == "__main__":
+    main()