plugin

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

>>> inp
tensor(
    [[[1.0000, 1.0000, 1.0000, 1.0000]],

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

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