Repository: pytorch/pytorch
Pull request: [ONNX] Add warning for inplace updates on tensor.shape in tracing mode
Related to #56162


Problem statement:
Tracer vs eager mode difference on tensor.shape
## 🐛 Bug

In eager mode, shape values obtained from tensor.shape are python int, so there are treated as different values. However in tracing both are traced as tensors, and then they share same memory. This causes mismatch in final outputs

## To Reproduce

```
# Tracer vs Eager mode difference in tensor.shape
import torch
import onnxruntime

class Model(torch.nn.Module):
    def forward(self, hidden_states):
        batch_size, seq_length = hidden_states.shape[:2]
        real_seq_length = seq_length
        real_seq_length += 2                                            => real_seq_length and seq_legth should have different values
        return real_seq_length + seq_length


# Call export

# hidden_states is tensor of rank 2
hidden_states = torch.randn(2, 3, 5)
model = Model()

import io
onnx_io = io.BytesIO()
tracemodel = torch.jit.trace(model, (hidden_states, ))
torch.onnx.export(model, (hidden_states, ), onnx_io, verbose=True)

# Call onnxruntime runtime and compare outputs on dynamic inputs
ort_session = onnxruntime.InferenceSession(onnx_io.getvalue())

print(model(hidden_states))
print(tracemodel(hidden_states))
print(ort_session.run(None, {
                                ort_session.get_inputs()[0].name: hidden_states.numpy()
                            }))
```


## Expected behavior

**Output in eager mode**
8 (real_seq_length(=3) + seq_length(=5))

**Output in tracing mode**
Should be 8 but is 10 (real_seq_length(=5) + seq_length(=5)) as tensors share same memory

**Graph output on exporting ONNX graph**
```
graph(%0 : Float(2, 3, 5, strides=[15, 5, 1], requires_grad=0, device=cpu)):
  %1 : Long(3, strides=[1], device=cpu) = onnx::Shape(%0)
  %2 : Long(device=cpu) = onnx::Constant[value={1}]()
  %3 : Long(device=cpu) = onnx::Gather[axis=0](%1, %2) # test/onnx/test_m.py:38:0
  %4 : Long(requires_grad=0, device=cpu) = onnx::Constant[value={2}]()
  %5 : Long(requires_grad=0, device=cpu) = onnx::Add(%3, %4) # test/onnx/test_m.py:40:0
  %6 : Long(requires_grad=0, device=cpu) = onnx::Add(%5, %5) # test/onnx/test_m.py:41:0   ===> %5 proves that both seq_length and real_seq_length have same values/memory
  return (%6)
```

cc @gmagogsfm @houseroad @spandantiwari @lara-hdr @BowenBao @neginraoof @SplitInfinity

Hints from issue discussion:
Thanks for reporting this issue, I think you analysis is correct. Converting and treating shape (int) to Tensor does have unintended effect.

Here is IR from `jit.trace`, it shows that real_seq_length and seq_length are treated as same value indeed:
```
graph(%self : __torch__.Model,
      %hidden_states : Float(2, 3, 5, strides=[15, 5, 1], requires_grad=0, device=cpu)):
  %7 : int = prim::Constant[value=1]() # trace_shape.py:6:0
  %8 : int = aten::size(%hidden_states, %7) # trace_shape.py:6:0
  %seq_length.1 : Long(device=cpu) = prim::NumToTensor(%8)
  %13 : Long(requires_grad=0, device=cpu) = prim::Constant[value={2}]() # trace_shape.py:8:0
  %14 : int = prim::Constant[value=1]() # trace_shape.py:8:0
  %seq_length : Long(requires_grad=0, device=cpu) = aten::add_(%seq_length.1, %13, %14) # trace_shape.py:8:0
  %16 : int = prim::Constant[value=1]() # trace_shape.py:9:0
  %17 : Long(requires_grad=0, device=cpu) = aten::add(%seq_length, %seq_length, %16) # trace_shape.py:9:0
  return (%17)
```

@jamesr66a Rumor says you are expert on `jit.trace`, any suggestions here?
Some possible action items:

1. Check why `jit.trace` sanity check doesn't detect this error.
2. Maybe we can do some post-processing after trace to rewrite patterns like %7 and %8 in my example from using TensorType to int type and getting rid of inplace updates.

If you need to be unblocked now, the immediate action is to not use augment assignment

`a += 2` -> `a = a + 2`
> @jamesr66a Rumor says you are expert on jit.trace

oh no

> original issue

I think the tracer is stuck between a rock and a hard place in terms of capturing networks with shape expressions at all v.s. capturing corner cases like this correctly. I'm not sure if there's anything fundamental we can do in the tracer to alleviate this. The workaround @gmagogsfm proposed seems reasonable.

Commit message: Modify sections

Patch to review:
diff --git a/docs/source/onnx.rst b/docs/source/onnx.rst
index f307aa251dd6..291a7a5a63b0 100644
--- a/docs/source/onnx.rst
+++ b/docs/source/onnx.rst
@@ -183,6 +183,25 @@ Using the Tensor.data field can produce an incorrect trace and therefore an inco
 Use :func:`torch.Tensor.detach` instead. (Work is ongoing to
 `remove Tensor.data entirely <https://github.com/pytorch/pytorch/issues/30987>`_).
 
