[ONNX] Handle multiple outputs of node

[titaiwangms]

pytorch/torch/onnx/_internal/fx/passes/fx_to_onnxscript.py

Line 351 in af440c4

fx_name_to_onnxscipt_value[node.name] = output

Let's say we have an Op like topk, which has multiple outputs. Currently, we treat the whole output as an unit, and the graph trying to use each of the output element to different nodes is not supported.

cc @justinchuby @BowenBao

[titaiwangms]

fx graph handles multiple outputs with getitem, so each output would be recorded independently.

For example,

class TopKModel(torch.nn.Module):
    def forward(self, x):
        values, _ = torch.topk(x, 3)
        return torch.sum(values)

sclass <lambda>(torch.nn.Module):
    def forward(self, arg0: f32[s0]):
        # File: /home/titaiwang/pytorch/test/onnx/test_fx_to_onnx.py:74, code: values, _ = torch.topk(x, 3)
        topk = torch.ops.aten.topk.default(arg0, 3);  arg0 = None
        getitem: f32[3] = topk[0]
        getitem_1: i64[3] = topk[1];  topk = None
        
        # File: /home/titaiwang/pytorch/test/onnx/test_fx_to_onnx.py:75, code: return torch.sum(values)
        sum_1: f32[] = torch.ops.aten.sum.default(getitem);  getitem = None
        return sum_1

[BowenBao]

Is there a covering test case already? If not, could you add the above

[titaiwangms]

Yes, I plan to include it in #96350 by extending test_multiple_outputs_op_with_evaluator in test_fx_to_onnx.py