Repository: pytorch/pytorch

Issue:
ONNX export fails with a `resolve_conj` op when a tensor slice is printed.
### 🐛 Describe the bug

Exporting a module as ONNX fails when the module involves a tensor slice being (maybe also any view) passed to `print()`. The export fails with a message stating that the `resolve_conj` op is not available.

It appears that printing a view invokes `resolve_conj`, which causes ONNX exports to fail even in cases where the branch is extraneous. As far as I understand, the ONNX optimiser can remove these kinds of extraneous branches, but in this case it fails before pruning.

Example:
```python
import torch
from torch import nn

class BadFirst(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        x_firsts = x[:, 0]
        print(f"x_firsts: {x_firsts}") # removing this line fixes the issue.
        return x_firsts


m = BadFirst().eval()
x = torch.rand(10, 5)

res = m(x) # this works
torch.onnx.export(m, x, "m.onnx") # this fails due to a resolve_conj op
```
Observed Output:
```
x_firsts: tensor([0.8206, 0.7061, 0.0366, 0.7012, 0.4380, 0.0920, 0.7127, 0.1474, 0.5061,
        0.9991])
x_firsts: tensor([0.8206, 0.7061, 0.0366, 0.7012, 0.4380, 0.0920, 0.7127, 0.1474, 0.5061,
        0.9991])
Traceback (most recent call last):
  File "~\scratch\pytorch\example.py", line 18, in <module>
    torch.onnx.export(m, x, "m.onnx") # this fails due to a resolve_conj op
  File "~\.virtualenvs\pytorch-jk_rFARN\lib\site-packages\torch\onnx\__init__.py", line 316, in export
    return utils.export(model, args, f, export_params, verbose, training,
  File "~\.virtualenvs\pytorch-jk_rFARN\lib\site-packages\torch\onnx\utils.py", line 107, in export
    _export(model, args, f, export_params, verbose, training, input_names, output_names,
  File "~\.virtualenvs\pytorch-jk_rFARN\lib\site-packages\torch\onnx\utils.py", line 724, in _export
    _model_to_graph(model, args, verbose, input_names,
  File "~\.virtualenvs\pytorch-jk_rFARN\lib\site-packages\torch\onnx\utils.py", line 497, in _model_to_graph
    graph = _optimize_graph(graph, operator_export_type,
  File "~\.virtualenvs\pytorch-jk_rFARN\lib\site-packages\torch\onnx\utils.py", line 216, in _optimize_graph
    graph = torch._C._jit_pass_onnx(graph, operator_export_type)
  File "~\.virtualenvs\pytorch-jk_rFARN\lib\site-packages\torch\onnx\__init__.py", line 373, in _run_symbolic_function
    return utils._run_symbolic_function(*args, **kwargs)
  File "~\.virtualenvs\pytorch-jk_rFARN\lib\site-packages\torch\onnx\utils.py", line 1028, in _run_symbolic_function
    symbolic_fn = _find_symbolic_in_registry(domain, op_name, opset_version, operator_export_type)
  File "~\.virtualenvs\pytorch-jk_rFARN\lib\site-packages\torch\onnx\utils.py", line 982, in _find_symbolic_in_registry
    return sym_registry.get_registered_op(op_name, domain, opset_version)
  File "~\.virtualenvs\pytorch-jk_rFARN\lib\site-packages\torch\onnx\symbolic_registry.py", line 125, in get_registered_op
    raise RuntimeError(msg)
RuntimeError: Exporting the operator resolve_conj to ONNX opset version 9 is not supported. Please feel free to request support or submit a pull request on PyTorch GitHub.
```
I think I understand why `resolve_conj` is invoked here, but causing an ONNX export failure feels like an unintentional side-effect, that could potentially be avoided by pruning the graph before invoking the ONNX optimiser.

If `torch.jit.script` is run first, then exporting the ONNX model succeeds, but `torch.jit.trace` fails in the same way.

### Versions

