    def __call__(self, a: Tensor, b: Tensor, b1: Tensor) -> Tensor:
        """Call BMM_RCR * scale + Softmax + BMM_RRR op

        Parameters
        ----------
        a : Tensor
            Tensor in shape of [B, M, K]
        b : Tensor
            Tensor in shape of [B, N, K]
        b1 : Tensor
            Tensor in shape of [B, N, O]

        Returns
        -------
        Tensor
            Tensor in shape of [B, M, D, O]

        Raises
        ------
        NotImplementedError
            Other permutation is not implemented yet
        """
        a, b = self._align_ab(a, b)
        self._attrs["inputs"] = [a, b, b1]
        self._attrs["input_accessors"] = [
            TensorAccessor(a),
            TensorAccessor(b),
            TensorAccessor(b1),
        ]
        self._set_depth()
        self._sanity_check(a, b)
        output_shape = self._infer_shapes(a, b, b1)

        output = Tensor(output_shape, src_ops={self}, dtype=a.dtype())
        self._attrs["outputs"] = [output]
        self._attrs["output_accessors"] = [TensorAccessor(output)]

        if self._attrs["layout"] == "Permute4DBMM_0213":
            b, m, o = output_shape
            d1 = self._attrs["shape"][0]
            output_shape = [b.value() // d1, m, d1, o]
            self._extract_epilogue_alignment(output_shape)
            return reshape()(output, output_shape)
        else:
            raise NotImplementedError(
                "{} is not implemented!".format(self._attrs["layout"])
            )
        return output
Call BMM_RCR * scale + Softmax + BMM_RRR op

Parameters
----------
a : Tensor
Tensor in shape of [B, M, K]
b : Tensor
Tensor in shape of [B, N, K]
b1 : Tensor
Tensor in shape of [B, N, O]

Returns
-------
Tensor
Tensor in shape of [B, M, D, O]

Raises
------
NotImplementedError
Other permutation is not implemented yet
python/aitemplate/compiler/ops/gemm_universal/bmm_softmax_bmm_permute.py