Repository: microsoft/onnxscript

Issue:
[torchlib] Assertion error from "aten.gather.default"
Discovered by onnxbench, failing in various huggingface models

  - 17 like: 
```
FX Node: call_function:aten.gather.default[name=gather]. Raised from:
    AssertionError: 
```
```
        DebertaForQuestionAnswering
        LayoutLMForMaskedLM
        RobertaForQuestionAnswering
        DistilBertForMaskedLM
        MegatronBertForQuestionAnswering
        AlbertForQuestionAnswering
        AlbertForMaskedLM
        MobileBertForQuestionAnswering
        MegatronBertForCausalLM
        ElectraForQuestionAnswering
        CamemBert
        ElectraForCausalLM
        BertForMaskedLM
        DebertaForMaskedLM
        RobertaForCausalLM
        BertForQuestionAnswering
        DistilBertForQuestionAnswering
```

Example traceback:
```
Traceback (most recent call last):
  File "/bert_ort/bowbao/repos/bench/pytorch/torch/onnx/_internal/diagnostics/infra/decorator.py", line 127, in wrapper
    return_values = fn(*args, **kwargs)
  File "/bert_ort/bowbao/repos/bench/pytorch/torch/onnx/_internal/fx/passes/fx_to_onnxscript.py", line 391, in _export_fx_node_to_onnxscript
    ] = symbolic_fn(*onnx_args, **onnx_kwargs)
  File "/bert_ort/bowbao/repos/bench/onnxscript/onnxscript/values.py", line 503, in __call__
    return evaluator.default().eval_function(self, args, kwargs)
  File "<@beartype(onnxscript.function_libs.torch_lib.graph_building.TorchScriptTracingEvaluator.eval_function) at 0x7f9f57233790>", line 113, in eval_function
  File "/bert_ort/bowbao/repos/bench/onnxscript/onnxscript/function_libs/torch_lib/graph_building.py", line 264, in eval_function
    return self._graph.add_function_call(function, inputs, attributes)
  File "<@beartype(onnxscript.function_libs.torch_lib.graph_building.TorchScriptGraph.add_function_call) at 0x7f9f57243a60>", line 114, in add_function_call
  File "/bert_ort/bowbao/repos/bench/onnxscript/onnxscript/function_libs/torch_lib/graph_building.py", line 585, in add_function_call
    result = self._add_torchscript_op_call(
  File "<@beartype(onnxscript.function_libs.torch_lib.graph_building.TorchScriptGraph._add_torchscript_op_call) at 0x7f9f57243550>", line 129, in _add_torchscript_op_call
  File "/bert_ort/bowbao/repos/bench/onnxscript/onnxscript/function_libs/torch_lib/graph_building.py", line 518, in _add_torchscript_op_call
    assert not isinstance(value, TorchScriptTensor)
AssertionError
```


Discussion:
Fundamental issue being we were not supporting intersected inputs and attrs for OnnxFunction. Hence they are defined in order of (all inputs) + (all attrs). On the other hand, PyTorch exporter provides arguments in aten signature order.

Fixed by PR stack from #759.

Resolving patch:
diff --git a/onnxscript/function_libs/torch_lib/graph_building.py b/onnxscript/function_libs/torch_lib/graph_building.py
--- a/onnxscript/function_libs/torch_lib/graph_building.py
+++ b/onnxscript/function_libs/torch_lib/graph_building.py
@@ -514,8 +514,10 @@ def _add_torchscript_op_call(
                 graph_inputs.append(self._add_constant_to_graph(input))
             else:
                 graph_inputs.append(input)
-        for value in onnx_attributes.values():
-            assert not isinstance(value, TorchScriptTensor)
+        for key, value in onnx_attributes.items():
+            assert not isinstance(
+                value, TorchScriptTensor
+            ), f"ONNX attribute must not be a TorchScriptTensor, got {key}: {value}."
         result = _create_op_call_in_torch_graph(
             self._torch_graph,
             name,
diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py
--- a/onnxscript/function_libs/torch_lib/ops/core.py
+++ b/onnxscript/function_libs/torch_lib/ops/core.py
@@ -2620,8 +2620,8 @@ def aten_fused_moving_avg_obs_fake_quant(
 @torch_op("aten::gather")
 def aten_gather(
     self: TReal,
-    index: TInt,
     dim: int,
+    index: TInt,
     sparse_grad: bool = False,  # pylint: disable=unused-argument
 ) -> TReal:
     """gather(Tensor self, int dim, Tensor index, *, bool sparse_grad=False) -> Tensor"""


Test patch:
diff --git a/onnxscript/tests/function_libs/torch_lib/ops_test_data.py b/onnxscript/tests/function_libs/torch_lib/ops_test_data.py
--- a/onnxscript/tests/function_libs/torch_lib/ops_test_data.py
+++ b/onnxscript/tests/function_libs/torch_lib/ops_test_data.py
@@ -154,14 +154,6 @@ def _flip_input_wrangler(
     return args, kwargs
 
 
-def _gather_input_wrangler(
-    args: list[Any], kwargs: dict[str, Any]
-) -> tuple[list[Any], dict[str, Any]]:
-    # Make the dim argument an attribute
-    kwargs["dim"] = args.pop(1)
-    return args, kwargs
-
-
 def _grid_sample_input_wrangler(
     args: list[Any], kwargs: dict[str, Any]
 ) -> tuple[list[Any], dict[str, Any]]:
@@ -382,7 +374,7 @@ def _where_input_wrangler(
     "full": core_ops.aten_full,
     "full_like_dtype": core_ops.aten_full_like_dtype,
     "full_like": core_ops.aten_full_like,
-    "gather": (core_ops.aten_gather, _gather_input_wrangler),
+    "gather": core_ops.aten_gather,
     "ge": core_ops.aten_ge,
     # "greater_equal": core_ops.aten_greater_equal,  # no test case in OPS_DB
     # "greater": core_ops.aten_greater,  # no test case in OPS_DB