Repository: pytorch/pytorch

Issue:
Tensor indexing issue in onnx
## 🐛 Bug
This is a sibling issue for https://github.com/microsoft/onnxruntime/issues/6910 as they suggested to report here too.

It seems that Tensor indexing is not fully supported once exported to ONNX:

## To Reproduce

```py
import io
import torch
from torch import Tensor
import onnxruntime

def f() -> Tensor:
    mask = torch.zeros(100, dtype=torch.bool)
    indices = (torch.rand(25) * mask.shape[0]).to(torch.int64)
    mask[indices] = True  # offending line
    return mask

class Module(torch.nn.Module):
    def forward(self, *args, **kwargs):
        return f()

model = Module()
model.eval()

model()  # works fine

onnx_io = io.BytesIO()
torch.onnx.export(model, [], onnx_io, opset_version=11)

ort_session = onnxruntime.InferenceSession(onnx_io.getvalue())
ort_outs = ort_session.run(None, {})  # errors
```

```
/Users/nicolashug/opt/miniconda3/envs/pt/lib/python3.8/site-packages/torch/onnx/utils.py:347: UserWarning: No input args
  warnings.warn("No input args")
2021-03-05 14:58:13.338019 [E:onnxruntime:, inference_session.cc:1293 operator()] Exception during initialization: /Users/runner/work/1/s/onnxruntime/core/providers/cpu/tensor/reshape_helper.h:43 onnxruntime::ReshapeHelper::ReshapeHelper(const onnxruntime::TensorShape &, std::vector<int64_t> &) gsl::narrow_cast<int64_t>(input_shape.Size()) == size was false. The input tensor cannot be reshaped to the requested shape. Input shape:{}, requested shape:{25}

Traceback (most recent call last):
  File "lol.py", line 24, in <module>
    ort_session = onnxruntime.InferenceSession(onnx_io.getvalue())
  File "/Users/nicolashug/opt/miniconda3/envs/pt/lib/python3.8/site-packages/onnxruntime/capi/onnxruntime_inference_collection.py", line 280, in __init__
    self._create_inference_session(providers, provider_options)
  File "/Users/nicolashug/opt/miniconda3/envs/pt/lib/python3.8/site-packages/onnxruntime/capi/onnxruntime_inference_collection.py", line 312, in _create_inference_session
    sess.initialize_session(providers, provider_options)
onnxruntime.capi.onnxruntime_pybind11_state.RuntimeException: [ONNXRuntimeError] : 6 : RUNTIME_EXCEPTION : Exception during initialization: /Users/runner/work/1/s/onnxruntime/core/providers/cpu/tensor/reshape_helper.h:43 onnxruntime::ReshapeHelper::ReshapeHelper(const onnxruntime::TensorShape &, std::vector<int64_t> &) gsl::narrow_cast<int64_t>(input_shape.Size()) == size was false. The input tensor cannot be reshaped to the requested shape. Input shape:{}, requested shape:{25}
```

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->



## Environment

Collecting environment information...
PyTorch version: 1.8.0a0+ad7d208
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 10.15.7 (x86_64)
GCC version: Could not collect
Clang version: 12.0.0 (clang-1200.0.32.29)
CMake version: version 3.18.2

Python version: 3.8 (64-bit runtime)
Is CUDA available: False
CUDA runtime version: No CUDA
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

Versions of relevant libraries:
[pip3] numpy==1.19.2
[pip3] pytorch-sphinx-theme==0.0.24
[pip3] torch==1.8.0a0
[pip3] torchaudio==0.8.0a0+2c8aad9
[pip3] torchtext==0.9.0a0+651c1f7
[pip3] torchvision==0.9.0a0+1438b0c
[conda] blas                      1.0                         mkl
[conda] mkl                       2019.4                      233
[conda] mkl-include               2020.2                      260
[conda] mkl-service               2.3.0            py38h9ed2024_0
[conda] mkl_fft                   1.3.0            py38ha059aab_0
[conda] mkl_random                1.1.1            py38h959d312_0
[conda] numpy                     1.19.2           py38h456fd55_0
[conda] numpy-base                1.19.2           py38hcfb5961_0
[conda] pytorch-sphinx-theme      0.0.24                    dev_0    <develop>
[conda] torch                     1.8.0a0                  pypi_0    pypi
[conda] torchaudio                0.8.0a0+f2da586          pypi_0    pypi
[conda] torchtext                 0.9.0a0+651c1f7           dev_0    <develop>
[conda] torchvision               0.9.0a0+1438b0c           dev_0    <develop>

cc @houseroad @spandantiwari @lara-hdr @BowenBao @neginraoof @SplitInfinity


Discussion:
@NicolasHug PR #53690 should resolve this issue, could you please confirm whether these changes resolved your issue

