Mojo function

zipped_divide

zipped_divide[LayoutType: TensorLayout, //, tile: Coord[tile.element_types]](layout: LayoutType) -> Layout[Coord[tile.element_types], Coord[#kgen.variadic.reduce(LayoutType._shape_types, base=, reducer=[PrevV: Variadic[CoordLike], VA: Variadic[CoordLike], idx: __mlir_type.index] #kgen.variadic.concat(PrevV, ComptimeInt[(VA[idx].static_value // tile.element_types[idx].static_value)]))], Coord[LayoutType._stride_types], Coord[#kgen.variadic.reduce(LayoutType._stride_types, base=, reducer=[PrevV: Variadic[CoordLike], VA: Variadic[CoordLike], idx: __mlir_type.index] #kgen.variadic.concat(PrevV, ComptimeInt[(VA[idx].static_value * tile.element_types[idx].static_value)]))]]

Divides a layout into inner (tile) and outer (number-of-tiles) parts.

Given a layout and a tile shape, produces a hierarchical layout where the inner component has the tile shape with the original strides, and the outer component has shape = original_shape / tile with scaled strides.

Parameters:

LayoutType (TensorLayout): The type of the input layout.
tile (Coord): The tile shape to divide by.

Args:

layout (LayoutType): The layout to divide.

Returns:

Layout: A ZippedDivideLayout with inner and outer components.