Repository: pytorch
Issue: `UnsupportedOperatorError`, `OnnxExporterError` and `SymbolicValueError` related to MultiheadAttention export to onnx with torch.jit.script
Labels: module: onnx,triaged

### 🐛 Describe the bug

When exporting a model to onnx (embedded into a `torch.jit.script` call) that uses MultiheadAttention, I am running into errors related to unsupported operators.

This is all part of a [bigger project](https://github.com/AnnikaStein/DeepJet/tree/adversarial_training/ParT) (I'm coming from High Energy Physics research where we use pytorch to identify the origin of certain objects produced in particle collisions) and the code to reproduce the issue can be found [in this gist](https://gist.github.com/AnnikaStein/500d1c7d7ac80633e331a5647c97187d), containing the model architecture as well as the wrapper code we use to export the model. The problems are traced down to nn.MultiheadAttention calls, all other operators have been made torch.jit.script & onnx-export -friendly by rewriting functions with missing operators wherever possible.

The export without `torch.jit.script` works, however, with `script`, I'm running into the following errors each suggesting that I open an issue.

a)
```
Traceback (most recent call last):
  File "/rwthfs/rz/cluster/home/um106329/aisafety/ParT/Part_export_minimalExample.py", line 71, in <module>
    conversion_loop(model)
  File "/rwthfs/rz/cluster/home/um106329/aisafety/ParT/Part_export_minimalExample.py", line 49, in conversion_loop
    torch.onnx.export(model, inputs, f"Example.onnx", verbose=True,
  File "/work/um106329/conda-storage/envs/VHccAllIn/lib/python3.10/site-packages/torch/onnx/utils.py", line 504, in export
    _export(
  File "/work/um106329/conda-storage/envs/VHccAllIn/lib/python3.10/site-packages/torch/onnx/utils.py", line 1529, in _export
    graph, params_dict, torch_out = _model_to_graph(
  File "/work/um106329/conda-storage/envs/VHccAllIn/lib/python3.10/site-packages/torch/onnx/utils.py", line 1115, in _model_to_graph
    graph = _optimize_graph(
  File "/work/um106329/conda-storage/envs/VHccAllIn/lib/python3.10/site-packages/torch/onnx/utils.py", line 663, in _optimize_graph
    graph = _C._jit_pass_onnx(graph, operator_export_type)
  File "/work/um106329/conda-storage/envs/VHccAllIn/lib/python3.10/site-packages/torch/onnx/utils.py", line 1867, in _run_symbolic_function
    return symbolic_fn(graph_context, *inputs, **attrs)
  File "/work/um106329/conda-storage/envs/VHccAllIn/lib/python3.10/site-packages/torch/onnx/symbolic_opset9.py", line 6567, in prim_if
    env = torch._C._jit_pass_onnx_block(
  File "/work/um106329/conda-storage/envs/VHccAllIn/lib/python3.10/site-packages/torch/onnx/utils.py", line 1867, in _run_symbolic_function
    return symbolic_fn(graph_context, *inputs, **attrs)
  File "/work/um106329/conda-storage/envs/VHccAllIn/lib/python3.10/site-packages/torch/onnx/symbolic_opset9.py", line 6594, in prim_if
    torch._C._jit_pass_onnx_block(
  File "/work/um106329/conda-storage/envs/VHccAllIn/lib/python3.10/site-packages/torch/onnx/utils.py", line 1867, in _run_symbolic_function
    return symbolic_fn(graph_context, *inputs, **attrs)
  File "/work/um106329/conda-storage/envs/VHccAllIn/lib/python3.10/site-packages/torch/onnx/symbolic_opset9.py", line 6567, in prim_if
    env = torch._C._jit_pass_onnx_block(
  File "/work/um106329/conda-storage/envs/VHccAllIn/lib/python3.10/site-packages/torch/onnx/utils.py", line 1867, in _run_symbolic_function
    return symbolic_fn(graph_context, *inputs, **attrs)
  File "/work/um106329/conda-storage/envs/VHccAllIn/lib/python3.10/site-packages/torch/onnx/symbolic_opset9.py", line 6567, in prim_if
    env = torch._C._jit_pass_onnx_block(
  File "/work/um106329/conda-storage/envs/VHccAllIn/lib/python3.10/site-packages/torch/onnx/utils.py", line 1889, in _run_symbolic_function
    raise errors.UnsupportedOperatorError(
torch.onnx.errors.UnsupportedOperatorError: Exporting the operator 'aten::format' to ONNX opset version 14 is not supported. Please feel free to request support or submit a pull request on PyTorch GitHub: https://github.com/pytorch/pytorch/issues
```

I was able to bypass this by modifying `nn/modules/activation.py` [like so in my private fork of pytorch](https://github.com/AnnikaStein/pytorch/commit/d91f31205e6284fb1d19ab63689a87b43e4c1e39) - the fix replaces string-formatting, and introduces a boolean flag for a simple check instead of checking the non-emptiness of a string directly with implicit true/false behavior.

b)
```
Traceback (most recent call last):
  File "/rwthfs/rz/cluster/home/um106329/aisafety/ParT/Part_export_minimalExample.py", line 71, in <module>
    conversion_loop(model)
  File "/rwthfs/rz/cluster/home/um106329/aisafety/ParT/Part_export_minimalExample.py", line 49, in conversion_loop
    torch.onnx.export(model, inputs, f"Example.onnx", verbose=True,
  File "/work/um106329/conda-storage/envs/VHccAllIn/lib/python3.10/site-packages/torch/onnx/utils.py", line 504, in export
    _export(
  File "/work/um106329/conda-storage/envs/VHccAllIn/lib/python3.10/site-packages/torch/onnx/utils.py", line 1529, in _export
    graph, params_dict, torch_out = _model_to_graph(
  File "/work/um106329/conda-storage/envs/VHccAllIn/lib/python3.10/site-packages/torch/onnx/utils.py", line 1115, in _model_to_graph
    graph = _optimize_graph(
  File "/work/um106329/conda-storage/envs/VHccAllIn/lib/python3.10/site-packages/torch/onnx/utils.py", line 663, in _optimize_graph
    graph = _C._jit_pass_onnx(graph, operator_export_type)
  File "/work/um106329/conda-storage/envs/VHccAllIn/lib/python3.10/site-packages/torch/onnx/utils.py", line 1867, in _run_symbolic_function
    return symbolic_fn(graph_context, *inputs, **attrs)
  File "/work/um106329/conda-storage/envs/VHccAllIn/lib/python3.10/site-packages/torch/onnx/symbolic_opset9.py", line 6594, in prim_if
    torch._C._jit_pass_onnx_block(
  File "/work/um106329/conda-storage/envs/VHccAllIn/lib/python3.10/site-packages/torch/onnx/utils.py", line 1867, in _run_symbolic_function
    return symbolic_fn(graph_context, *inputs, **attrs)
  File "/work/um106329/conda-storage/envs/VHccAllIn/lib/python3.10/site-packages/torch/onnx/symbolic_opset9.py", line 6594, in prim_if
    torch._C._jit_pass_onnx_block(
  File "/work/um106329/conda-storage/envs/VHccAllIn/lib/python3.10/site-packages/torch/onnx/utils.py", line 1867, in _run_symbolic_function
    return symbolic_fn(graph_context, *inputs, **attrs)
  File "/work/um106329/conda-storage/envs/VHccAllIn/lib/python3.10/site-packages/torch/onnx/symbolic_opset11.py", line 508, in add
    return symbolic_helper._unimplemented(
  File "/work/um106329/conda-storage/envs/VHccAllIn/lib/python3.10/site-packages/torch/onnx/symbolic_helper.py", line 577, in _unimplemented
    _onnx_unsupported(f"{op}, {msg}", value)
  File "/work/um106329/conda-storage/envs/VHccAllIn/lib/python3.10/site-packages/torch/onnx/symbolic_helper.py", line 592, in _onnx_unsupported
    raise errors.OnnxExporterError(message)
torch.onnx.errors.OnnxExporterError: Unsupported: ONNX export of operator add, does not support adding dynamic tensor list to another. Please feel free to request support or submit a pull request on PyTorch GitHub: https://github.com/pytorch/pytorch/issues
```

