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

Function meshgrid2d

tensorrt_llm/functional.py:6341–6364  ·  view source on GitHub ↗

Creates grids (2D) of coordinates specified by the 1D inputs (only supports `indexing=\'xy\'`). Parameters: x : Tensor The first input (1D) tensor. y : Tensor The second input (1D) tensor. Returns: The tuple of two tensors produced.

(x: Tensor, y: Tensor)

Source from the content-addressed store, hash-verified

6339
6340
6341def meshgrid2d(x: Tensor, y: Tensor) -> Tuple[Tensor]:
6342 '''
6343 Creates grids (2D) of coordinates specified by the 1D inputs (only supports `indexing=\'xy\'`).
6344
6345 Parameters:
6346 x : Tensor
6347 The first input (1D) tensor.
6348 y : Tensor
6349 The second input (1D) tensor.
6350
6351 Returns:
6352 The tuple of two tensors produced.
6353
6354 TODO: Add full support for torch.meshgrid.
6355 See https://pytorch.org/docs/stable/generated/torch.meshgrid.html#torch-meshgrid
6356 '''
6357 if x.ndim() == 1:
6358 x = expand_dims(x, 0)
6359 if y.ndim() == 1:
6360 y = expand_dims(y, 0)
6361 grid_x = repeat_interleave(x, shape(y, 1),
6362 1).view([x.shape[-1], y.shape[-1]])
6363 grid_y = repeat(y, (x.shape[-1], 1))
6364 return (grid_x, grid_y)
6365
6366
6367def generate_logn_scaling(seq_length: int = 8192,

Callers 2

get_2d_sincos_pos_embedFunction · 0.85
_get_cached_embMethod · 0.85

Calls 6

expand_dimsFunction · 0.85
repeat_interleaveFunction · 0.85
repeatFunction · 0.85
shapeFunction · 0.70
ndimMethod · 0.45
viewMethod · 0.45

Tested by

no test coverage detected