def __call__(self, a: Tensor, b: Tensor) -> Tensor:
        """Performs sanity checks, offline shape inference and returns an output tensor."""

        a, b = self._align_ab(a, b)
        self._attrs["inputs"] = [a, b]

        self._sanity_check(a, b)

        output_shape = self._infer_shapes(a, b)
        self._extract_epilogue_alignment(output_shape)

        self._attrs["input_accessors"] = [
            TensorAccessor(tensor) for tensor in self._attrs["inputs"]
        ]

        self._set_depth()

        output = Tensor(output_shape, src_ops={self}, dtype=a._attrs["dtype"])
        self._attrs["outputs"] = [output]
        self._attrs["output_accessors"] = [TensorAccessor(output)]
        return output
Performs sanity checks, offline shape inference and returns an output tensor.
python/aitemplate/compiler/ops/gemm_epilogue_vistor/bmm_rcr_softmax.py