I was able to bypass this by modifying `nn/functional.py` [like so in my private fork of pytorch](https://github.com/AnnikaStein/pytorch/commit/40cf95aed62a78e785a90bfc14ce17b411301f62) - here adding dynamic tensor list is replaced with stepwise adding components into output list.

c) Now I am left with a batched matrix multiplication causing an Unknown Rank error, which I was not able to resolve myself:

```
Traceback (most recent call last):
  File "/rwthfs/rz/cluster/home/um106329/aisafety/ParT/Part_export_minimalExample.py", line 71, in <module>
    conversion_loop(model)
  File "/rwthfs/rz/cluster/home/um106329/aisafety/ParT/Part_export_minimalExample.py", line 49, in conversion_loop
    torch.onnx.export(model, inputs, f"Example.onnx", verbose=True,
  File "/work/um106329/conda-storage/envs/VHccAllIn/lib/python3.10/site-packages/torch/onnx/utils.py", line 504, in export
    _export(
  File "/work/um106329/conda-storage/envs/VHccAllIn/lib/python3.10/site-packages/torch/onnx/utils.py", line 1529, in _export
    graph, params_dict, torch_out = _model_to_graph(
  File "/work/um106329/conda-storage/envs/VHccAllIn/lib/python3.10/site-packages/torch/onnx/utils.py", line 1115, in _model_to_graph
    graph = _optimize_graph(
  File "/work/um106329/conda-storage/envs/VHccAllIn/lib/python3.10/site-packages/torch/onnx/utils.py", line 663, in _optimize_graph
    graph = _C._jit_pass_onnx(graph, operator_export_type)
  File "/work/um106329/conda-storage/envs/VHccAllIn/lib/python3.10/site-packages/torch/onnx/utils.py", line 1867, in _run_symbolic_function
    return symbolic_fn(graph_context, *inputs, **attrs)
  File "/work/um106329/conda-storage/envs/VHccAllIn/lib/python3.10/site-packages/torch/onnx/symbolic_helper.py", line 380, in wrapper
    return fn(g, *args, **kwargs)
  File "/work/um106329/conda-storage/envs/VHccAllIn/lib/python3.10/site-packages/torch/onnx/symbolic_helper.py", line 303, in wrapper
    return fn(g, *args, **kwargs)
  File "/work/um106329/conda-storage/envs/VHccAllIn/lib/python3.10/site-packages/torch/onnx/symbolic_opset9.py", line 1001, in transpose
    raise errors.SymbolicValueError(
torch.onnx.errors.SymbolicValueError: Unsupported: ONNX export of transpose for tensor of unknown rank.  [Caused by the value 'attn_output defined in (%attn_output : FloatTensor(device=cpu) = onnx::MatMul(%ret, %v.3), scope: ParT_minimalExample.ParticleTransformer::/ParT_minimalExample.HF_TransformerEncoder::Encoder/ParT_minimalExample.HF_TransformerEncoderLayer::layers.0/torch.nn.modules.activation.MultiheadAttention::self_attn # /work/um106329/conda-storage/envs/VHccAllIn/lib/python3.10/site-packages/torch/nn/functional.py:5168:18
)' (type 'Tensor') in the TorchScript graph. The containing node has kind 'onnx::MatMul'.] 
    (node defined in   File "/work/um106329/conda-storage/envs/VHccAllIn/lib/python3.10/site-packages/torch/nn/functional.py", line 5168
        attn_output_weights = dropout(attn_output_weights, p=dropout_p)

    attn_output = torch.bmm(attn_output_weights, v)
                  ~~~~~~~~~ <--- HERE
    #attn_output = torch.matmul(attn_output_weights, v)
)

    Inputs:
        #0: ret defined in (%ret : FloatTensor(device=cpu) = onnx::Softmax[axis=-1](%attn_output_weights), scope: ParT_minimalExample.ParticleTransformer::/ParT_minimalExample.HF_TransformerEncoder::Encoder/ParT_minimalExample.HF_TransformerEncoderLayer::layers.0/torch.nn.modules.activation.MultiheadAttention::self_attn # /work/um106329/conda-storage/envs/VHccAllIn/lib/python3.10/site-packages/torch/nn/functional.py:1841:14
    )  (type 'Tensor')
        #1: v.3 defined in (%v.3 : Tensor(*, *, *) = onnx::Transpose[perm=[1, 0, 2]](%2010), scope: ParT_minimalExample.ParticleTransformer::/ParT_minimalExample.HF_TransformerEncoder::Encoder/ParT_minimalExample.HF_TransformerEncoderLayer::layers.0/torch.nn.modules.activation.MultiheadAttention::self_attn # /work/um106329/conda-storage/envs/VHccAllIn/lib/python3.10/site-packages/torch/nn/functional.py:5109:12
    )  (type 'Tensor')
    Outputs:
        #0: attn_output defined in (%attn_output : FloatTensor(device=cpu) = onnx::MatMul(%ret, %v.3), scope: ParT_minimalExample.ParticleTransformer::/ParT_minimalExample.HF_TransformerEncoder::Encoder/ParT_minimalExample.HF_TransformerEncoderLayer::layers.0/torch.nn.modules.activation.MultiheadAttention::self_attn # /work/um106329/conda-storage/envs/VHccAllIn/lib/python3.10/site-packages/torch/nn/functional.py:5168:18
    )  (type 'Tensor')
```

