Repository: pytorch
Issue: [ONNX] Using values from a different tensor to index a tensor returns a tensor with incorrect shape in exported ONNX model
Labels: oncall: jit,module: onnx,triaged,bug

### 🐛 Describe the bug

# Bug Description

**Note: The complete code for reproducing the issue is at the very bottom.**

**PyTorch Version: 1.12.1 (installed using the following command)**
```
pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu116
```

When I use the values of a different tensor to index a tensor, the returned result has the incorrect shape. For example, for the following two tensors:
```
input = torch.tensor([[[0, 4, 2],
                       [3, 0, 7]],

                       [[0, 4, 6],
                        [3, 7, 5]]], dtype=torch.int32)     # shape: (2, 2, 3)

indices = torch.tensor([[2, 0, 1]], dtype=torch.int64)      # shape: (1, 3)
```
the PyTorch expression `input[:, :, indices[0, 0]]` returns a tensor with the shape of (2, 2). However, when I export the TorchScript of this expression into ONNX model, the ONNX model's inference output has the shape of (2, 2, 1). The `SimpleFailureModel` module in the code at the very bottom is a minimal example that is used to reproduced this issue.

This issue may seem trivial and may seem to have limited impact, but when it's used in conjunction with other functions or operators, it could trigger consequential runtime issues during ONNX inference. For example, if I were to use `torch.cat()` on `input[:, :, indices[0, 0]]`, `input[:, :, indices[0, 1]]`, `input[:, :, indices[0, 2]]` to concatenate the three tensors along `dim=-1` (full expression being `torch.cat((input[:, :, indices[0, 0]], input[:, :, indices[0, 1]], input[:, :, indices[0, 2]]), dim=-1)`), the correct shape of the returned tensor in PyTorch should be (2, 6), and ONNX Runtime, on the other hand, would return a tensor with the incorrect shape of (2, 2, 3). And now if I were to access the element at indices `[0, 4]` in the resulting concatenated tensor, PyTorch would return perfectly valid results while ONNX Runtime would return an "indices element out of data bounds" runtime error during inference because in ONNX the concatenated result **only has size of 2 in axis-1** and we are trying to access the element at index 0 in axis-0 and **index 4 in axis-1**. This is the direct result of the incorrect shape issues mentioned earlier when we try to obtain a new tensor by using the values of a different to index a tensor. More specifically, the error message in this example case is the following:
```
[ONNXRuntimeError] : 2 : INVALID_ARGUMENT : Non-zero status code returned while running Gather node. Name:'Gather_19' Status Message: indices element out of data bounds, idx=4 must be within the inclusive range [-2,1]
```
The `RuntimeFailureModel` module in the code at the very bottom is an example that is used to reproduced this runtime error and illustrate how the indexing issue mentioned in earlier can have consequential impact.

# Expected Behavior
The inference output tensor of the exported ONNX model from TorchScript should have the same shape as the one in PyTorch when we use values from a different tensor to index a tensor.

