MCPcopy Index your code
hub / github.com/NVIDIA/TensorRT-LLM / __init__

Method __init__

tensorrt_llm/parameter.py:37–65  ·  view source on GitHub ↗
(self,
                 value: Optional[Union[np.ndarray, torch.Tensor]] = None,
                 shape: Sequence[int] = None,
                 dtype: Union[str, trt.DataType] = None,
                 is_buffer: bool = False,
                 prefer_managed=False)

Source from the content-addressed store, hash-verified

35 _DEFAULT_DTYPE = trt.DataType.FLOAT
36
37 def __init__(self,
38 value: Optional[Union[np.ndarray, torch.Tensor]] = None,
39 shape: Sequence[int] = None,
40 dtype: Union[str, trt.DataType] = None,
41 is_buffer: bool = False,
42 prefer_managed=False):
43 if dtype is None:
44 logger.warning(
45 f'Parameter dtype is None, using default dtype: {self._DEFAULT_DTYPE}, it is recommended to always specify dtype explicitly'
46 )
47 dtype = self._DEFAULT_DTYPE if dtype is None else dtype
48 if isinstance(dtype, str):
49 dtype = str_dtype_to_trt(dtype)
50 self._dtype: trt.DataType = dtype
51 if value is None:
52 assert isinstance(shape, (
53 list,
54 tuple)), f"shape must be list or tuple, receive {(type(shape))}"
55 self._shape = tuple(shape)
56 self._value = None
57 else:
58 self._shape = value.shape
59 self._value = self._regularize_value(value)
60 self.is_buffer = is_buffer
61 self._prefer_managed = prefer_managed
62 self._tensor: Tensor = None
63 self._network: weakref.ref = None
64 self._name = None
65 self.need_transpose = False
66
67 @property
68 def shape(self):

Callers

nothing calls this directly

Calls 3

_regularize_valueMethod · 0.95
str_dtype_to_trtFunction · 0.85
warningMethod · 0.45

Tested by

no test coverage detected