    def __init__(
        self,
        source: Source,
        tensor_idx: int,
        dim_idx: Union[int, List[int]],
        placeholder: bool = False,
    ):
        """
        source:
            Source.INPUT or Source.OUTPUT
        tensor_idx:
            Depending on source, extract info from inputs[tensor_idx] or outputs[tensor_idx]
        dim_idx:
            Extract shape from inputs/outputs[tensor_idx][dim_idx]
        placeholder:
            If True, the diminfo might not be accurate in compile time, just a placeholder to be filled afterwards
            This is useful to handle issue such as broadcasting which B might not be exact.
        """
        self.source = source
        self.tensor_idx = tensor_idx
        if isinstance(dim_idx, int):
            dim_idx = [dim_idx]
        self.dim_idx = dim_idx
        self.placeholder = placeholder
source:
Source.INPUT or Source.OUTPUT
tensor_idx:
Depending on source, extract info from inputs[tensor_idx] or outputs[tensor_idx]
dim_idx:
Extract shape from inputs/outputs[tensor_idx][dim_idx]
placeholder:
If True, the diminfo might not be accurate in compile time, just a placeholder to be filled afterwards
This is useful to handle issue such as broadcasting which B might not be exact.
python/aitemplate/compiler/ops/gemm_universal/gemm_common.py