Function constant_to_tensor_

tensorrt_llm/functional.py:2854–2881

(input: Union[Tensor, int, float, bool],
                        dtype: Union[trt.DataType, str] = None,
                        to_array=True)

def constant_to_tensor_(input: Union[Tensor, int, float, bool],
                        dtype: Union[trt.DataType, str] = None,
                        to_array=True) -> Tensor:
    if dtype is None:
        # deduce the type from the given value
        # NOTE: bool is a subtype of int, so bool needs to be checked first
        if isinstance(input, bool):
            dtype = trt.bool
        elif isinstance(input, int):
            dtype = trt.int32
        else:
            dtype = trt.float32

    if not isinstance(input, Tensor):
        if isinstance(dtype, str):
            dtype = str_dtype_to_trt(dtype)
        array_fn_dict = {
            trt.int64: int64_array,
            trt.int32: int32_array,
            trt.float32: fp32_array,
            trt.float16: fp16_array,
            trt.bfloat16: bf16_array,
            trt.bool: bool_array,
        }
        assert dtype in array_fn_dict
        return constant(array_fn_dict[dtype]([input] if to_array else input))

    return input


def constants_to_tensors_(