Repository: pytorch/pytorch

Issue:
Improper model conversion from PyTorch to ONNX with torch.onnx.OperatorExportTypes.ONNX_ATEN flag 
### 🐛 Describe the bug

I have spotted improper model conversion to ONNX when using `torch.onnx.OperatorExportTypes.ONNX_ATEN` flag with `PyTorch >=1.12.0`

My understanding from the [documentation](https://pytorch.org/docs/stable/onnx.html#aten-operators)

> OperatorExportTypes.ONNX_ATEN: All ATen ops (in the TorchScript namespace “aten”) are exported as ATen ops (in opset domain “org.pytorch.aten”). [ATen](https://pytorch.org/cppdocs/#aten) is PyTorch’s built-in tensor library, so this instructs the runtime to use PyTorch’s implementation of these ops.

is that when one uses `torch.onnx.OperatorExportTypes.ONNX_ATEN` flag, all ATen operators has to be exported to ONNX graph. When I run the code snippet below: 

```python
import torch


class ModelWithAtenFmod(torch.nn.Module):
    def forward(self, x, y):
        return torch.fmod(x, y)


x = torch.randn(3, 4, dtype=torch.float32)
y = torch.randn(3, 4, dtype=torch.float32)
OUT = torch.onnx.export_to_pretty_string(
    ModelWithAtenFmod(), (x, y),
    add_node_names=False,
    do_constant_folding=False,
    operator_export_type=torch.onnx.OperatorExportTypes.ONNX_ATEN)
print(OUT)
```
with `PyTorch 1.11.0` I obtain ONNX graph with ATen operator that is correct:
```
### PyTorch 1.11.0

ModelProto {
  producer_name: "pytorch"
  domain: ""
  doc_string: ""
  graph:
    GraphProto {
      name: "torch-jit-export"
      inputs: [{name: "aten::ATen_0", type:Tensor dtype: 1, Tensor dims: 3 4},{name: "aten::ATen_1", type:Tensor dtype: 1, Tensor dims: 3 4}]
      outputs: [{name: "2", type:Tensor dtype: 1, Tensor dims: ? ?}]
      value_infos: []
      initializers: []
      nodes: [
        Node {type: "ATen", inputs: [aten::ATen_0,aten::ATen_1], outputs: [2], attributes: [{ name: 'operator', type: string, value: 'fmod'}]}
      ]
    }
  opset_import: [OperatorSetIdProto { domain: , version: 9}OperatorSetIdProto { domain: org.pytorch.aten, version: 1}],
}
```

However, when I run the same code with `PyTorch 1.12.1`, ATen operator got substituted with ONNX native operator and that is not expected behaviour:
```
### PyTorch 1.12.1

ModelProto {
  producer_name: "pytorch"
  domain: ""
  doc_string: ""
  graph:
    GraphProto {
      name: "torch_jit"
      inputs: [{name: "onnx::Mod_0", type:Tensor dtype: 1, Tensor dims: 3 4},{name: "onnx::Mod_1", type:Tensor dtype: 1, Tensor dims: 3 4}]
      outputs: [{name: "2", type:Tensor dtype: 1, Tensor dims: 3 4}]
      value_infos: []
      initializers: []
      nodes: [
        Node {type: "Mod", inputs: [onnx::Mod_0,onnx::Mod_1], outputs: [2], attributes: [{ name: 'fmod', type: int, value: 1}]}
      ]
    }
  opset_import: [OperatorSetIdProto { domain: , version: 13}],
}
```

### Versions

PyTorch version: 1.12.1
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 12.6 (arm64)
GCC version: Could not collect
Clang version: 14.0.6
CMake version: version 3.24.1
Libc version: N/A

Python version: 3.8.13 | packaged by conda-forge | (default, Mar 25 2022, 06:04:14)  [Clang 12.0.1 ] (64-bit runtime)
Python platform: macOS-12.6-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
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
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.23.4
[pip3] torch==1.12.1
[pip3] torchvision==0.13.1
[conda] numpy                     1.23.4           py38h09ac2d9_0    conda-forge
[conda] pytorch                   1.12.1                  py3.8_0    pytorch
[conda] torchvision               0.13.1                 py38_cpu    pytorch


Discussion:
I got the example of the code snippet from one existing unit test:
https://github.com/pytorch/pytorch/blob/664058fa83f1d8eede5d66418abff6e20bd76ca8/test/jit/test_export_modes.py#L107-L121
Operator substitution is happening inside the `_optimize_graph` function in the code below:

https://github.com/pytorch/pytorch/blob/664058fa83f1d8eede5d66418abff6e20bd76ca8/torch/onnx/utils.py#L726-L743

The `graph` before `_optimize_graph` function call looks like:
```
graph(%0 : Float(3, 4, strides=[4, 1], requires_grad=0, device=cpu),
      %1 : Float(3, 4, strides=[4, 1], requires_grad=0, device=cpu)):
  %6 : Float(3, 4, strides=[4, 1], requires_grad=0, device=cpu) = aten::fmod(%0, %1)
  return (%6)
```

And after the call:
```
graph(%0 : Float(3, 4, strides=[4, 1], requires_grad=0, device=cpu),
      %1 : Float(3, 4, strides=[4, 1], requires_grad=0, device=cpu)):
  %2 : Float(3, 4, strides=[4, 1], requires_grad=0, device=cpu) = onnx::Mod[fmod=1](%0, %1)
  return (%2)
```
despite `operator_export_type` being equal to `OperatorExportTypes.ONNX_ATEN`
This is happening inside `graph = _C._jit_pass_onnx(graph, operator_export_type)`
ca374773b4 - [ONNX] update default opset_version to 13 (#73898) introduced this regression. Looking into it further
Fix: https://github.com/pytorch/pytorch/pull/87735

Resolving patch:
diff --git a/torch/onnx/utils.py b/torch/onnx/utils.py
--- a/torch/onnx/utils.py
+++ b/torch/onnx/utils.py
@@ -1739,17 +1739,22 @@ def _add_output_to_block(block: _C.Block, value: _C.Value) -> int:
 
 @_beartype.beartype
 def _should_aten_fallback(
-    name: str,
-    opset_version: int,
-    operator_export_type: _C_onnx.OperatorExportTypes,
+    name: str, opset_version: int, operator_export_type: _C_onnx.OperatorExportTypes
 ):
+    # For BUILD_CAFFE2=0 builds, if domain=="aten" and operator_export_type==ONNX_ATEN,
+    #   an aten::ATen operator is created regardless of symbolics existence
+    # For BUILD_CAFFE2=1, the same applies only if there is no symbolic available
+
     is_exportable_aten_op = registration.registry.is_registered_op(name, opset_version)
     is_onnx_aten_export = operator_export_type == _C_onnx.OperatorExportTypes.ONNX_ATEN
     is_aten_fallback_export = (
         operator_export_type == _C_onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK
     )
-    return is_onnx_aten_export or (
-        not is_exportable_aten_op and is_aten_fallback_export
+    is_caffe2_build = _C_onnx._CAFFE2_ATEN_FALLBACK
+
+    return name.startswith("aten::") and (
+        ((is_onnx_aten_export or is_aten_fallback_export) and not is_caffe2_build)
+        or (not is_exportable_aten_op and is_aten_fallback_export)
     )
 
 
@@ -1844,6 +1849,21 @@ def _run_symbolic_function(
         env=env,
     )
 
+    # Direct ATen export requested
+    if _should_aten_fallback(ns_op_name, opset_version, operator_export_type):
+        attrs = {
+            k + "_" + node.kindOf(k)[0]: symbolic_helper._node_get(node, k)
+            for k in node.attributeNames()
+        }
+        outputs = node.outputsSize()
+        attrs["outputs"] = outputs
+        return graph_context.at(
+            op_name,
+            *inputs,
+            overload_name=_get_aten_op_overload_name(node),
+            **attrs,
+        )
+
     try:
         # Caffe2-specific: Quantized op symbolics are registered for opset 9 only.
         if symbolic_helper.is_caffe2_aten_fallback() and opset_version == 9:
@@ -1861,6 +1881,7 @@ def _run_symbolic_function(
         if symbolic_function_group is not None:
             symbolic_fn = symbolic_function_group.get(opset_version)
             if symbolic_fn is not None:
+                # TODO Wrap almost identical attrs assignment or comment the difference.
                 attrs = {
                     k: symbolic_helper._node_get(node, k) for k in node.attributeNames()
                 }
@@ -1874,18 +1895,6 @@ def _run_symbolic_function(
             # Clone node to trigger ONNX shape inference
             return graph_context.op(op_name, *inputs, **attrs, outputs=node.outputsSize())  # type: ignore[attr-defined]
 
-        if _should_aten_fallback(ns_op_name, opset_version, operator_export_type):
-            # Direct ATen export requested
-            outputs = node.outputsSize()
-            attrs["outputs"] = outputs
-            # `overload_name` is set for non-Caffe2 builds only
-            return graph_context.at(
-                op_name,
-                *inputs,
-                overload_name=_get_aten_op_overload_name(node),
-                **attrs,
-            )
-
         raise errors.UnsupportedOperatorError(
             domain,
             op_name,


Test patch:
diff --git a/test/onnx/test_pytorch_onnx_no_runtime.py b/test/onnx/test_pytorch_onnx_no_runtime.py
--- a/test/onnx/test_pytorch_onnx_no_runtime.py
+++ b/test/onnx/test_pytorch_onnx_no_runtime.py
@@ -32,6 +32,7 @@ def export_to_onnx(
     mocks: Optional[Iterable] = None,
     operator_export_type: torch.onnx.OperatorExportTypes = torch.onnx.OperatorExportTypes.ONNX,
     opset_version: int = GLOBALS.export_onnx_opset_version,
+    **torch_onnx_export_kwargs,
 ) -> onnx.ModelProto:
     """Exports `model(input)` to ONNX and returns it.
 
@@ -44,6 +45,7 @@ def export_to_onnx(
         mocks: list of mocks to use during export
         operator_export_type: export type as described by `torch.onnx.export(...operator_export_type,...)`
         opset_version: ONNX opset version as described by `torch.onnx.export(...opset_version,...)`
+        torch_onnx_export_kwargs: extra torch.onnx.export kwargs arguments
     Returns:
         A valid ONNX model (`onnx.ModelProto`)
     """
@@ -60,6 +62,7 @@ def export_to_onnx(
             f,
             operator_export_type=operator_export_type,
             opset_version=opset_version,
+            **torch_onnx_export_kwargs,
         )
 
     # Validate ONNX graph before returning it
@@ -777,6 +780,91 @@ def forward(self, x):
             model, inputs, f, dynamic_axes={"x": [0, 1]}, input_names=["x"]
         )
 
+    @common_utils.skipIfNoCaffe2
+    def test_caffe2_aten_fallback(self):
+        class ModelWithAtenNotONNXOp(torch.nn.Module):
+            def forward(self, x, y):
+                abcd = x + y
+                defg = torch.linalg.qr(abcd)
+                return defg
+
+        x = torch.rand(3, 4)
+        y = torch.rand(3, 4)
+        f = io.BytesIO()
+        torch.onnx.export(
+            ModelWithAtenNotONNXOp(),
+            (x, y),
+            f,
+            do_constant_folding=False,
+            operator_export_type=torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK,
+            # support for linalg.qr was added in later op set versions.
+            opset_version=9,
+        )
+        onnx_model = onnx.load(io.BytesIO(f.getvalue()))
+        self.assertAtenOp(onnx_model, "linalg_qr")
+
+    @common_utils.skipIfNoCaffe2
+    def test_caffe2_onnx_aten(self):
+        class ModelWithAtenFmod(torch.nn.Module):
+            def forward(self, x, y):
+                return torch.fmod(x, y)
+
+        x = torch.randn(3, 4, dtype=torch.float32)
+        y = torch.randn(3, 4, dtype=torch.float32)
+        f = io.BytesIO()
+        torch.onnx.export(
+            ModelWithAtenFmod(),
+            (x, y),
+            f,
+            do_constant_folding=False,
+            operator_export_type=torch.onnx.OperatorExportTypes.ONNX_ATEN,
+            opset_version=10,  # or higher
+        )
+        onnx_model = onnx.load(io.BytesIO(f.getvalue()))
+        assert onnx_model.graph.node[0].op_type == "Mod"
+
+    @common_utils.skipIfCaffe2
+    def test_aten_fallback(self):
+        class ModelWithAtenNotONNXOp(torch.nn.Module):
+            def forward(self, x, y):
+                abcd = x + y
+                defg = torch.linalg.qr(abcd)
+                return defg
+
+        x = torch.rand(3, 4)
+        y = torch.rand(3, 4)
+        f = io.BytesIO()
+        torch.onnx.export(
+            ModelWithAtenNotONNXOp(),
+            (x, y),
+            f,
+            do_constant_folding=False,
+            operator_export_type=torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK,
+            # support for linalg.qr was added in later op set versions.
+            opset_version=9,
+        )
+        onnx_model = onnx.load(io.BytesIO(f.getvalue()))
+        self.assertAtenOp(onnx_model, "linalg_qr")
+
+    @common_utils.skipIfCaffe2
+    def test_onnx_aten(self):
+        class ModelWithAtenFmod(torch.nn.Module):
+            def forward(self, x, y):
+                return torch.fmod(x, y)
+
+        x = torch.randn(3, 4, dtype=torch.float32)
+        y = torch.randn(3, 4, dtype=torch.float32)
+        f = io.BytesIO()
+        torch.onnx.export(
+            ModelWithAtenFmod(),
+            (x, y),
+            f,
+            do_constant_folding=False,
+            operator_export_type=torch.onnx.OperatorExportTypes.ONNX_ATEN,
+        )
+        onnx_model = onnx.load(io.BytesIO(f.getvalue()))
+        self.assertAtenOp(onnx_model, "fmod", "Tensor")
+
 
 if __name__ == "__main__":
     common_utils.run_tests()