For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).

Mojo function

repeat_interleave_shape

def repeat_interleave_shape[type_repeats: DType](input: TileTensor[Storage=input.Storage, address_space=input.address_space, linear_idx_type=input.linear_idx_type], repeats: TileTensor[type_repeats, Storage=repeats.Storage, address_space=repeats.address_space, linear_idx_type=repeats.linear_idx_type], axis: Int) -> IndexList[input.LayoutType.rank]

Computes the output shape of repeat_interleave for the given input, repeats, and axis.

The returned IndexList matches input's rank with the size along axis replaced by the summed repeats values, or by input[axis] multiplied by the single repeat value when repeats is a size-1 broadcast.

Args:

input (TileTensor[Storage=input.Storage, address_space=input.address_space, linear_idx_type=input.linear_idx_type]): The input tensor whose shape is being transformed.
repeats (TileTensor[type_repeats, Storage=repeats.Storage, address_space=repeats.address_space, linear_idx_type=repeats.linear_idx_type]): A one-dimensional integral tensor of per-element repeat counts, either size 1 or equal to input.dim(axis).
axis (Int): The axis along which elements are repeated.

Returns:

IndexList[input.LayoutType.rank]: An IndexList matching input.rank with the axis dimension updated to the total repeated size.