Brevitas `make_fx` generating different graph about brevitas

[vivekkhandelwal1]

@volcacius, PTAL at this issue.
CC: @jinchen62

[volcacius]

Which Pytorch version are you on? With 2.1.1 I see:

Torch Fx graph: 
 graph():
    %arg0_1 : [num_users=1] = placeholder[target=arg0_1]
    %arg1_1 : [num_users=1] = placeholder[target=arg1_1]
    %arg2_1 : [num_users=1] = placeholder[target=arg2_1]
    %_scaled_dot_product_flash_attention : [num_users=9] = call_function[target=torch.ops.aten._scaled_dot_product_flash_attention.default](args = (%arg0_1, %arg1_1, %arg2_1, 0.0, True), kwargs = {})
    %getitem : [num_users=1] = call_function[target=operator.getitem](args = (%_scaled_dot_product_flash_attention, 0), kwargs = {})
    %getitem_1 : [num_users=0] = call_function[target=operator.getitem](args = (%_scaled_dot_product_flash_attention, 1), kwargs = {})
    %getitem_2 : [num_users=0] = call_function[target=operator.getitem](args = (%_scaled_dot_product_flash_attention, 2), kwargs = {})
    %getitem_3 : [num_users=0] = call_function[target=operator.getitem](args = (%_scaled_dot_product_flash_attention, 3), kwargs = {})
    %getitem_4 : [num_users=0] = call_function[target=operator.getitem](args = (%_scaled_dot_product_flash_attention, 4), kwargs = {})
    %getitem_5 : [num_users=0] = call_function[target=operator.getitem](args = (%_scaled_dot_product_flash_attention, 5), kwargs = {})
    %getitem_6 : [num_users=0] = call_function[target=operator.getitem](args = (%_scaled_dot_product_flash_attention, 6), kwargs = {})
    %getitem_7 : [num_users=0] = call_function[target=operator.getitem](args = (%_scaled_dot_product_flash_attention, 7), kwargs = {})
    %getitem_8 : [num_users=0] = call_function[target=operator.getitem](args = (%_scaled_dot_product_flash_attention, 8), kwargs = {})
    return getitem
Brevitas Fx graph: 
 graph():
    %arg0_1 : [#users=1] = placeholder[target=arg0_1]
    %arg1_1 : [#users=1] = placeholder[target=arg1_1]
    %arg2_1 : [#users=1] = placeholder[target=arg2_1]
    %_scaled_dot_product_flash_attention_default : [#users=9] = call_function[target=torch.ops.aten._scaled_dot_product_flash_attention.default](args = (%arg0_1, %arg1_1, %arg2_1, 0.0, True), kwargs = {})
    %getitem : [#users=1] = call_function[target=operator.getitem](args = (%_scaled_dot_product_flash_attention_default, 0), kwargs = {})
    %getitem_1 : [#users=0] = call_function[target=operator.getitem](args = (%_scaled_dot_product_flash_attention_default, 1), kwargs = {})
    %getitem_2 : [#users=0] = call_function[target=operator.getitem](args = (%_scaled_dot_product_flash_attention_default, 2), kwargs = {})
    %getitem_3 : [#users=0] = call_function[target=operator.getitem](args = (%_scaled_dot_product_flash_attention_default, 3), kwargs = {})
    %getitem_4 : [#users=0] = call_function[target=operator.getitem](args = (%_scaled_dot_product_flash_attention_default, 4), kwargs = {})
    %getitem_5 : [#users=0] = call_function[target=operator.getitem](args = (%_scaled_dot_product_flash_attention_default, 5), kwargs = {})
    %getitem_6 : [#users=0] = call_function[target=operator.getitem](args = (%_scaled_dot_product_flash_attention_default, 6), kwargs = {})
    %getitem_7 : [#users=0] = call_function[target=operator.getitem](args = (%_scaled_dot_product_flash_attention_default, 7), kwargs = {})
    %getitem_8 : [#users=0] = call_function[target=operator.getitem](args = (%_scaled_dot_product_flash_attention_default, 8), kwargs = {})
    return getitem

[vivekkhandelwal1]

Which Pytorch version are you on? With 2.1.1 I see:

torch 2.2.0.dev20231115+cpu

[volcacius]

It should be fixed in this PR #763 I just merged in dev. Thanks for spotting it.