Test 1 (python 3.9.6/torch 1.10.2):
```
PyTorch version: 1.10.2+cpu
Is debug build: False
CUDA used to build PyTorch: Could not collect
ROCM used to build PyTorch: N/A

OS: Microsoft Windows 10 Enterprise
GCC version: Could not collect
Clang version: Could not collect
CMake version: version 3.21.0-rc3
Libc version: N/A

Python version: 3.9.6 (tags/v3.9.6:db3ff76, Jun 28 2021, 15:26:21) [MSC v.1929 64 bit (AMD64)] (64-bit runtime)
Python platform: Windows-10-10.0.19042-SP0
Is CUDA available: False
CUDA runtime version: 10.0.130
GPU models and configuration: GPU 0: Quadro P400
Nvidia driver version: 461.55

cuDNN version: C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v10.0\bin\cudnn64_7.dll
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.22.2
[pip3] torch==1.10.2
[conda] Could not collect
```
Test 2 (python 3.9.10/torch 1.9.1):
```
PyTorch version: 1.9.1+cpu
Is debug build: False
CUDA used to build PyTorch: Could not collect
ROCM used to build PyTorch: N/A

OS: Microsoft Windows 10 Enterprise
GCC version: Could not collect
Clang version: Could not collect
CMake version: Could not collect
Libc version: N/A

Python version: 3.9.10 (tags/v3.9.10:f2f3f53, Jan 17 2022, 15:14:21) [MSC v.1929 64 bit (AMD64)] (64-bit runtime)
Python platform: Windows-10-10.0.19042-SP0
Is CUDA available: False
CUDA runtime version: Could not collect
GPU models and configuration: GPU 0: Quadro P2000
Nvidia driver version: 462.30
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.22.2
[pip3] torch==1.9.1
[conda] Could not collect
```


Discussion:
This is being tracked internally by Microsoft at https://msdata.visualstudio.com/Vienna/_workitems/edit/1685669
I am not able to open the Microsoft link, but I am encountering the same error when exporting the `MarianMTModel`, specifically `Helsinki-NLP/opus-mt-fr-en`, from the HF transformers package to onnx.

I seem to be unable to find the line(s) causing it, though. Is there a work-around, besides `torch.jit.script`, which is not possible for different reasons?

Additionally, @mvhv, can you explain why this function is triggered with the print statement?
@Kroshtan `torch.conj()` creates a view of a tensor with the conjugate bit set to avoid excessive copying. If a later op can't operate on a view then `torch.resolve_conj()` is called which materialises the conjugate on-demand.

I don't know the exact workings (and don't really care enough to check further), but my guess is that whatever string formatter is used for displaying tensors makes a copy to avoid hogging memory from the runtime while blocking on IO, and that copy triggers resolve_conj because you may as well materialise it now in-case you need it later.

The ops should just be pruned off the tree by ONNX, but it can't know that until it's loaded the graph, and it can't load the graph while it has unknown ops in it.
@BowenBao reproducible in nightly

```
Torch IR graph at exception: graph(%0 : Float(10, 5, strides=[5, 1], requires_grad=0, device=cpu)):
  %174 : Long(device=cpu) = prim::Constant[value={0}](), scope: __main__.BadFirst::
  %175 : Long(device=cpu) = prim::Constant[value={9223372036854775807}](), scope: __main__.BadFirst::
  %176 : Long(device=cpu) = prim::Constant[value={1}](), scope: __main__.BadFirst::
  %93 : Float(10, 5, strides=[5, 1], requires_grad=0, device=cpu) = aten::slice(%0, %174, %174, %175, %176), scope: __main__.BadFirst:: # <ipython-input-3-cc5d0ee6a95d>:9:0
  %inp : Float(10, strides=[5], requires_grad=0, device=cpu) = aten::select(%93, %176, %174), scope: __main__.BadFirst:: # <ipython-input-3-cc5d0ee6a95d>:9:0
  %172 : Float(10, strides=[5], requires_grad=0, device=cpu) = aten::resolve_conj(%inp), scope: __main__.BadFirst:: # pytorch/torch/_tensor_str.py:232:0
  %173 : Float(10, strides=[5], requires_grad=0, device=cpu) = aten::resolve_neg(%172), scope: __main__.BadFirst:: # pytorch/torch/_tensor_str.py:232:0
  return (%173)
```

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
@@ -31,6 +31,7 @@
     "check_training_mode",
     "dequantize_helper",
     "is_caffe2_aten_fallback",
+    "is_complex_value",
     "parse_args",
     "pytorch_name_to_type",
     "quantize_helper",
@@ -503,6 +504,16 @@ def _is_tuple_construct(x: _C.Value) -> bool:
     return x.node().kind() == "prim::TupleConstruct"
 
 
+@_beartype.beartype
+def is_complex_value(x: _C.Value) -> bool:
+    assert _is_value(x)
+    return _type_utils.JitScalarType.from_name(x.type().scalarType()) in {
+        _type_utils.JitScalarType.COMPLEX32,
+        _type_utils.JitScalarType.COMPLEX64,
+        _type_utils.JitScalarType.COMPLEX128,
+    }
+
+
 @_beartype.beartype
 def is_caffe2_aten_fallback() -> bool:
     return (
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
@@ -268,6 +268,8 @@
     "unsafe_split_with_sizes",
     "unsafe_split",
     "unsqueeze",
+    "unsupported_complex_operators",
+    "noop_complex_operators",
     "unused",
     "var_mean",
     "var",
@@ -6660,3 +6662,35 @@ def onnx_placeholder(g: jit_utils.GraphContext, *inputs, **attrs):
     env = g.env
 
     return torch._C._jit_onnx_convert_pattern_from_subblock(block, node, env)
+
+
+@_onnx_symbolic("aten::resolve_conj")
+@_onnx_symbolic("aten::resolve_neg")
+@_beartype.beartype
+def noop_complex_operators(g: jit_utils.GraphContext, input: _C.Value):
+    # ONNX does not have operators to *directly* manipulate real/imaginary components
+    # However, a few torch APIs (e.g. .tolist()) use complex operations when input is real,
+    # which results in failures due to missing operators for complex numbers
+
+    # `aten::resolve_conj` and `aten::resolve_neg` can safely be implemented as no-op
+    return input
+
+
+@_onnx_symbolic("aten::_conj")
+@_onnx_symbolic("aten::conj_physical")
+@_beartype.beartype
+def unsupported_complex_operators(g: jit_utils.GraphContext, input: _C.Value):
+    # ONNX does not have operators to *directly* manipulate real/imaginary components
+    # However, a few torch APIs (e.g. .tolist()) use complex operations when input is real,
+    # which results in failures due to missing operators for complex numbers
+
+    # While `aten::_conj` and `aten::conj_phisical` raise exception when input is complex
+    if symbolic_helper.is_complex_value(input):
+        # FIXME(justinchuby): report correct name for symbolic being executed
+        return symbolic_helper._onnx_unsupported(
+            "aten::_conj, aten::conj_physical",
+            input,
+        )
+
+    # they can safely be implemented as no-op for real numbers only
+    return noop_complex_operators(g, input)


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
@@ -12476,6 +12476,33 @@ def forward(self, x):
 
         self.run_test(LerpModel(), torch.rand(5, 4, 3))
 
+    @common_utils.parametrize("input_dtype", [torch.cfloat, torch.float])
+    @skipIfUnsupportedMinOpsetVersion(9)
+    def test_print_tensor_within_torch_nn_module(self, input_dtype: torch.dtype):
+        class PrintTensorOnMyModel(torch.nn.Module):
+            def forward(self, x):
+                # 'print' has side effect calling 'resolve_conj' and 'resolve_neg'.
+                x_firsts = x[:, 0]
+                print(f"x_firsts: {x_firsts}")
+                # 'tolist' has side effect calling 'resolve_conj' and 'resolve_neg'.
+                # Annotation added to pass torch script.
+                _: List[float] = x.tolist()
+                return x_firsts
+
+        m = PrintTensorOnMyModel()
+        x = torch.randn(10, 5, dtype=input_dtype)
+        if input_dtype == torch.cfloat:
+            with self.assertRaises(RuntimeError):
+                self.run_test(
+                    m,
+                    x,
+                )
+        else:
+            self.run_test(
+                m,
+                x,
+            )
+
     # Cannot export with older opsets because of "ConstantFill" op
     # ConstantFill was a temp op removed at opset 8. This is no longer supported by onnxruntime
     # There are still some issues prevent us from enabling script test for these scenarios: