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

Function view

tensorrt_llm/functional.py:1745–1790  ·  view source on GitHub ↗

Add an operation to create a view of a tensor. That operation adds a tensorrt.IShuffleLayer to the network. If the 'shape' parameter is a Tensor, that view is dynamic. Otherwise, it is a static view. Note that TensorRT limits the number of inferred dimensions to 1. It means

(input: Tensor,
         shape: Union[Tensor, Sequence[int]],
         zero_is_placeholder: bool = True)

Source from the content-addressed store, hash-verified

1743
1744
1745def view(input: Tensor,
1746 shape: Union[Tensor, Sequence[int]],
1747 zero_is_placeholder: bool = True) -> Tensor:
1748 '''
1749 Add an operation to create a view of a tensor.
1750
1751 That operation adds a tensorrt.IShuffleLayer to the network. If the 'shape'
1752 parameter is a Tensor, that view is dynamic. Otherwise, it is a static
1753 view.
1754
1755 Note that TensorRT limits the number of inferred dimensions to 1. It means
1756 that the shape sequence or tensor cannot contain more than one -1. This
1757 function enforces that constraint and will assert if it is not respected.
1758
1759 Parameters:
1760 input : Tensor
1761 The input tensor to transpose.
1762
1763 shape : Union[Tensor, Sequence[int]]
1764 The shape of the new tensor.
1765
1766 zero_is_placeholder : bool
1767 When that parameter is True, the 0s in 'shape' are replaced by the
1768 sizes of the corresponding dimensions from the 'input'. Otherwise,
1769 the dimensions corresponding to 0s are shrunk.
1770
1771 Returns:
1772 The tensor produced by the view/shuffle layer.
1773 '''
1774
1775 # TensorRT demands that at most one dimension is permitted to be specified as -1
1776 def assert_no_more_than_one_inferred_dim(list):
1777 inferred_dim_list = [i for i in list if i == -1]
1778 assert len(inferred_dim_list) <= 1
1779
1780 layer = default_trtnet().add_shuffle(input.trt_tensor)
1781 layer.zero_is_placeholder = zero_is_placeholder
1782 if isinstance(shape, Tensor):
1783 assert_no_more_than_one_inferred_dim(shape.shape)
1784 layer.set_input(1, shape.trt_tensor)
1785 elif isinstance(shape, (list, tuple)):
1786 assert_no_more_than_one_inferred_dim(shape)
1787 layer.reshape_dims = tuple(shape)
1788 else:
1789 raise TypeError("%s is not supported" % type(shape))
1790 return _create_tensor(layer.get_output(0), layer)
1791
1792
1793def flatten(input: Tensor, start_dim: int = 0, end_dim: int = -1):

Callers 15

_unpack_beamsFunction · 0.90
_get_draft_token_indicesFunction · 0.90
_get_draft_token_arrayFunction · 0.90
_gather_beamsFunction · 0.90
_flatten_decoding_dimFunction · 0.90
_unflatten_decoding_dimFunction · 0.90
_beam_search_candidatesFunction · 0.90
_unpack_gen_dataFunction · 0.90
viewMethod · 0.85
__init__Method · 0.85
flattenFunction · 0.85

Calls 4

default_trtnetFunction · 0.85
_create_tensorFunction · 0.85
get_outputMethod · 0.45

Tested by

no test coverage detected