+Avoid Inplace Operations when using tensor.shape in tracing mode
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+In tracing mode, shape values obtained from tensor.shape are traced as tensors,
+and share the same memory. This might cause a mismatch in values of the final outputs.
+As a workaround, avoid use of inplace operations in these scenarios.
+For example, in the model::
+
+    class Model(torch.nn.Module):
+      def forward(self, states):
+          batch_size, seq_length = states.shape[:2]
+          real_seq_length = seq_length
+          real_seq_length += 2
+          return real_seq_length + seq_length
+
+``real_seq_length`` and ``seq_length`` share the same memory in tracing mode.
+This could be avoiding by rewriting the inplace operation such that::
+
+    real_seq_length = real_seq_length + 2
 
 Limitations
 -----------
@@ -537,6 +556,7 @@ Q: Are lists of Tensors exportable to ONNX?
 
   Yes, for ``opset_version`` >= 11, since ONNX introduced the Sequence type in opset 11.
 
+
 Functions
 --------------------------
 .. autofunction:: export
diff --git a/torch/csrc/jit/passes/remove_inplace_ops.cpp b/torch/csrc/jit/passes/remove_inplace_ops.cpp
index 9260f0e5a651..16200b7dfe51 100644
--- a/torch/csrc/jit/passes/remove_inplace_ops.cpp
+++ b/torch/csrc/jit/passes/remove_inplace_ops.cpp
@@ -112,6 +112,16 @@ void ImplicitCastForBinaryInplaceOps(Block* b) {
       if (originalInputs.at(0) == originalInputs.at(1)) {
         continue;
       }
+
+      auto shape_node = originalInputs.at(0)->node();
+      if ((shape_node->kind() == prim::NumToTensor) &&
+          (shape_node->inputs().at(0)->node()->kind() == aten::size)) {
+        std::cerr
+            << "In tracing mode, shape values obtained from tensor.shape might cause a mismatch "
+            << "final output values. Avoid use of inplace operations in this situation."
+            << std::endl;
+      }
+
       TensorTypePtr firstInp_tensor =
           originalInputs.at(0)->type()->cast<TensorType>();
       TensorTypePtr secondInp_tensor =


Code review comments:
diff_hunk: @@ -112,6 +112,16 @@ void ImplicitCastForBinaryInplaceOps(Block* b) {
       if (originalInputs.at(0) == originalInputs.at(1)) {
         continue;
       }
+
+      auto shape_node = originalInputs.at(0)->node();
+      if ((shape_node->kind() == prim::NumToTensor) &&
+          (shape_node->inputs().at(0)->node()->kind() == aten::size)) {
+        std::cerr
+            << "In tracing mode, shape values obtained from tensor.shape might cause a mismatch "
path: torch/csrc/jit/passes/remove_inplace_ops.cpp
text: @user1:
I'd replace this with a very short text and the URL of the doc you added.

Maybe:
"In-place op on output of tensor.shape. See https://pytorch.org/docs/master/onnx.html#..."

Also, any way to indicate the source location (file name and line)?

@author:
I tried accessing the errant node's sourceRange object, but seems like it is empty and usually sourceRange is used to indicate line/file information.

@user1:
Hmm, that's too bad. It'd be really helpful to add that in the future, but no need to block this on doing that.

diff_hunk: @@ -183,6 +183,25 @@ Using the Tensor.data field can produce an incorrect trace and therefore an inco
 Use :func:`torch.Tensor.detach` instead. (Work is ongoing to
 `remove Tensor.data entirely <https://github.com/pytorch/pytorch/issues/30987>`_).
 
+Avoid Inplace Operations when using tensor.shape in tracing mode
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+In tracing mode, shape values obtained from tensor.shape are traced as tensors,
+and share the same memory. This might cause a mismatch in values of the final outputs.
+As a workaround, avoid use of inplace operations in these scenarios.
+For example, in the model::
+
+    class Model(torch.nn.Module):
+      def forward(self, states):
+          batch_size, seq_length = states.shape[:2]
+          real_seq_length = seq_length
+          real_seq_length += 2
+          return real_seq_length + seq_length
+
+``real_seq_length`` and ``seq_length`` share the same memory in tracing mode.
+This could be avoiding by rewriting the inplace operation such that::
path: docs/source/onnx.rst
text: @user1:
remove "such that"

diff_hunk: @@ -183,6 +183,25 @@ Using the Tensor.data field can produce an incorrect trace and therefore an inco
 Use :func:`torch.Tensor.detach` instead. (Work is ongoing to
 `remove Tensor.data entirely <https://github.com/pytorch/pytorch/issues/30987>`_).
 
+Avoid Inplace Operations when using tensor.shape in tracing mode
path: docs/source/onnx.rst
text: @user1:
capitalize only the first word, and write "in-place" hyphenated
`Avoid in-place operations when using tensor.shape in tracing mode`