Mojo function

batched_matmul_shape

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.

Was this page helpful?

Thank you! We'll create more content like this.

Thank you for helping us improve!

😔 What went wrong?

Some code doesn’t work

It includes inaccurate information

It's missing information I need

It was difficult to understand

Other

Submit