Skip to main content

Mojo function

batched_matmul

batched_matmul[rank: Int, a_type: DType, b_type: DType, c_type: DType, //, *, transpose_a: Bool, transpose_b: Bool, elementwise_epilogue_fn: OptionalReg[elementwise_epilogue_type] = None, saturated_vnni: Bool = False, single_thread_blocking_override: Bool = False, target: StringSlice[StaticConstantOrigin] = "cpu"](c_buf: NDBuffer[c_type, rank, origin, shape, strides], a_buf: NDBuffer[a_type, rank, origin, shape, strides], b_buf: NDBuffer[b_type, rank, origin, shape, strides], *, context: DeviceContextPtr = DeviceContextPtr())

batched_matmul[rank: Int, a_type: DType, b_type: DType, c_type: DType, //, *, transpose_b: Bool, elementwise_epilogue_fn: OptionalReg[elementwise_epilogue_type] = None, saturated_vnni: Bool = False, target: StringSlice[StaticConstantOrigin] = "cpu"](c_buf: NDBuffer[c_type, rank, origin, shape, strides], a_buf: NDBuffer[a_type, rank, origin, shape, strides], b_buf: NDBuffer[b_type, rank, origin, shape, strides], *, context: DeviceContextPtr = DeviceContextPtr())

Was this page helpful?