# Interesting Observations On Previous PyTorch Releases
**For experiments on previous PyTorch versions, I used the wheel files from [https://download.pytorch.org/whl/torch/](https://download.pytorch.org/whl/torch/)**
I also tried running the same export and inference code on PyTorch `1.10.0`, `1.10.1`, `1.10.2`, `1.11.0`, `1.12.0`, `1.12.1`, and they all exhibit the same erroneous behavior. However, there's a difference between versions before and after `1.12`. In the `SimpleFailureModel` example, ONNX models exported from TorchScript using PyTorch versions before `1.12` will have the returned shape of `int32[Gatheroutput_dim_0,Gatheroutput_dim_1,Gatheroutput_dim_2]` (visualized in Netron); however, starting from `1.12`, the exported ONNX model will ahve the returned shape of `int32[input_dynamic_axes_1,input_dynamic_axes_2,1]`. Even though all the versions listed above exhibit the incorrect behavior, the problematic effects of the ones from versions before `1.12` tend to be "hidden" when it's used in conjunction with a large number of other operators and functions. But still, the runtime error in `RuntimeFailureModel` is encountered and can be reproduced in all versions listed above. For your convenience, I have also attached the exported ONNX files (both `SimpleFailureModel` and `RuntimeFailureModel`) from TorchScript using all the listed versions above.

# Complete Code
```python
import torch
import torch.nn as nn
import numpy as np
import onnx
import onnxruntime as ort
from typing import List


# A simple and minimal model for reproducing the issue
class SimpleFailureModel(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x, indices):
        return x[:, :, indices[0, 0]]


# A more complex model demonstrating how the errorneous export (from torch.jit.script to ONNX) issue 
# shown in "SimpleFailureModel" above may trigger issues (e.g. "indices element out of data bounds") 
# during inference.
class RuntimeFailureModel(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x, indices):
        a = x[:, :, indices[0, 0]]
        b = x[:, :, indices[0, 1]]
        c = x[:, :, indices[0, 2]]
        cat = torch.cat((a, b, c), dim=-1)
        return cat[0, 4]


TORCH_VERSION = "torch_1.12.1"
SIMPLE_FAILURE_ONNX_PATH = "SimpleFailure_" + TORCH_VERSION + ".onnx"
RUNTIME_FAILURE_ONNX_PATH = "RuntimeFailure_" + TORCH_VERSION + ".onnx"


input = torch.tensor([[[0, 4, 2],
                       [3, 0, 7]],

                       [[0, 4, 6],
                        [3, 7, 5]]], dtype=torch.int32)     # shape: (2, 2, 3)

print("input.size():", input.size())


indices = torch.tensor([[2, 0, 1]], dtype=torch.int64)      # shape: (1, 3)

print("indices.size():", indices.size())


simple_failure_model = SimpleFailureModel()
simple_failure_model_script = torch.jit.script(simple_failure_model)

torch.onnx.export(simple_failure_model_script,                     # the model's script module
                 (input, indices),                                 # model input (or a tuple for multiple inputs)
                 SIMPLE_FAILURE_ONNX_PATH,                         # where to save the model (can be a file or file-like object)
                 export_params=True,                               # store the trained parameter weights inside the model file
                 opset_version=14,                                 # the ONNX version to export the model to
                 input_names = ["input", "indices"],               # the model's input names
                 output_names = ["output"],                        # the model's output names
                 dynamic_axes={"input": [0, 1, 2],
                               "indices": [0, 1]},
                 do_constant_folding=True)


runtime_failure_model = RuntimeFailureModel()
runtime_failure_model_script = torch.jit.script(runtime_failure_model)

torch.onnx.export(runtime_failure_model_script,                    # the model's script module
                 (input, indices),                                 # model input (or a tuple for multiple inputs)
                 RUNTIME_FAILURE_ONNX_PATH,                        # where to save the model (can be a file or file-like object)
                 export_params=True,                               # store the trained parameter weights inside the model file
                 opset_version=14,                                 # the ONNX version to export the model to
                 input_names = ["input", "indices"],               # the model's input names
                 output_names = ["output"],                        # the model's output names
                 dynamic_axes={"input": [0, 1, 2],
                               "indices": [0, 1]},
                 do_constant_folding=True)


EP_LIST = ["CPUExecutionProvider"]  # A list of ONNX execution providers


# SimpleFailureModel Inference
onnx.checker.check_model(onnx.load(SIMPLE_FAILURE_ONNX_PATH))
simple_failure_model_ort_session = ort.InferenceSession(SIMPLE_FAILURE_ONNX_PATH, providers=EP_LIST)

output_exp = simple_failure_model(input, indices)
output_act = simple_failure_model_ort_session.run(None, {"input": input.numpy(), "indices": indices.numpy()})[0]

print("---------- Expected v.s. Actual (SimpleFailureModel) ----------")
print("output_exp.size():", output_exp.size())
print("output_exp.dtype:", output_exp.dtype)
# print("output_exp:", output_exp)

print("output_act.shape:", output_act.shape)
print("output_act.dtype:", output_act.dtype)
# print("output_act:", output_act)



# RuntimeFailureModel Inference
onnx.checker.check_model(onnx.load(RUNTIME_FAILURE_ONNX_PATH))
runtime_failure_model_ort_session = ort.InferenceSession(RUNTIME_FAILURE_ONNX_PATH, providers=EP_LIST)

output_exp = runtime_failure_model(input, indices)
output_act = runtime_failure_model_ort_session.run(None, {"input": input.numpy(), "indices": indices.numpy()})[0]

print("---------- Expected v.s. Actual (RuntimeFailureModel) ----------")
print("output_exp.size():", output_exp.size())
print("output_exp.dtype:", output_exp.dtype)
# print("output_exp:", output_exp)

print("output_act.shape:", output_act.shape)
print("output_act.dtype:", output_act.dtype)
# print("output_act:", output_act)
```

### Versions

Collecting environment information...
PyTorch version: 1.12.1+cu116
Is debug build: False
CUDA used to build PyTorch: 11.6
ROCM used to build PyTorch: N/A

OS: Microsoft Windows 10 Home
GCC version: Could not collect
Clang version: Could not collect
CMake version: version 3.23.1
Libc version: N/A

Python version: 3.9.12 (tags/v3.9.12:b28265d, Mar 23 2022, 23:52:46) [MSC v.1929 64 bit (AMD64)] (64-bit runtime)
Python platform: Windows-10-10.0.19044-SP0
Is CUDA available: True
CUDA runtime version: 11.6.124
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 3090
Nvidia driver version: 516.59
cuDNN version: C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.6\bin\cudnn_ops_train64_8.dll
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] mypy==0.961
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.22.3
[pip3] torch==1.12.1+cu116
[pip3] torchaudio==0.12.1+cu116
[pip3] torchvision==0.13.1+cu116
[conda] Could not collect