For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).
Mojo function
tpool_patch_merger_shape
def tpool_patch_merger_shape(input: T, _grid_thws: T, _kH: Int32, _kW: Int32, _max_h: Int32, _max_w: Int32, total_output_patches: Int32) -> IndexList[Int(2)]
Computes the output shape for the tpool_patch_merger graph op.
Args:
- input (
T): Rank-2 input tensor of shape(n_tokens, D); the second dimension is preserved as the output's second dimension. - _grid_thws (
T): Rank-2 grid dimensions tensor of shape(n_vids, 3)holding the(T, H, W)triple per video. - _kH (
Int32): Merge kernel height, in elements. - _kW (
Int32): Merge kernel width, in elements. - _max_h (
Int32): MaximumHacross all videos, used for launch grid sizing. - _max_w (
Int32): MaximumWacross all videos, used for launch grid sizing. - total_output_patches (
Int32): Total number of output patches across all videos; becomes the first dimension of the output shape.
Returns: