Repository: pytorch
Issue: [onnx] export repeat_interleave TypeError: z_(): incompatible function arguments
Labels: needs reproduction,module: onnx,triaged,onnx-triaged

### 🐛 Describe the bug

I was able to export my model to TorchScript but trying to export it to ONNX gives me a cryptic error I don't understand. I managed to reproduce the error with this small example:
```python
import torch
import torch.nn.functional as F
from torch import nn

print(torch.__version__)

class RepeatedTokenEmbedding(nn.Module):
    def __init__(self, label_count=16):
        super().__init__()
        self.embed = nn.Embedding(label_count, 2)

    def forward(self, x, durs):
        # problem goes away with pad lines removed
        x = F.pad(x,       [0, 2, 0, 0], value=0)
        durs = F.pad(durs, [0, 2, 0, 0], value=0)

        repeats = torch.repeat_interleave(x[0], durs[0]).unsqueeze(0)

        embedding = self.embed(repeats)

        return embedding


rte = RepeatedTokenEmbedding()

x = torch.randint(low=1, high=16, size=(1, 16)).long()
durs = torch.randint(low=1, high=4, size=(1, 16)).long()

out = rte(x, durs)

torch.onnx.export(rte, (x, durs), 'test.onnx', verbose=True)
```

It works with the pad line removed, but breaks again if I try to specify the axis as dynamic like this:
```python
torch.onnx.export(
    rte,
    (x, durs),
    'test.onnx',
    verbose=True,
    input_names=["x", "durs"],
    output_names=["embedding"],
    
    # problem goes away with this removed
    dynamic_axes={
        "x": {1: "sequence"}
    }
)
```


