def _make_input_view_shape(
    cat_input: Tensor,
    original_view_shape: List[IntVar],
    cat_dim: int,
    input_idx: int,
) -> Optional[List[IntVar]]:
    """
    Assumes that there is a pattern like concat + view_op in the graph, we tries
    to transform it into view_op + concat. However, it's not always valid to
    perform such a transformation, because the concat's original inputs may
    not be shape-compatible with the moved view_op. Currently, we only support
    cases where a view_op only changes the dims after cat_dim and the dims after
    the cat_dim must be static.

    For example, for the following code:

        x1 = Tensor(batch, 3 * 4)
        x2 = Tensor(batch, 5 * 4)
        concat_0 = concat([x1, x2], cat_dim=1)
        reshape_1 = reshape(concat_0, [batch, 8, 4])

    This function will generate shape [batch, 3, 4] for x1 and [batch, 5, 4]
    for x2, respectively.

    In contrast, if we have code like below:

        x1 = tensor([batch, 16])
        x2 = tensor([batch, 8])
        cat_1 = concatenate([x1, x2], cat_dim=1)
        reshape_2 = reshape(cat_1, [batch, 4, 6])

    We would return None for both x1 and x2, because we cannot make valid reshape
    ops for x1 and x2 while keeping the original semantics.

    Parameters
    ----------
    cat_input: Tensor
        a concat op's input for which we will generate a view op, e.g. the x1 or
        x2 tensor in the example above
    original_view_shape: List[IntVar]
        the shape of the view op's output, where the view op consumes the concat's
        output, e.g. the reshape op in the example above
    cat_dim: int
        the value of the cat_dim attribute of the concate op
    input_idx: int
        the index of the cat_input in the concat's inputs list
    """
    cat_input_shape = cat_input.shape()
    if cat_dim >= len(cat_input_shape) or cat_dim >= len(original_view_shape):
        return None
    # make sure each dimension at the same index in front of the cat_dim is the same for
    # both cat_input_shape and original_view_shape
    for curr_cat_dim, orig_dim in zip(
        cat_input_shape[:cat_dim], original_view_shape[:cat_dim]
    ):
        if curr_cat_dim != orig_dim:
            return None
    input_stride_at_cat_dim = shape_utils.get_static_stride(cat_input_shape, cat_dim)
    # make sure all dimensions are static after cat_dim
    if input_stride_at_cat_dim is None:
        return None
    orig_view_stride_at_cat_dim = shape_utils.get_static_stride(
        original_view_shape, cat_dim
    )
    # make sure all dimensions are static after cat_dim
    if orig_view_stride_at_cat_dim is None:
        return None
    new_input_view_shape = copy.deepcopy(original_view_shape)
    cat_stride = cat_input_shape[cat_dim].value() * input_stride_at_cat_dim
    if cat_stride % orig_view_stride_at_cat_dim != 0:
        return None
    orig_dim_name = original_view_shape[cat_dim]._attrs["name"]
    new_input_view_shape[cat_dim] = IntImm(
        cat_stride // orig_view_stride_at_cat_dim,
        name=f'{orig_dim_name}_{cat_input._attrs["name"]}_{input_idx}',
    )
    return new_input_view_shape

    Assumes that there is a pattern like concat + view_op in the graph, we tries
    to transform it into view_op + concat. However, it's not always valid to
    perform such a transformation, because the concat's original inputs may
    not be shape-compatible with the moved view_op. Currently, we only support
    cases where a view_op only changes the dims after cat_dim and the dims after
    the cat_dim must be static.

    For example, for the following code:

        x1 = Tensor(batch, 3 * 4)
        x2 = Tensor(batch, 5 * 4)
        concat_0 = concat([x1, x2], cat_dim=1)
        reshape_1 = reshape(concat_0, [batch, 8, 4])

    This function will generate shape [batch, 3, 4] for x1 and [batch, 5, 4]
    for x2, respectively.

    In contrast, if we have code like below:

        x1 = tensor([batch, 16])
        x2 = tensor([batch, 8])
        cat_1 = concatenate([x1, x2], cat_dim=1)
        reshape_2 = reshape(cat_1, [batch, 4, 6])

    We would return None for both x1 and x2, because we cannot make valid reshape
    ops for x1 and x2 while keeping the original semantics.

    Parameters
    ----------
    cat_input: Tensor
        a concat op's input for which we will generate a view op, e.g. the x1 or
        x2 tensor in the example above
    original_view_shape: List[IntVar]
        the shape of the view op's output, where the view op consumes the concat's
        output, e.g. the reshape op in the example above
    cat_dim: int
        the value of the cat_dim attribute of the concate op
    input_idx: int
        the index of the cat_input in the concat's inputs list
    
python/aitemplate/compiler/transform/move_view_ops.py