Mojo function

test_warp_specialize_gemm

test_warp_specialize_gemm[a_type: DType, b_type: DType, c_type: DType, cluster_shape: IndexList[3], block_tile_shape: IndexList[3], wgmma_shape: IndexList[3], num_consumer: Int = 1, num_pipeline_stages: Int = 4, transpose_b: Bool = True, partitioned_multicast: Bool = False, grid_shape: OptionalReg[IndexList[2]] = OptionalReg[IndexList[2]]({:i1 0, 1}), use_tma_store: Bool = False, schedule: MatmulSchedule = MatmulSchedule(-1), default_epilogue: Bool = False, elementwise_compute_lambda_fn: OptionalReg[fn[DType, Int, Int](IndexList[2], SIMD[$0, $1]) capturing -> SIMD[$0, $1]] = OptionalReg[fn[DType, Int, Int](IndexList[2], SIMD[$0, $1]) capturing -> SIMD[$0, $1]]({:i1 0, 1}), measure_threshold: OptionalReg[SIMD[float64, 1]] = OptionalReg[SIMD[float64, 1]]({:i1 0, 1})](ctx: DeviceContext, m: ValOrDim[dim], n: ValOrDim[dim], k: ValOrDim[dim], handle: Optional[Handle] = Optional(None))