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
 1inp = tp.iota((2, 1, 4))
 2out = tp.plugin(
 3    name="CustomGeluPluginDynamic",
 4    inputs=[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, 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).

Parameters:
  • op – The ID of plugin to call, in the form "<namespace>::<name>".

  • inputs – The inputs to the plugin.

  • **kwargs (Dict[str, Any]) – Additional arguments to pass to the plugin as attributes. Supported attribute types are documented here.

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]