Mojo function
dispatch_materialized_mask_and_score_mod
dispatch_materialized_mask_and_score_mod[dtype: DType, layout: Layout, //, score_mod_type: String, callback_fn: fn[mask_t: MHAMask, score_mod_t: ScoreModTrait](mask: mask_t, score_mod: score_mod_t) raises capturing -> None, num_heads: Int = -1](mask_nd: LayoutTensor[dtype, layout, MutableAnyOrigin], start_pos_nd: OptionalReg[LayoutTensor[DType.uint32, Layout.row_major(-1), MutableAnyOrigin]] = None)
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!