Ignore shape inference exception from Caffe2 ATen fallback (pytorch#90408)

Fixes pytorch#87318

Pull Request resolved: pytorch#90408
Approved by: https://github.com/BowenBao

Author: Thiago Crepaldi

test/onnx/test_export_modes.py

@@ -5,7 +5,6 @@
 import shutil
 import sys
 import tempfile
-import unittest
 
 import torch
 import torch.nn as nn
@@ -87,10 +86,6 @@ def foo(a):
         x = torch.ones(3)
         torch.onnx.export(foo, (x,), f)
 
-    # TODO(87318): Can't pass even with Caffe2
-    @unittest.skip(
-        "RuntimeError: ScalarType UNKNOWN_SCALAR is an unexpected tensor scalar type"
-    )
     @common_utils.skipIfNoCaffe2
     @common_utils.skipIfNoLapack
     def test_caffe2_aten_fallback(self):

torch/onnx/utils.py

@@ -686,9 +686,19 @@ def _optimize_graph(
     graph = _C._jit_pass_canonicalize(graph)
     _C._jit_pass_lint(graph)
     if GLOBALS.onnx_shape_inference:
-        _C._jit_pass_onnx_graph_shape_type_inference(
-            graph, params_dict, GLOBALS.export_onnx_opset_version
-        )
+        try:
+            _C._jit_pass_onnx_graph_shape_type_inference(
+                graph, params_dict, GLOBALS.export_onnx_opset_version
+            )
+        except RuntimeError as exc:
+            if (
+                _C_onnx._CAFFE2_ATEN_FALLBACK
+                and exc.args[0]
+                == "ScalarType UNKNOWN_SCALAR is an unexpected tensor scalar type!"
+            ):
+                # Caffe2 builds can have UNKNOWN_SCALAR for some tensors
+                pass
+
     return graph
 
 
@@ -1183,9 +1193,18 @@ def _model_to_graph(
         _C._jit_pass_dce_allow_deleting_nodes_with_side_effects(graph)
 
     if GLOBALS.onnx_shape_inference:
-        _C._jit_pass_onnx_graph_shape_type_inference(
-            graph, params_dict, GLOBALS.export_onnx_opset_version
-        )
+        try:
+            _C._jit_pass_onnx_graph_shape_type_inference(
+                graph, params_dict, GLOBALS.export_onnx_opset_version
+            )
+        except RuntimeError as exc:
+            if (
+                _C_onnx._CAFFE2_ATEN_FALLBACK
+                and exc.args[0]
+                == "ScalarType UNKNOWN_SCALAR is an unexpected tensor scalar type!"
+            ):
+                # Caffe2 builds can have UNKNOWN_SCALAR for some tensors
+                pass
 
     params_dict = _C._jit_pass_onnx_eliminate_unused_items(graph, params_dict)