plugin¶

nvtripy.plugin(*args, **kwargs: Dict[str, Any]) → Tensor | List[Tensor][source]¶

This function has multiple overloads:

nvtripy.plugin(name: str, inputs: Sequence[nvtripy.Tensor], output_info: List[Tuple[int, nvtripy.dtype]], version: str = 1, namespace: str = , kwargs: Dict[str, typing.Any]) -> Tensor | List[Tensor]

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

Parameters:

name – The name of the plugin to call.
inputs – The inputs to the plugin.
output_info – A list of tuples that indicate the rank and data type for each output.
version – The version of the plugin to call.
namespace – The namespace of the plugin.
**kwargs (Dict[str, Any]) – 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

inp = tp.iota((2, 1, 4))
out = tp.plugin(
    name="CustomGeluPluginDynamic",
    inputs=[inp],
    # GELU has a single output which always has the same rank and data
    # type as the input.
    output_info=[(inp.rank, inp.dtype)],
    # The GELU plugin expects a `type_id` parameter indicating the precision
    # to use. `0` indicates float32.
    type_id=0,
)

Local Variables¶

>>> inp
tensor(
    [[[0, 0, 0, 0]],

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

>>> out
tensor(
    [[[0, 0, 0, 0]],

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

nvtripy.plugin(op: str, inputs: Sequence[nvtripy.Tensor], kwargs: Dict[str, typing.Any]) -> nvtripy.Tensor | List[nvtripy.Tensor]

Calls a TensorRT quickly deployable plugin (QDP).