@bowang007

When I test flatten via torch.flatten(inputs, start_dim, end_dim), got error AssertionError: False is not true : expected ops {<OpOverload(op='aten.flatten', overload='using_ints')>}, actuall ops {<OpOverload(op='aten.view', overload='default')>, <OpOverloadPacket(op='aten.sym_size')>, <built-in function mul>}

Is there a specific reason - performance or otherwise, why flatten should need a different implementation than reshape, when using static shapes?

Thanks for the advice! I did this because I noticed there's a flatten op in this schema. I thought our goal is to support these native_functions as much as possible. But anyways, I'll comment out the flatten converter and then wrap the reshape to do flatten op.

Generally, the focus is to cover as much of this operation set as possible: https://pytorch.org/docs/stable/ir.html#core-aten-ir, though if there are operators that show up which we can directly convert as opposed to lowering, that is certainly a good thing to have.

I do still think flatten_dims can be a utility which gives the shape to pass to reshape.

I tried implementing flatten_dims in converter_utils.py. Since flatten_dims needs to call reshape, it caused circular import. That's why I moved to shuffle file. Anyways, I will rewrite in flatten_dims.

@zewenli98 I see - thanks for the details - to clarify, I was intending for flatten_dims to not change the network at all since changing the network requires the function to be in impl/. I was thinking it could be instead something like:

def flatten_dims(
    input: Sequence[Union[TRTTensor, torch.Tensor, np.ndarray]],
    start_dim: int,
    end_dim: int,
) -> Tuple[int]:
    shape = input.shape
    dim_size = len(shape)
    start_dim = get_positive_dim(start_dim, dim_size)
    end_dim = get_positive_dim(end_dim, dim_size)

    if not isinstance(input, TRTTensor):
        input = get_trt_tensor(ctx, input, f"{name}_flatten")

    num_elements = 1
    for i in range(start_dim, end_dim + 1):
        num_elements *= shape[i]

    new_shape = tuple(shape[:start_dim]) + (num_elements,) + tuple(shape[end_dim + 1 :])

    return new_shape

Then, the user can call flatten_dims to get back the flattened dimension shape, which they can then pass to reshape themselves.

@gs-olive Thanks a lot! This makes more sense. Modified!

Relevant tests pass locally on Torch 2.1.0. Merging to main.

@bowang007 Please consult this PR for the shuffle op.

@zewenli98 Thanks! Let me update PR accordingly