    def __init__(self, shape: Tuple[int], scale=1.0, causal=False, layout="0213"):
        """Constructor for BMM_RCR * scale + Softmax + BMM_RRR

        Parameters
        ----------
        shape : Tuple[int]
            The reshape dim info, in attention family, the shape is [num_heads]
        scale : float, optional
            The norm scale, in attention family, it is 1.0 / sqrt(num_heads)
        layout : str, optional
            output permute layout, by default "0213"
        """
        super().__init__()
        causal_mask = "_causal" if causal else ""
        self._attrs["op"] = "bmm_softmax_bmm_permute" + causal_mask
        self._attrs["scale"] = scale
        self._attrs["shape"] = shape
        self._attrs["layout"] = "Permute4DBMM_{}".format(layout)

        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 for BMM_RCR * scale + Softmax + BMM_RRR

Parameters
----------
shape : Tuple[int]
The reshape dim info, in attention family, the shape is [num_heads]
scale : float, optional
The norm scale, in attention family, it is 1.0 / sqrt(num_heads)
layout : str, optional
output permute layout, by default "0213"
python/aitemplate/compiler/ops/gemm_universal/bmm_softmax_bmm_permute.py