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

flatten_leading

def flatten_leading[dtype: DType, layout: TensorLayout, //](tensor: TileTensor[dtype, layout, Storage=tensor.Storage, address_space=tensor.address_space, linear_idx_type=tensor.linear_idx_type]) -> TileTensor[dtype, Layout[*?, *?], tensor.origin, Storage=tensor.Storage, address_space=tensor.address_space]

Merge the first two dimensions of a rank-3 TileTensor: (A, B, C) -> (A*B, C).

Returns a new TileTensor sharing the same pointer with row-major strides computed from the merged shape. Zero-cost operation.

Common use case: converting 3D batched tensors (num_experts, N, K) to 2D (num_experts*N, K) for TMA descriptor creation in MoE kernels.

Parameters:

• ​dtype (DType): Element type (inferred from tensor).
• ​layout (TensorLayout): Layout type (inferred from tensor).

Args:

• ​tensor (TileTensor[dtype, layout, Storage=tensor.Storage, address_space=tensor.address_space, linear_idx_type=tensor.linear_idx_type]): A rank-3 TileTensor.

Returns:

TileTensor[dtype, Layout[*?, *?], tensor.origin, Storage=tensor.Storage, address_space=tensor.address_space]: A rank-2 TileTensor where dim[0] = old dim[0] * dim[1].