    def _extract_dims(self, for_profiling: bool = False) -> Dict[str, List[DimInfo]]:
        """Extracts a mapping between dim names and a list of DimInfo.
        This function will be used in gemm shape inference, gemm padding graph
        transformation, gemm profiling, etc.

        All subclasses must implement this API.

        An example result from gemm_rcr:
        {
            "M": [
                DimInfo(source=INPUT, tensor_idx=0, dim_idx=0),
                DimInfo(source=OUTPUT, tensor_idx=0, dim_idx=0),
            ],
            "K": [
                DimInfo(source=INPUT, tensor_idx=0, dim_idx=1),
                DimInfo(source=INPUT, tensor_idx=1, dim_idx=1),
            ],
            "N": [
                DimInfo(source=INPUT, tensor_idx=1, dim_idx=0),
                DimInfo(source=OUTPUT, tensor_idx=0, dim_idx=1),
            ],
        }


        Parameters
        ----------
        for_profiling: bool
            Whether this function is used for generating profiling source codes.
            If yes, some DimInfo are simplified. e.g. For gemm, we treat all tensors
            as 2d.
        """

        raise NotImplementedError("extract_dims() is not implemented!")
Extracts a mapping between dim names and a list of DimInfo.
This function will be used in gemm shape inference, gemm padding graph
transformation, gemm profiling, etc.

All subclasses must implement this API.

An example result from gemm_rcr:
{
"M": [
DimInfo(source=INPUT, tensor_idx=0, dim_idx=0),
DimInfo(source=OUTPUT, tensor_idx=0, dim_idx=0),
],
"K": [
DimInfo(source=INPUT, tensor_idx=0, dim_idx=1),
DimInfo(source=INPUT, tensor_idx=1, dim_idx=1),
],
"N": [
DimInfo(source=INPUT, tensor_idx=1, dim_idx=0),
DimInfo(source=OUTPUT, tensor_idx=0, dim_idx=1),
],
}


Parameters
----------
for_profiling: bool
Whether this function is used for generating profiling source codes.
If yes, some DimInfo are simplified. e.g. For gemm, we treat all tensors
as 2d.
python/aitemplate/compiler/ops/gemm_universal/gemm_common.py