feat: support flatten and reshape via shuffle_layer

[zewenli98]

Description

Support flatten and reshape via shuffle_layer

Fixes #2214

Type of change

Please delete options that are not relevant and/or add your own.

• New feature (non-breaking change which adds functionality)

Checklist:

• My code follows the style guidelines of this project (You can use the linters)
• I have performed a self-review of my own code
• I have commented my code, particularly in hard-to-understand areas and hacks
• I have made corresponding changes to the documentation
• I have added tests to verify my fix or my feature
• New and existing unit tests pass locally with my changes
• I have added the relevant labels to my PR in so that relevant reviewers are notified

[zewenli98]

@bowang007

[zewenli98]

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>}

[gs-olive]

Is there a specific reason - performance or otherwise, why flatten should need a different implementation than reshape, when using static shapes? Specifically, we can comment out the flatten implementation for now, and for any converters needing flatten for static shapes, they can just use a reshape and flatten the dimensions themselves.

As an alternative, @zewenli98, you can add a utility flatten_dims, which will flatten the dimensions of an input tensor into a reshape-usable form, then you can have @bowang007's converter test that utility.

[zewenli98]

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.

[gs-olive]

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.

[gs-olive]

I do still think flatten_dims can be a utility which gives the shape to pass to reshape. That way, it can get tested as a utility and not as a converter (see tests/py/dynamo/conversion/test_converter_utils.py). Added a suggestion on syntax.

[zewenli98]

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.

[gs-olive]

@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.

[zewenli98]

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

[gs-olive]

This line can be removed, in accordance with the new testing PR

[gs-olive]

Looks good to me, pending CI pass.

[gs-olive]

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

[zewenli98]

@bowang007 Please consult this PR for the shuffle op.

[bowang007]

@zewenli98 Thanks! Let me update PR accordingly