From 7a16e5a23c0597ceda41a7077567ebbbba8ad9af Mon Sep 17 00:00:00 2001 From: LeiWang1999 Date: Wed, 7 Aug 2024 02:37:20 +0000 Subject: [PATCH] bug fix --- bitblas/ops/general_matmul/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bitblas/ops/general_matmul/__init__.py b/bitblas/ops/general_matmul/__init__.py index c63a33bd3..d384ee83d 100644 --- a/bitblas/ops/general_matmul/__init__.py +++ b/bitblas/ops/general_matmul/__init__.py @@ -128,7 +128,7 @@ def __initialize_propagate(self, propagate_a: Optional[TransformKind], self.optimize_stratety == OptimizeStrategy.ContigousBatching # TODO(lei): Should add ladder stage 3 inverse layout propagation in the expr. # And recover the layout in the schedule template. - and self.M != 1 and 1 not in self.M): + and (self.M != 1 or (isinstance(self.M, Tuple) and 1 not in self.M))): object.__setattr__(self, "propagate_b", TransformKind.LDMatrixTransform) # TODO(lei): This is a limitation arose by pytorch and llvm