def create_input_batch_diminfo(input_shapes, batch_dims, output_batch):
    """
    Create inputs' batch diminfo.
    Provided input_shapes and the corresponding batch_dims, this function
    returns a list of batch's DimInfo of the inputs.

    input_shapes:
        A list of input shapes.
    batch_dims:
        The batch dimension for the corresponding input_shapes.
        If length of corresponding input's shape is less than 2, neglected.
    output_batch:
        The batch size for output.
    """
    assert len(input_shapes) == len(batch_dims)

    batch_diminfo = []
    for idx, input_shape in enumerate(input_shapes):
        if len(input_shape) > 2:
            batch_diminfo.append(
                DimInfo(
                    Source.INPUT,
                    tensor_idx=idx,
                    dim_idx=batch_dims[idx],
                    placeholder=input_shape[batch_dims[idx]] != output_batch,
                )
            )
    return batch_diminfo

    Create inputs' batch diminfo.
    Provided input_shapes and the corresponding batch_dims, this function
    returns a list of batch's DimInfo of the inputs.

    input_shapes:
        A list of input shapes.
    batch_dims:
        The batch dimension for the corresponding input_shapes.
        If length of corresponding input's shape is less than 2, neglected.
    output_batch:
        The batch size for output.
    
python/aitemplate/compiler/ops/gemm_universal/gemm_common.py