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

def batched_matmul[*, transpose_a: Bool = False, transpose_b: Bool = False, elementwise_epilogue_fn: Optional[def[c_type: DType, width: SIMDLength, rank: Int, *, alignment: Int = Int(1)](IndexList[rank], SIMD[c_type, width]) capturing thin -> None] = None, saturated_vnni: Bool = False, target: StringSlice[ImmStaticOrigin] = StringSlice("cpu")](c_buf: TileTensor[Storage=c_buf.Storage, linear_idx_type=c_buf.linear_idx_type], a_buf: TileTensor[Storage=a_buf.Storage, linear_idx_type=a_buf.linear_idx_type], b_buf: TileTensor[Storage=b_buf.Storage, linear_idx_type=b_buf.linear_idx_type], *, context: Optional[DeviceContext] = None)

TileTensor primary implementation of batched_matmul.

Parameters:

transpose_a (Bool): Whether the LHS input is transposed (defaults to False; not yet supported).
transpose_b (Bool): Whether the RHS input is transposed (defaults to False).
elementwise_epilogue_fn (Optional[def[c_type: DType, width: SIMDLength, rank: Int, *, alignment: Int = Int(1)](IndexList[rank], SIMD[c_type, width]) capturing thin -> None]): Optional epilogue applied to each output element (defaults to None).
saturated_vnni (Bool): Whether to use saturated VNNI accumulation on CPU (defaults to False; not applicable on GPU).
target (StringSlice[ImmStaticOrigin]): Target hardware for the operation (defaults to "cpu").

Args:

c_buf (TileTensor[Storage=c_buf.Storage, linear_idx_type=c_buf.linear_idx_type]): Output tensor of shape (..., m, n); rank must be at least 2 and match the inputs.
a_buf (TileTensor[Storage=a_buf.Storage, linear_idx_type=a_buf.linear_idx_type]): LHS input tensor of shape (..., m, k).
b_buf (TileTensor[Storage=b_buf.Storage, linear_idx_type=b_buf.linear_idx_type]): RHS input tensor of shape (..., k, n), or (..., n, k) when transpose_b is set.
context (Optional[DeviceContext]): Optional device context used for dispatch and parallelism (defaults to None).