    def __call__(self, a: Tensor, b: Tensor) -> Tensor:
        """Call the gemm op.

        Parameters
        ----------
        a : Tensor
            Tensor with correct shape for the gemm operand A.
        b : Tensor
            Tensor with correct shape for the gemm operand B.

        Returns
        -------
        Tensor
            Output tensor for the gemm operation.
        """
        a, b = self._align_ab(a, b)
        self._attrs["inputs"] = [a, b]
        # TensorAccessor(b) is for bmm or rare cases of gemm where b is not constant weight
        self._attrs["input_accessors"] = [TensorAccessor(a), TensorAccessor(b)]
        self._set_depth()
        self._sanity_check(a, b)
        output_shape = self._infer_shapes(a, b)
        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 gemm op.

Parameters
----------
a : Tensor
Tensor with correct shape for the gemm operand A.
b : Tensor
Tensor with correct shape for the gemm operand B.

Returns
-------
Tensor
Output tensor for the gemm operation.
python/aitemplate/compiler/ops/gemm_universal/gemm_common.py