    def __call__(self, a: Tensor, b: Tensor, b1: Tensor) -> Tensor:
        """Call the BMM_RCR + 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
            Tensors in shape of [B, M, O]
        """
        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)
        self._extract_epilogue_alignment(output_shape)
        output = Tensor(output_shape, src_ops={self}, dtype=a.dtype())
        self._attrs["outputs"] = [output]
        self._attrs["output_accessors"] = [TensorAccessor(output)]
        return output
Call the BMM_RCR + 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
Tensors in shape of [B, M, O]
python/aitemplate/compiler/ops/gemm_universal/bmm_softmax_bmm.py