MCPcopy
hub / github.com/NVIDIA/TensorRT-LLM / get_inputs

Method get_inputs

tensorrt_llm/graph_rewriting.py:29–51  ·  view source on GitHub ↗

Get the input tensors of the layer. Parameters: idx: the indices of the input tensor, will return all inputs if left empty Returns: List[Tensor]

(self, *indices: int)

Source from the content-addressed store, hash-verified

27 return self._network()
28
29 def get_inputs(self, *indices: int):
30 """Get the input tensors of the layer.
31
32 Parameters:
33 idx: the indices of the input tensor, will return all inputs if left empty
34
35 Returns:
36 List[Tensor]
37 """
38 from .functional import Tensor
39
40 indices = indices if indices else range(self.trt_layer.num_inputs)
41
42 ret = []
43 for i in indices:
44 assert i < self.trt_layer.num_inputs, (
45 f"Invalid input index {i} for layer {self.trt_layer.name}"
46 )
47
48 tensor = self.trt_layer.get_input(i)
49 tensor = Tensor(trt_tensor=tensor, network=self.network, is_network_input=False)
50 ret.append(tensor)
51 return ret
52
53 def get_outputs(self, *indices: int):
54 """Get the output tensor of the layer.

Callers 2

const_foldableMethod · 0.45
get_eltwise_inputsMethod · 0.45

Calls 3

TensorClass · 0.70
get_inputMethod · 0.45
appendMethod · 0.45

Tested by

no test coverage detected