Mojo function

expand_modes_alike

expand_modes_alike(shape_a: IntTuple[origin], stride_a: IntTuple[origin], shape_b: IntTuple[origin], stride_b: IntTuple[origin]) -> InlineArray[IntTuple, 3]

Aligns two shape-stride pairs to have the same hierarchical structure.

This function is used to make two layouts compatible for operations by ensuring they have the same hierarchical structure, expanding scalar values into tuples as needed.

Args:

shape_a (IntTuple[origin]): The first shape tuple.
stride_a (IntTuple[origin]): The first stride tuple.
shape_b (IntTuple[origin]): The second shape tuple.
stride_b (IntTuple[origin]): The second stride tuple.

Returns:

An array containing three tuples: the common shape, the expanded stride_a, and the expanded stride_b.

expand_modes_alike(layout_a: Layout, layout_b: Layout) -> InlineArray[Layout, 2]

Aligns two layouts to have the same hierarchical structure.

This function tiles both layouts so they mirror each other's structure, making them compatible for operations that require matching hierarchies.

Example:

Given layouts with different structures:

layout_0: (((3, (5, 2)), 4):((1, (24, 12)), 3))
layout_1: ((30, (2, 2)):(2, (60, 1)))

The result would be two layouts with matching structures:

(((3, (5, 2)), (2, 2)):((1, (24, 12)), (3, 6)))
(((3, (5, 2)), (2, 2)):((2, (6, 30)), (60, 1)))

from layout import Layout, IntTuple
from layout.layout import expand_modes_alike

alias layout_0 = Layout(
    IntTuple(IntTuple(3, IntTuple(5, 2)), 4),
    IntTuple(IntTuple(1, IntTuple(24, 12)), 3),
)
alias layout_1 = Layout(
    IntTuple(30, IntTuple(2, 2)), IntTuple(2, IntTuple(60, 1))
)
alias uc = expand_modes_alike(layout_0, layout_1)
print(uc[0])
# (((3, (5, 2)), (2, 2)):((1, (24, 12)), (3, 6)))
print(uc[1])
# (((3, (5, 2)), (2, 2)):((2, (6, 30)), (60, 1)))
from layout import Layout, IntTuple
from layout.layout import expand_modes_alike

alias layout_0 = Layout(
    IntTuple(IntTuple(3, IntTuple(5, 2)), 4),
    IntTuple(IntTuple(1, IntTuple(24, 12)), 3),
)
alias layout_1 = Layout(
    IntTuple(30, IntTuple(2, 2)), IntTuple(2, IntTuple(60, 1))
)
alias uc = expand_modes_alike(layout_0, layout_1)
print(uc[0])
# (((3, (5, 2)), (2, 2)):((1, (24, 12)), (3, 6)))
print(uc[1])
# (((3, (5, 2)), (2, 2)):((2, (6, 30)), (60, 1)))

Args:

layout_a (Layout): The first layout to align.
layout_b (Layout): The second layout to align.

Returns:

An array containing two layouts with matching hierarchical structures.