Function identity

tensorrt_llm/functional.py:3278–3300

Add an identity operation. Parameters: input : Tensor The input tensor. Returns: The tensor produced by this identity operation.

(input: Tensor)

def identity(input: Tensor) -> Tensor:
    '''
    Add an identity operation.

    Parameters:
        input : Tensor
            The input tensor.

    Returns:
        The tensor produced by this identity operation.
    '''
    if not default_net().plugin_config.identity_plugin:
        layer = default_trtnet().add_identity(input.trt_tensor)
    else:
        plg_creator = trt.get_plugin_registry().get_plugin_creator(
            'Identity', '1', TRT_LLM_PLUGIN_NAMESPACE)
        assert plg_creator is not None
        pfc = trt.PluginFieldCollection()
        id_plug = plg_creator.create_plugin("identity", pfc)
        plug_inputs = [input.trt_tensor]
        layer = default_trtnet().add_plugin_v2(plug_inputs, id_plug)
        _add_plugin_info(layer, plg_creator, "identity", pfc)
    return _create_tensor(layer.get_output(0), layer)


def argmax(input: Tensor, dim: int, keepdim: bool = False) -> Tensor: