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

pack_matmul_b_shape_func_shape

def pack_matmul_b_shape_func_shape[a_type: DType, a_shape: IntTuple, b_type: DType, b_shape: IntTuple, c_type: DType, c_shape: IntTuple, transpose_in_0: Bool](b_input: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=b_input.static_spec]) -> IndexList[Int(2)]

Computes the output shape for the pack_matmul_b_shape_func graph op.

Parameters:

a_type (DType): Element type of the A (activation) operand of the matmul the packed B will be used in.
a_shape (IntTuple): Static shape of the A operand; a_shape[0] is the M dimension used to select the matmul kernel variant (UNKNOWN_VALUE for dynamic M).
b_type (DType): Element type of the B (weight) operand being packed.
b_shape (IntTuple): Static shape of the B operand.
c_type (DType): Element type of the C (output) operand of the matmul the packed B will be used in.
c_shape (IntTuple): Static shape of the C (output) operand.
transpose_in_0 (Bool): True if the B operand is transposed, stored as [N, K] instead of [K, N].

Args:

b_input (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=b_input.static_spec]): Rank-2 B input tensor whose packed output shape is computed.

Returns:

IndexList[Int(2)]: The packed output shape for the B operand.