Masks for block matrix

Author: gsvgit

cfpq_matrix/block/block_matrix.py

@@ -24,10 +24,7 @@ def optimize_similarly(self, other: Matrix) -> "OptimizedMatrix":
         )
 
     def to_mask(self):
-        if self.block_matrix_space.is_single_cell(self.shape):
-            return self.base.to_mask()
-        else:
-            return None
+        return self.base.to_mask()
 
 
 class CellBlockMatrix(BlockMatrix):
@@ -38,6 +35,8 @@ def __init__(self, base: OptimizedMatrix, block_matrix_space: BlockMatrixSpace):
     def mxm(self, other: Matrix, op: Semiring, mask:Matrix, swap_operands: bool = False) -> Matrix:
         if self.block_matrix_space.is_single_cell(other.shape):
             return self.base.mxm(other, op, mask, swap_operands=swap_operands)
+        if not mask is None:
+            mask = self.block_matrix_space.repeat_into_hyper_column(mask)
         return self.base.mxm(
             self.block_matrix_space.hyper_rotate(
                 other,
@@ -92,6 +91,8 @@ def _force_init_orientation(
 
     def mxm(self, other: Matrix, op: Semiring, mask:Matrix, swap_operands: bool = False) -> Matrix:
         if self.block_matrix_space.is_single_cell(other.shape):
+            if not mask is None:
+                mask = self.block_matrix_space.hyper_rotate(self.block_matrix_space.repeat_into_hyper_column(mask),BlockMatrixOrientation.HORIZONTAL)
             return self._force_init_orientation(
                 BlockMatrixOrientation.HORIZONTAL
                 if swap_operands