    def is_valid_shape(a: Tensor, b: Tensor):
        """
        Check input a/b is valid for bmm_rcr_n1.
        Requirements:
            1) matching dimension of a/b (where a is row major, b is column major)
            2) dim N of b needs to be 1
            3) dim K of b needs to be multiple of 8
        """
        if len(a.shape()) != 3 or len(b.shape()) != 3:
            return False

        valid = True
        valid &= a.shape()[0] == b.shape()[0]
        valid &= a.shape()[2] == b.shape()[2]

        # check N = 1
        BN = b.shape()[1]
        if not isinstance(BN, IntImm):
            return False
        valid &= BN.value() == 1

        # check BK is static dim
        BK = b.shape()[2]
        if not isinstance(BK, IntImm):
            return False

        return valid
Check input a/b is valid for bmm_rcr_n1.
Requirements:
1) matching dimension of a/b (where a is row major, b is column major)
2) dim N of b needs to be 1
3) dim K of b needs to be multiple of 8
python/aitemplate/compiler/ops/gemm_special/bmm_rcr_n1.py