def __init__(self, scale=1.0):
        """Constructor of BMM_RCR + Softmax + BMM_RRR op

        Parameters
        ----------
        scale : float, optional
            normalization factor, by default 1.0

        """
        super().__init__()
        self._attrs["op"] = "bmm_softmax_bmm"
        self._attrs["scale"] = scale

        def cal_align_ab(m, n, k):
            return common.default_align_ab(k, k, self._attrs["inputs"][0].dtype())

        self._attrs["f_ab_alignment"] = cal_align_ab
Constructor of BMM_RCR + Softmax + BMM_RRR op

        Parameters
        ----------
        scale : float, optional
            normalization factor, by default 1.0

        
python/aitemplate/compiler/ops/gemm_universal/bmm_softmax_bmm.py