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:

Resolving patch:
diff --git a/torch/onnx/symbolic_opset11.py b/torch/onnx/symbolic_opset11.py
--- a/torch/onnx/symbolic_opset11.py
+++ b/torch/onnx/symbolic_opset11.py
@@ -136,6 +136,10 @@ def index_put(g, self, indices_list_value, values, accumulate=False):
     sub_data_shape = sym_help._slice_helper(
         g, g.op("Shape", self), axes=[0], starts=[len(indices_list)], ends=[maxsize])
     values_shape = g.op("Concat", broadcast_index_shape, sub_data_shape, axis_i=0)
+    # Check if values is a singular value and expand accordingly
+    rank = sym_help._get_tensor_rank(values)
+    if rank is not None and rank == 0:
+        values = expand(g, values, values_shape, None)
     values = g.op("Reshape", values, values_shape)
 
     if accumulate:


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
@@ -1754,6 +1754,26 @@ def forward(self, x, ind, update):
         update = torch.ones(4)
         self.run_test(IndexPutModel(), (x, ind, update))
 
+    @skipIfUnsupportedMinOpsetVersion(11)
+    def test_index_put_singular(self):
+        class IndexPutBoolModel(torch.nn.Module):
+            def forward(self, mask, indices):
+                mask[indices] = True
+                return mask
+
+        mask = torch.zeros(100, dtype=torch.bool)
+        indices = (torch.rand(25) * mask.shape[0]).to(torch.int64)
+        self.run_test(IndexPutBoolModel(), (mask, indices))
+
+        class IndexPutFloatModel(torch.nn.Module):
+            def forward(self, mask, indices):
+                mask[indices] = torch.tensor(5.5)
+                return mask
+
+        mask = torch.rand(100, dtype=torch.float)
+        indices = (torch.rand(50) * mask.shape[0]).to(torch.int64)
+        self.run_test(IndexPutFloatModel(), (mask, indices))
+
     @skipIfUnsupportedMinOpsetVersion(11)
     def test_index_put_accumulate(self):
         class IndexPutModel(torch.nn.Module):