BlockMaskedMM¶

Matrix multiplication with block masking.

Abstract Signature:

BlockMaskedMM(a: Tensor, b: Tensor, block_size: int = 64, mask_out: Tensor, mask_lhs: Tensor, mask_rhs: Tensor)

API: torch.sparse.bsr_spmm

Strategy: Direct Mapping

API: mlx.core.block_masked_mm

Strategy: Direct Mapping