Mojo function
dispatch_materialized_mask_and_score_mod
dispatch_materialized_mask_and_score_mod[dtype: DType, layout: Layout, //, score_mod_type: String, callback_fn: callback_fn_type, num_heads: Int = -1](mask_nd: LayoutTensor[dtype, layout, MutAnyOrigin], start_pos_nd: OptionalReg[LayoutTensor[DType.uint32, Layout.row_major(-1), MutAnyOrigin]] = None)
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!