So I am wondering if I should submit a PR for the things that I could successfully circumvent (so two of the three things I mentioned), and on the other hand, if there are suggestions on how to proceed with the "Unknown rank" problem.

### Versions

```
Collecting environment information...
PyTorch version: 1.13.0+cu117
Is debug build: False
CUDA used to build PyTorch: 11.7
ROCM used to build PyTorch: N/A

OS: CentOS Linux release 7.9.2009 (Core) (x86_64)
GCC version: (GCC) 4.8.5 20150623 (Red Hat 4.8.5-44)
Clang version: Could not collect
CMake version: version 3.10.1
Libc version: glibc-2.17

Python version: 3.10.8 | packaged by conda-forge | (main, Nov 22 2022, 08:23:14) [GCC 10.4.0] (64-bit runtime)
Python platform: Linux-3.10.0-1160.83.1.el7.x86_64-x86_64-with-glibc2.17
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:          x86_64
CPU op-mode(s):        32-bit, 64-bit
Byte Order:            Little Endian
CPU(s):                48
On-line CPU(s) list:   0-47
Thread(s) per core:    1
Core(s) per socket:    24
Socket(s):             2
NUMA node(s):          4
Vendor ID:             GenuineIntel
CPU family:            6
Model:                 85
Model name:            Intel(R) Xeon(R) Platinum 8160 CPU @ 2.10GHz
Stepping:              4
CPU MHz:               2799.957
CPU max MHz:           3700.0000
CPU min MHz:           1000.0000
BogoMIPS:              4200.00
Virtualization:        VT-x
L1d cache:             32K
L1i cache:             32K
L2 cache:              1024K
L3 cache:              33792K
NUMA node0 CPU(s):     0-2,6-8,12-14,18-20
NUMA node1 CPU(s):     3-5,9-11,15-17,21-23
NUMA node2 CPU(s):     24-26,30-32,36-38,42-44
NUMA node3 CPU(s):     27-29,33-35,39-41,45-47
Flags:                 fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc art arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc aperfmperf eagerfpu pni pclmulqdq dtes64 monitor ds_cpl vmx smx est tm2 ssse3 sdbg fma cx16 xtpr pdcm pcid dca sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch epb cat_l3 cdp_l3 invpcid_single intel_ppin ssbd mba rsb_ctxsw ibrs ibpb stibp tpr_shadow vnmi flexpriority ept vpid fsgsbase tsc_adjust bmi1 hle avx2 smep bmi2 erms invpcid rtm cqm mpx rdt_a avx512f avx512dq rdseed adx smap clflushopt clwb intel_pt avx512cd avx512bw avx512vl xsaveopt xsavec xgetbv1 cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local dtherm ida arat pln pts hwp hwp_act_window hwp_epp hwp_pkg_req pku ospke md_clear spec_ctrl intel_stibp flush_l1d arch_capabilities

Versions of relevant libraries:
[pip3] numpy==1.23.5
[pip3] torch==1.13.0
[conda] numpy                     1.23.5                   pypi_0    pypi
[conda] torch                     1.13.0                   pypi_0    pypi
```