The traceback:
```
1.14.0.dev20221021+cu116
/usr/local/lib/python3.7/dist-packages/torch/onnx/_internal/jit_utils.py:258: UserWarning: Constant folding - Only steps=1 can be constant folded for opset >= 10 onnx::Slice op. Constant folding not applied. (Triggered internally at ../torch/csrc/jit/passes/onnx/constant_fold.cpp:179.)
  _C._jit_pass_onnx_node_shape_type_inference(node, params_dict, opset_version)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
[<ipython-input-1-aedf6a6a2a21>](https://localhost:8080/#) in <module>
     27 out = rte(x, durs)
     28 
---> 29 torch.onnx.export(rte, (x, durs), 'test.onnx', verbose=True)

16 frames
[/usr/local/lib/python3.7/dist-packages/torch/onnx/utils.py](https://localhost:8080/#) in export(model, args, f, export_params, verbose, training, input_names, output_names, operator_export_type, opset_version, do_constant_folding, dynamic_axes, keep_initializers_as_inputs, custom_opsets, export_modules_as_functions)
    517         keep_initializers_as_inputs=keep_initializers_as_inputs,
    518         custom_opsets=custom_opsets,
--> 519         export_modules_as_functions=export_modules_as_functions,
    520     )
    521 

[/usr/local/lib/python3.7/dist-packages/torch/onnx/utils.py](https://localhost:8080/#) in _export(model, args, f, export_params, verbose, training, input_names, output_names, operator_export_type, export_type, opset_version, do_constant_folding, dynamic_axes, keep_initializers_as_inputs, fixed_batch_size, custom_opsets, add_node_names, onnx_shape_inference, export_modules_as_functions)
   1537                 fixed_batch_size=fixed_batch_size,
   1538                 training=training,
-> 1539                 dynamic_axes=dynamic_axes,
   1540             )
   1541 

[/usr/local/lib/python3.7/dist-packages/torch/onnx/utils.py](https://localhost:8080/#) in _model_to_graph(model, args, verbose, input_names, output_names, operator_export_type, do_constant_folding, _disable_torch_constant_prop, fixed_batch_size, training, dynamic_axes)
   1121             dynamic_axes=dynamic_axes,
   1122             input_names=input_names,
-> 1123             module=module,
   1124         )
   1125     except Exception as e:

[/usr/local/lib/python3.7/dist-packages/torch/onnx/utils.py](https://localhost:8080/#) in _optimize_graph(graph, operator_export_type, _disable_torch_constant_prop, fixed_batch_size, params_dict, dynamic_axes, input_names, module)
    661     _C._jit_pass_onnx_lint(graph)
    662 
--> 663     graph = _C._jit_pass_onnx(graph, operator_export_type)
    664     _C._jit_pass_onnx_lint(graph)
    665     _C._jit_pass_lint(graph)

[/usr/local/lib/python3.7/dist-packages/torch/onnx/utils.py](https://localhost:8080/#) in _run_symbolic_function(graph, block, node, inputs, env, operator_export_type)
   1865                     k: symbolic_helper._node_get(node, k) for k in node.attributeNames()
   1866                 }
-> 1867                 return symbolic_fn(graph_context, *inputs, **attrs)
   1868 
   1869         attrs = {

[/usr/local/lib/python3.7/dist-packages/torch/onnx/symbolic_opset13.py](https://localhost:8080/#) in repeat_interleave(g, self, repeats, dim, output_size)
    536     # If input size is dynamic or repeats vector is dynamic
    537     if output_sizes[dim] == 0 or cond_dynamic_repeats:
--> 538         reps = symbolic_helper._size_helper(g, input, dim)
    539         reps = opset11.unsqueeze(g, reps, 0)
    540         # Check if repeats vector is a single integer value

[/usr/local/lib/python3.7/dist-packages/torch/onnx/symbolic_helper.py](https://localhost:8080/#) in _size_helper(g, self, dim)
   1307     from torch.onnx.symbolic_opset9 import select
   1308 
-> 1309     return select(g, full_shape, g.op("Constant", value_t=torch.tensor([0])), dim)
   1310 
   1311 

[/usr/local/lib/python3.7/dist-packages/torch/onnx/symbolic_helper.py](https://localhost:8080/#) in wrapper(g, *args, **kwargs)
    378                 for descriptor, arg in descriptor_args
    379             ):
--> 380                 return fn(g, *args, **kwargs)
    381 
    382             # Dequantize arguments that are quantized

[/usr/local/lib/python3.7/dist-packages/torch/onnx/symbolic_helper.py](https://localhost:8080/#) in wrapper(g, *args, **kwargs)
    301                     f"{FILE_BUG_MSG}"
    302                 )
--> 303             return fn(g, *args, **kwargs)
    304 
    305         return wrapper

[/usr/local/lib/python3.7/dist-packages/torch/onnx/symbolic_opset9.py](https://localhost:8080/#) in select(g, self, dim, index)
   1138     else:
   1139         # FIXME(justinchuby): can index be an int and not a value?
-> 1140         return g.op("Gather", self, index, axis_i=dim)
   1141 
   1142 

[/usr/local/lib/python3.7/dist-packages/torch/onnx/_internal/jit_utils.py](https://localhost:8080/#) in op(self, opname, outputs, *raw_args, **kwargs)
     83         """
     84         # FIXME(justinchuby): Add the return type back once we know how to handle mypy
---> 85         return _add_op(self, opname, *raw_args, outputs=outputs, **kwargs)
     86 
     87     @_beartype.beartype

[/usr/local/lib/python3.7/dist-packages/torch/onnx/_internal/jit_utils.py](https://localhost:8080/#) in _add_op(graph_context, opname, outputs, *args, **kwargs)
    187         keyword argument for multi-return nodes).
    188     """
--> 189     inputs = [_const_if_tensor(graph_context, arg) for arg in args]
    190     # Filter out None attributes, this can be convenient client side because
    191     # now they can pass through None attributes, and have them not show up

[/usr/local/lib/python3.7/dist-packages/torch/onnx/_internal/jit_utils.py](https://localhost:8080/#) in <listcomp>(.0)
    187         keyword argument for multi-return nodes).
    188     """
--> 189     inputs = [_const_if_tensor(graph_context, arg) for arg in args]
    190     # Filter out None attributes, this can be convenient client side because
    191     # now they can pass through None attributes, and have them not show up

[/usr/local/lib/python3.7/dist-packages/torch/onnx/_internal/jit_utils.py](https://localhost:8080/#) in _const_if_tensor(graph_context, arg)
    218         return arg
    219 
--> 220     return _add_op(graph_context, "onnx::Constant", value_z=arg)
    221 
    222 

[/usr/local/lib/python3.7/dist-packages/torch/onnx/_internal/jit_utils.py](https://localhost:8080/#) in _add_op(graph_context, opname, outputs, *args, **kwargs)
    203         opset_version=graph_context.opset,
    204         n_outputs=outputs,
--> 205         shape_inference=GLOBALS.onnx_shape_inference,
    206     )
    207 

[/usr/local/lib/python3.7/dist-packages/torch/onnx/_internal/jit_utils.py](https://localhost:8080/#) in _create_node(graph_or_block, domain_op, inputs, attributes, params_dict, opset_version, n_outputs, shape_inference)
    254         if key in _SKIP_NODE_ATTRIBUTES:
    255             continue
--> 256         _add_attribute(node, key, value, aten=aten)
    257     if shape_inference:
    258         _C._jit_pass_onnx_node_shape_type_inference(node, params_dict, opset_version)

[/usr/local/lib/python3.7/dist-packages/torch/onnx/_internal/jit_utils.py](https://localhost:8080/#) in _add_attribute(node, key, value, aten)
    307             else:
    308                 kind = "i"
--> 309     return getattr(node, f"{kind}_")(name, value)
    310 
    311 

TypeError: z_(): incompatible function arguments. The following argument types are supported:
    1. (self: torch._C.Node, arg0: str, arg1: at::Tensor) -> torch._C.Node

Invoked with: %41 : Tensor = onnx::Constant(), scope: __main__.RepeatedTokenEmbedding::
, 'value', 0 
(Occurred when translating repeat_interleave).
```

### Versions

```
Collecting environment information...
PyTorch version: 1.14.0.dev20221021+cu116
Is debug build: False
CUDA used to build PyTorch: 11.6
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.6 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: 6.0.0-1ubuntu2 (tags/RELEASE_600/final)
CMake version: version 3.22.6
Libc version: glibc-2.26

Python version: 3.7.15 (default, Oct 12 2022, 19:14:55)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.10.133+-x86_64-with-Ubuntu-18.04-bionic
Is CUDA available: False
CUDA runtime version: 11.2.152
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: Could not collect
Nvidia driver version: Could not collect
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.1.1
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.1.1
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.1.1
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.1.1
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.1.1
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.1.1
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.1.1
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.21.6
[pip3] pytorch-lightning==1.7.7
[pip3] torch==1.14.0.dev20221021+cu116
[pip3] torch-stft==0.1.4
[pip3] torchaudio==0.14.0.dev20221021+cu116
[pip3] torchmetrics==0.10.0
[pip3] torchsummary==1.5.1
[pip3] torchtext==0.14.0.dev20221021
[pip3] torchvision==0.15.0.dev20221021+cu116
[conda] Could not collect
```