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
andIPluginV3
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 expectDims
should be provided as atuple
ofint
s. Fields that expect multiple values can be provided aslist
s ortuple
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:
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: