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: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=input.static_spec], _grid_thws: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=_grid_thws.static_spec], _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 (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=input.static_spec]): Rank-2 input tensor of shape (n_tokens, D); the second dimension is preserved as the output's second dimension.
• ​_grid_thws (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=_grid_thws.static_spec]): 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): Maximum H across all videos, used for launch grid sizing.
• ​_max_w (Int32): Maximum W across 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:

IndexList[Int(2)]