Resolving patch:
diff --git a/torch/onnx/symbolic_helper.py b/torch/onnx/symbolic_helper.py
--- a/torch/onnx/symbolic_helper.py
+++ b/torch/onnx/symbolic_helper.py
@@ -296,6 +296,11 @@ def _is_fp(value):
             return (type == 'Float') or (type == 'Double') or (type == 'Half')
     return False
 
+def _dtype_is_fp(type_value):
+    if type_value:
+        return (type_value == torch.float16) or (type_value == torch.float32) or (type_value == torch.float64)
+    return False
+
 def _generate_wrapped_number(g, scalar):
     """
     Create a wrapped number based on https://github.com/pytorch/pytorch/issues/9515
diff --git a/torch/onnx/symbolic_opset9.py b/torch/onnx/symbolic_opset9.py
--- a/torch/onnx/symbolic_opset9.py
+++ b/torch/onnx/symbolic_opset9.py
@@ -3014,3 +3014,22 @@ def linear(g, input, weight, bias):
             output = add(g, bias, output)
 
     return output
+
+@parse_args('v', 'b', 'i', 'v', 'v', 'v', 'v')
+def hann_window(g, window_length, periodic=True, dtype=None, layout=None, device=None, pin_memory=None, requires_grad=False):
+    if dtype is None:
+        dtype = torch.get_default_dtype()
+        if sym_help._dtype_is_fp(dtype) is False:
+            dtype = torch.float
+        dtype = sym_help.scalar_type_to_pytorch_type.index(dtype)
+
+    n_array = arange(g, window_length, 4, None, None, None)
+    output = g.op('Cast', n_array, to_i=sym_help.cast_pytorch_to_onnx['Float'])
+    output = mul(g, g.op('Constant', value_t=torch.tensor(math.pi, dtype=torch.float)), output)
+
+    if periodic is False:
+        window_length = sub(g, window_length, g.op("Constant", value_t=torch.tensor(1, dtype=torch.int)))
+    output = div(g, output, window_length)
+    output = g.op("Cast", square(g, sin(g, output)), to_i=sym_help.scalar_type_to_onnx[dtype])
+
+    return output


Test patch:
diff --git a/test/onnx/test_pytorch_onnx_onnxruntime.py b/test/onnx/test_pytorch_onnx_onnxruntime.py
--- a/test/onnx/test_pytorch_onnx_onnxruntime.py
+++ b/test/onnx/test_pytorch_onnx_onnxruntime.py
@@ -7777,6 +7777,60 @@ def forward(self, input_ids):
         self.run_test(M(), (x,), input_names=['input_ids'],
                       dynamic_axes={'input_ids': {0: 'batch', 1: 'sequence'}})
 
+    @skipIfUnsupportedMinOpsetVersion(9)
+    def test_hann_window_periodic(self):
+        class HannWindowModule_Periodic(torch.nn.Module):
+            def __init__(self):
+                super(HannWindowModule_Periodic, self).__init__()
+                self.window_length = 0
+
+            def forward(self, x, window_length: int):
+                self.window_length = window_length
+                return torch.add(x, torch.hann_window(self.window_length, periodic=True, dtype=torch.float))
+
+        win_length = 100
+        x = torch.randn(win_length)
+
+        module = HannWindowModule_Periodic()
+        self.run_test(module, (x, win_length))
+
+    @skipIfUnsupportedMinOpsetVersion(9)
+    def test_hann_window_not_periodic(self):
+        class HannWindowModule_NotPeriodic(torch.nn.Module):
+            def __init__(self):
+                super(HannWindowModule_NotPeriodic, self).__init__()
+                self.window_length = 0
+
+            def forward(self, x, window_length: int):
+                self.window_length = window_length
+                return torch.add(x, torch.hann_window(self.window_length, periodic=False, dtype=torch.float))
+
+        win_length = 100
+        x = torch.randn(win_length)
+
+        module = HannWindowModule_NotPeriodic()
+        self.run_test(module, (x, win_length))
+
+    @skipIfUnsupportedMinOpsetVersion(9)
+    @disableScriptTest()
+    def test_hann_window_default_values(self):
+        class HannWindowModule(torch.nn.Module):
+            def __init__(self):
+                super(HannWindowModule, self).__init__()
+                self.window_length = 0
+
+            def forward(self, x, window_length: int):
+                import torch.nn.functional as F
+                self.window_length = window_length
+                return torch.add(x, F.relu(torch.hann_window(self.window_length)))
+
+        win_length = 100
+        x = torch.randn(win_length, dtype=torch.float)
+        module = HannWindowModule()
+
+        output = module(x, win_length)
+        self.run_test(module, (x, win_length))
+
 def make_test(name, base, layer, bidirectional, initial_state,
               variable_length, dropout,
               **extra_kwargs):