[Tracer] Add flatten argument to .trace()

[chhzh123]

Description

This PR adds flatten argument to the .trace() primitive, so users can determine whether they want the traced graph to be flattened or not. Combining with leaf_module, users can have fine-grained control on how many and which levels of submodules they want to be flattened or hierarchical. For example, to only flatten the BertAttention layer but not the inner BertSelfAttention and BertSelfOutput modules, users can write the following code.

    input_names = ["hidden_states", "attention_mask"]
    sig = inspect.signature(sub_sch.mod.forward)
    concrete_args = {
        p.name: p.default for p in sig.parameters.values() if p.name not in input_names
    }
    sub_sch.trace(
        tracer="pytorch",
        flatten=True,
        concrete_args=concrete_args,
        leaf_modules=["BertSelfAttention", "BertSelfOutput"],
    )

As shown below, only attention.self and attention.output are generated in the dataflow graph, but not the inner attention.self.query/key/value layers.

def forward(self, hidden_states : torch.Tensor, attention_mask : typing_Union[torch.FloatTensor,NoneType] = None, head_mask = None, encoder_hidden_states = None, encoder_attention_mask = None, past_key_value = None, output_attentions = False) -> typing_Tuple[torch.Tensor]:
    eq = output_attentions == False;  output_attentions = None
    _assert = torch._assert(eq, 'output_attentions has been specialized to have value False but got another value');  eq = None
    attention_self = self.attention.self(hidden_states, attention_mask, None, None, None, None, False);  attention_mask = None
    getitem = attention_self[0]
    attention_output = self.attention.output(getitem, hidden_states);  getitem = hidden_states = None
    getitem_1 = attention_self[slice(1, None, None)];  attention_self = None
    add = (attention_output,) + getitem_1;  attention_output = getitem_1 = None
    getitem_2 = add[0]
    getitem_3 = add[slice(1, None, None)];  add = None
    intermediate_dense = self.intermediate.dense(getitem_2)
    gelu = torch._C._nn.gelu(intermediate_dense);  intermediate_dense = None
    output_dense = self.output.dense(gelu);  gelu = None
    output_dropout = self.output.dropout(output_dense);  output_dense = None
    add_1 = output_dropout + getitem_2;  output_dropout = getitem_2 = None
    output_layer_norm = self.output.LayerNorm(add_1);  add_1 = None
    add_2 = (output_layer_norm,) + getitem_3;  output_layer_norm = getitem_3 = None
    return add_2

Checklist

• PR's title starts with a category (e.g. [Bugfix], [Model], [Tutorial], etc)
• Changes are complete (i.e. I finished coding on this PR)
• All changes have test coverage
• Code is well-documented

[comaniac]

LGTM

[comaniac]

Oh btw, you may need to update README.

[chhzh123]

Which part? The current README has already presented the flatten arg

[comaniac]

I didn't notice that...then it's good to go.

[comaniac]

Thanks @chhzh123