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

batched_matmul_shape

def batched_matmul_shape[rank: Int](a_buff: TileTensor[address_space=a_buff.address_space, linear_idx_type=a_buff.linear_idx_type, element_size=a_buff.element_size], b_buff: TileTensor[address_space=b_buff.address_space, linear_idx_type=b_buff.linear_idx_type, element_size=b_buff.element_size]) -> IndexList[rank]

Compute the output shape of a batch_matmul operation, and assert the inputs are compatible.

Parameters:

rank (Int): Rank of the input and output tensors.

Args:

a_buff (TileTensor[address_space=a_buff.address_space, linear_idx_type=a_buff.linear_idx_type, element_size=a_buff.element_size]): The lhs input tensor.
b_buff (TileTensor[address_space=b_buff.address_space, linear_idx_type=b_buff.linear_idx_type, element_size=b_buff.element_size]): The rhs input tensor.

Returns:

IndexList[rank]: The output shape.