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

Function pad

tensorrt_llm/functional.py:1327–1398  ·  view source on GitHub ↗

Add a pad layer. The padding layer adds zero-padding at the start and end of the input tensor. And the padding size by which to pad some dimensions of input are described starting from the last dimension and moving forward. `[len(pad) / 2]` dimensions of input will be padded.

(input: Tensor,
        pad: Union[Sequence[int], Tensor],
        mode: str = 'constant',
        value: Optional[float] = None)

Source from the content-addressed store, hash-verified

1325
1326
1327def pad(input: Tensor,
1328 pad: Union[Sequence[int], Tensor],
1329 mode: str = 'constant',
1330 value: Optional[float] = None) -> Tensor:
1331 '''
1332 Add a pad layer.
1333
1334 The padding layer adds zero-padding at the start and end of the input tensor. And the
1335 padding size by which to pad some dimensions of input are described starting from the
1336 last dimension and moving forward.
1337
1338 `[len(pad) / 2]` dimensions of input will be padded. For example, to pad only the last
1339 dimension of the input tensor, then pad has the form [padding_left, padding_right]; to
1340 pad the last 2 dimensions of the input tensor, then use [padding_left, padding_right,
1341 padding_top, padding_bottom]; to pad the last 3 dimensions, use [padding_left,
1342 padding_right, padding_top, padding_bottom, padding_front, padding_back].
1343
1344 Parameters:
1345 input : Tensor
1346 The input tensor on which the padding_2d is performed.
1347 pad : sequence of int
1348 An m-elements tuple for padding, where its length m meets the requirement that
1349 m <= 2*input dimensions, and m is even.
1350 mode : str
1351 Only \'constant\' is supported.
1352 value : float
1353 Fill value for 'constant' padding. Default: 0.
1354
1355 Returns:
1356 The tensor produced by the inserted layer.
1357 ''&#x27;
1358 assert mode == "constant", "Only `'constant'` is supported now."
1359
1360 if isinstance(pad, list) or isinstance(pad, tuple):
1361 assert (
1362 len(pad) % 2 == 0 and len(pad) <= 2 * input.ndim()
1363 ), "The length of `pad` should be even and less than 2*input.ndim"
1364 pad = constant(np.array(pad).astype(np.int32)).view([-1, 2])
1365 elif isinstance(pad, Tensor):
1366 pad = pad.flatten()
1367 assert (
1368 pad.size(0) % 2 == 0 and pad.size(0) <= 2 * input.ndim()
1369 ), "The length of `pad` should be even and less than 2*input.ndim"
1370 pad = pad.cast("int32").view([-1, 2])
1371 else:
1372 raise NotImplementedError(f"pad type {type(pad)} not supported")
1373 if value is None:
1374 value = 0
1375
1376 pad = concat([constant(np.zeros((1, 2), dtype=np.int32)),
1377 pad]) # pre-padding the indices
1378 padding_index = [0] * input.ndim()
1379 padding_index[-(pad.size(0) - 1):] = list(range(pad.size(0) - 1, 0,
1380 -1)) # reverse the indices
1381 pad = index_select(pad,
1382 dim=0,
1383 index=constant(np.array(padding_index, dtype=np.int32)))
1384 pre_padding, post_padding = chunk(pad, chunks=2, dim=1)

Callers 7

dynamic_preprocessFunction · 0.90
get_timestep_embeddingFunction · 0.50
fillMethod · 0.50
forwardMethod · 0.50
timestep_embeddingMethod · 0.50
forwardMethod · 0.50

Calls 14

constantFunction · 0.85
concatFunction · 0.85
index_selectFunction · 0.85
chunkFunction · 0.85
default_trtnetFunction · 0.85
constant_to_tensor_Function · 0.85
_create_tensorFunction · 0.85
flattenMethod · 0.80
castMethod · 0.80
shapeFunction · 0.70
ndimMethod · 0.45
viewMethod · 0.45

Tested by

no test coverage detected