feat(atenlib): index_select; trace_only ops

[justinchuby]

Implement aten_index_select; enable tests for transpose.

This change also added the trace_only argument to mark functions as trace only to work around functions that cannot be compiled. The logic can then be tested, but the argument should be removed later when onnxscript extends support for the syntax.

[justinchuby]

cc @BowenBao

[BowenBao]

We need to fine a way to call the function that have the attribute dim before the input index.

Sounds like a restriction posed from onnx script function. Is there something fundamentally blocking from mixing the order of inputs and attributes? cc @gramalingam

[gramalingam]

We need to fine a way to call the function that have the attribute dim before the input index.

Sounds like a restriction posed from onnx script function. Is there something fundamentally blocking from mixing the order of inputs and attributes? cc @gramalingam

No, I don't think there is anything fundamentally blocking it. But we need to add the logic to onnxscript translator to process the parameters accordingly, so it is a missing feature.

[gramalingam]

Otherwise, is index expected to be 1-dimensional? Or, can it have any other rank?

[justinchuby]

Yes, it should be 1d I think.

[justinchuby]

should I do op.Reshape(index, [-1]) instead?

[gramalingam]

Yes, if we expect it to be 0d or 1d, I think reshape, as you mention, is a better option.

[justinchuby]

Done. thanks!

[fatcat-z]

Approved.

[codecov]

Codecov Report

Merging #274 (e6b8a9f) into main (47fe75f) will increase coverage by 0.07%.
The diff coverage is 100.00%.

@@            Coverage Diff             @@
##             main     #274      +/-   ##
==========================================
+ Coverage   72.97%   73.04%   +0.07%     
==========================================
  Files          95       95              
  Lines        9113     9131      +18     
==========================================
+ Hits         6650     6670      +20     
+ Misses       2463     2461       -2     

Impacted Files	Coverage Δ
onnxscript/function_libs/torch_aten/ops/core.py	59.07% <100.00%> (+0.57%)	⬆️
...nnxscript/function_libs/torch_aten/registration.py	88.57% <100.00%> (+0.69%)	⬆️
...t/function_libs/torch_aten/ops_correctness_test.py	93.79% <100.00%> (+0.08%)	⬆️

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.