plugin

nvtripy.plugin(name: str, inputs: Sequence[Tensor], output_info: List[Tuple[int, dtype]], version: str = '1', namespace: str = '', **kwargs) → Tensor | List[Tensor][source]

Calls a TensorRT plugin. Only the IPluginV2DynamicExt and IPluginV3 interfaces are supported.

Parameters:

  • name (str) – The name of the plugin to call.

  • inputs (Sequence[Tensor]) – The inputs to the plugin.

  • output_info (List[Tuple[int, dtype]]) – A list of tuples that indicate the rank and data type for each output.

  • version (str) – The version of the plugin to call.

  • namespace (str) – The namespace of the plugin.

  • **kwargs – Additional arguments to pass to the plugin as plugin fields. These should be primitive Python types like int s, float s, str s etc. Fields that expect Dims should be provided as a tuple of int s. Fields that expect multiple values can be provided as list s or tuple s.

Returns:

The output(s) of the plugin either as a single tensor if there is only one output, or a list of tensors otherwise.

Return type:

Tensor | List[Tensor]

Example
 1inp = tp.iota((2, 1, 4))
 2out = tp.plugin(
 3    "CustomGeluPluginDynamic",
 4    [inp],
 5    # GELU has a single output which always has the same rank and data
 6    # type as the input.
 7    output_info=[(inp.rank, inp.dtype)],
 8    # The GELU plugin expects a `type_id` parameter indicating the precision
 9    # to use. `0` indicates float32.
10    type_id=0,
11)
Local Variables
>>> inp
tensor(
    [[[0.0000, 0.0000, 0.0000, 0.0000]],

     [[1.0000, 1.0000, 1.0000, 1.0000]]], 
    dtype=float32, loc=gpu:0, shape=(2, 1, 4))

>>> out
tensor(
    [[[0.0000, 0.0000, 0.0000, 0.0000]],

     [[0.8412, 0.8412, 0.8412, 0.8412]]], 
    dtype=float32, loc=gpu:0, shape=(2, 1, 4))