[ONNX] Speed up unit tests
Labels: module: onnx, triaged, better-engineering

Reduce runtime for tests in onnx.

## Profile

```
Program: pytest test/onnx/test_pytorch_onnx_onnxruntime.py -k version_17

789.755 <module>  <string>:1
   [2773 frames hidden]  <string>, runpy, pytest, _pytest, plu...
      778.248 __call__  unittest/case.py:650
      └─ 778.248 run  torch/testing/_internal/common_utils.py:2034
         └─ 778.242 _run_with_retry  torch/testing/_internal/common_utils.py:1961
            └─ 778.242 run  unittest/case.py:558
                  [57 frames hidden]  unittest, contextlib, <built-in>, _py...
                     776.002 _callTestMethod  unittest/case.py:549
                     ├─ 303.927 instantiated_test  torch/testing/_internal/common_utils.py:247
                     │  └─ 295.628 wrapper  pytorch_test_common.py:112
                     │     └─ 295.627 wrapper  pytorch_test_common.py:51
                     │        └─ 295.627 test_rnn  test_pytorch_onnx_onnxruntime.py:12500
                     │           └─ 295.627 _dispatch_rnn_test  test_pytorch_onnx_onnxruntime.py:9173
                     │              ├─ 109.771 _lstm_test  test_pytorch_onnx_onnxruntime.py:9282
                     │              │  └─ 109.450 run_test  onnx_test_common.py:82
                     │              │     └─ 109.450 _run_test  onnx_test_common.py:99
                     │              │        └─ 109.450 run_model_test  onnx_test_common.py:36
                     │              │           └─ 108.890 verify  torch/onnx/verification.py:605
                     │              │              └─ 103.088 _export  torch/onnx/utils.py:1380
                     │              │                 └─ 102.224 _model_to_graph  torch/onnx/utils.py:1040
                     │              │                    ├─ 76.864 _optimize_graph  torch/onnx/utils.py:541
                     │              │                    │  ├─ 38.422 _run_symbolic_function  torch/onnx/utils.py:1730
                     │              │                    │  │  └─ 34.782 lstm  torch/onnx/symbolic_opset9.py:4466
                     │              │                    │  │     └─ 34.776 wrapper  torch/onnx/symbolic_helper.py:301
                     │              │                    │  │        ├─ 23.189 _lstm_packed  torch/onnx/symbolic_opset9.py:4434
                     │              │                    │  │        │  └─ 23.126 _generic_rnn  torch/onnx/symbolic_opset9.py:4171
                     │              │                    │  │        │     └─ 20.277 transform_weights  torch/onnx/symbolic_opset9.py:4286
                     │              │                    │  │        │        └─ 17.787 <genexpr>  torch/onnx/symbolic_opset9.py:4292
                     │              │                    │  │        │           └─ 17.779 reform_weights  torch/onnx/symbolic_opset9.py:4265
                     │              │                    │  │        │              └─ 17.590 <listcomp>  torch/onnx/symbolic_opset9.py:4267
                     │              │                    │  │        │                 └─ 17.533 _slice_helper  torch/onnx/symbolic_helper.py:704
                     │              │                    │  │        │                    └─ 17.297 _slice  torch/onnx/symbolic_opset10.py:336
                     │              │                    │  │        │                       └─ 14.634 _graph_op  torch/onnx/_patch_torch.py:19
                     │              │                    │  │        │                          └─ 10.123 PyCapsule._jit_pass_onnx_node_shape_type_inference  <built-in>:0
                     │              │                    │  │        │                                [2 frames hidden]  <built-in>
                     │              │                    │  │        └─ 11.568 _lstm_full  torch/onnx/symbolic_opset9.py:4402
                     │              │                    │  │           └─ 11.548 _generic_rnn  torch/onnx/symbolic_opset9.py:4171
                     │              │                    │  │              └─ 10.117 transform_weights  torch/onnx/symbolic_opset9.py:4286
                     │              │                    │  │                 └─ 8.870 <genexpr>  torch/onnx/symbolic_opset9.py:4292
                     │              │                    │  │                    └─ 8.867 reform_weights  torch/onnx/symbolic_opset9.py:4265
                     │              │                    │  │                       └─ 8.769 <listcomp>  torch/onnx/symbolic_opset9.py:4267
                     │              │                    │  │                          └─ 8.742 _slice_helper  torch/onnx/symbolic_helper.py:704
                     │              │                    │  │                             └─ 8.633 _slice  torch/onnx/symbolic_opset10.py:336
                     │              │                    │  └─ 26.808 PyCapsule._jit_pass_onnx_graph_shape_type_inference  <built-in>:0
                     │              │                    │        [2 frames hidden]  <built-in>
                     │              │                    └─ 18.103 _create_jit_graph  torch/onnx/utils.py:916
                     │              │                       └─ 18.047 _trace_and_get_graph_from_model  torch/onnx/utils.py:848
                     │              │                          └─ 18.015 _get_trace_graph  torch/jit/_trace.py:1138
                     │              │                             └─ 18.005 _call_impl  torch/nn/modules/module.py:1184
                     │              │                                └─ 18.005 forward  torch/jit/_trace.py:94
                     │              │                                   └─ 16.789 wrapper  torch/jit/_trace.py:104
                     │              │                                      └─ 12.148 _call_impl  torch/nn/modules/module.py:1184
                     │              │                                         └─ 12.147 _slow_forward  torch/nn/modules/module.py:1164
                     │              ├─ 98.451 _gru_test  test_pytorch_onnx_onnxruntime.py:9350
                     │              │  └─ 98.128 run_test  onnx_test_common.py:82
                     │              │     └─ 98.128 _run_test  onnx_test_common.py:99
                     │              │        └─ 98.128 run_model_test  onnx_test_common.py:36
                     │              │           └─ 97.712 verify  torch/onnx/verification.py:605
                     │              │              └─ 92.431 _export  torch/onnx/utils.py:1380
                     │              │                 └─ 91.787 _model_to_graph  torch/onnx/utils.py:1040
                     │              │                    ├─ 71.784 _optimize_graph  torch/onnx/utils.py:541
                     │              │                    │  ├─ 36.852 _run_symbolic_function  torch/onnx/utils.py:1730
                     │              │                    │  │  └─ 33.353 symbolic  torch/onnx/symbolic_opset9.py:4564
                     │              │                    │  │     └─ 33.350 wrapper  torch/onnx/symbolic_helper.py:301
                     │              │                    │  │        ├─ 22.248 _rnn_packed  torch/onnx/symbolic_opset9.py:4536
                     │              │                    │  │        │  └─ 22.202 _generic_rnn  torch/onnx/symbolic_opset9.py:4171
                     │              │                    │  │        │     └─ 20.135 transform_weights  torch/onnx/symbolic_opset9.py:4286
                     │              │                    │  │        │        └─ 17.671 <genexpr>  torch/onnx/symbolic_opset9.py:4292
                     │              │                    │  │        │           └─ 17.669 reform_weights  torch/onnx/symbolic_opset9.py:4265
                     │              │                    │  │        │              └─ 17.510 <listcomp>  torch/onnx/symbolic_opset9.py:4267
                     │              │                    │  │        │                 └─ 17.477 _slice_helper  torch/onnx/symbolic_helper.py:704
                     │              │                    │  │        │                    └─ 17.342 _slice  torch/onnx/symbolic_opset10.py:336
                     │              │                    │  │        │                       └─ 15.157 _graph_op  torch/onnx/_patch_torch.py:19
                     │              │                    │  │        │                          └─ 10.545 PyCapsule._jit_pass_onnx_node_shape_type_inference  <built-in>:0
                     │              │                    │  │        │                                [2 frames hidden]  <built-in>
                     │              │                    │  │        └─ 11.101 _rnn_full  torch/onnx/symbolic_opset9.py:4507
                     │              │                    │  │           └─ 11.081 _generic_rnn  torch/onnx/symbolic_opset9.py:4171
                     │              │                    │  │              └─ 10.036 transform_weights  torch/onnx/symbolic_opset9.py:4286
                     │              │                    │  │                 └─ 8.806 <genexpr>  torch/onnx/symbolic_opset9.py:4292
                     │              │                    │  │                    └─ 8.806 reform_weights  torch/onnx/symbolic_opset9.py:4265
                     │              │                    │  │                       └─ 8.737 <listcomp>  torch/onnx/symbolic_opset9.py:4267
                     │              │                    │  │                          └─ 8.712 _slice_helper  torch/onnx/symbolic_helper.py:704
                     │              │                    │  │                             └─ 8.634 _slice  torch/onnx/symbolic_opset10.py:336
                     │              │                    │  └─ 24.573 PyCapsule._jit_pass_onnx_graph_shape_type_inference  <built-in>:0
                     │              │                    │        [2 frames hidden]  <built-in>
                     │              │                    └─ 14.916 _create_jit_graph  torch/onnx/utils.py:916
                     │              │                       └─ 14.862 _trace_and_get_graph_from_model  torch/onnx/utils.py:848
                     │              │                          └─ 14.817 _get_trace_graph  torch/jit/_trace.py:1138
                     │              │                             └─ 14.811 _call_impl  torch/nn/modules/module.py:1184
                     │              │                                └─ 14.811 forward  torch/jit/_trace.py:94
                     │              │                                   └─ 13.683 wrapper  torch/jit/_trace.py:104
                     │              │                                      └─ 9.260 _call_impl  torch/nn/modules/module.py:1184
                     │              │                                         └─ 9.260 _slow_forward  torch/nn/modules/module.py:1164
                     │              └─ 87.405 _elman_rnn_test  test_pytorch_onnx_onnxruntime.py:9181
                     │                 └─ 86.774 run_test  onnx_test_common.py:82
                     │                    └─ 86.774 _run_test  onnx_test_common.py:99
                     │                       └─ 86.773 run_model_test  onnx_test_common.py:36
                     │                          └─ 85.647 verify  torch/onnx/verification.py:605
                     │                             └─ 78.367 _export  torch/onnx/utils.py:1380
                     │                                └─ 77.105 _model_to_graph  torch/onnx/utils.py:1040
                     │                                   ├─ 42.625 _optimize_graph  torch/onnx/utils.py:541
                     │                                   │  ├─ 20.128 _run_symbolic_function  torch/onnx/utils.py:1730
                     │                                   │  │  └─ 13.245 symbolic  torch/onnx/symbolic_opset9.py:4564
                     │                                   │  │     └─ 13.240 wrapper  torch/onnx/symbolic_helper.py:301
                     │                                   │  │        └─ 8.815 _rnn_packed  torch/onnx/symbolic_opset9.py:4536
                     │                                   │  │           └─ 8.725 _generic_rnn  torch/onnx/symbolic_opset9.py:4171
                     │                                   │  └─ 12.593 PyCapsule._jit_pass_onnx_graph_shape_type_inference  <built-in>:0
                     │                                   │        [2 frames hidden]  <built-in>
                     │                                   └─ 26.437 _create_jit_graph  torch/onnx/utils.py:916
                     │                                      └─ 26.333 _trace_and_get_graph_from_model  torch/onnx/utils.py:848
                     │                                         └─ 26.261 _get_trace_graph  torch/jit/_trace.py:1138
                     │                                            └─ 26.248 _call_impl  torch/nn/modules/module.py:1184
                     │                                               └─ 26.246 forward  torch/jit/_trace.py:94
                     │                                                  └─ 24.035 wrapper  torch/jit/_trace.py:104
                     │                                                     ├─ 15.227 _call_impl  torch/nn/modules/module.py:1184
                     │                                                     │  └─ 15.226 _slow_forward  torch/nn/modules/module.py:1164
                     │                                                     └─ 8.795 <genexpr>  torch/jit/_trace.py:114
                     ├─ 271.394 wrapper  pytorch_test_common.py:51
                     │  ├─ 29.974 test_index_put_loop  test_pytorch_onnx_onnxruntime.py:2351
                     │  │  └─ 29.930 run_test  onnx_test_common.py:82
                     │  │     └─ 29.913 _run_test  onnx_test_common.py:99
                     │  │        └─ 29.913 run_model_test  onnx_test_common.py:36
                     │  │           └─ 29.888 verify  torch/onnx/verification.py:605
                     │  │              └─ 24.322 _export  torch/onnx/utils.py:1380
                     │  │                 └─ 23.918 _model_to_graph  torch/onnx/utils.py:1040
                     │  │                    └─ 14.552 _optimize_graph  torch/onnx/utils.py:541
                     │  │                       └─ 9.139 _run_symbolic_function  torch/onnx/utils.py:1730
                     │  │                          └─ 9.027 prim_loop  torch/onnx/symbolic_opset9.py:6298
                     │  ├─ 10.137 wrapper  pytorch_test_common.py:112
                     │  ├─ 8.031 test_index_put_slice_index  test_pytorch_onnx_onnxruntime.py:2235
                     │  │  └─ 8.031 run_test  onnx_test_common.py:82
                     │  └─ 7.961 test_instancenorm3d_runningstats  test_pytorch_onnx_onnxruntime.py:3771
                     ├─ 11.971 test_avgpool_3d_ceil  test_pytorch_onnx_onnxruntime.py:1420
                     │  └─ 9.299 _VariableFunctionsClass.randn  <built-in>:0
                     │        [2 frames hidden]  <built-in>
                     └─ 10.681 wrapper  pytorch_test_common.py:112
```