Revert "[ONNX] Fix scalar type promotion between fp16 tensor and fp32 scalar (#113404)"

pytorchmergebot · pytorchmergebot · commit 0fd856ca22e9 · 2023-11-13T14:56:35.000Z
This reverts commit 39ca5a3. Reverted #113404 on behalf of https://github.com/jeanschmidt due to sorry it is breaking CI jobs on main ([comment](#113404 (comment)))
diff --git a/test/onnx/test_pytorch_onnx_onnxruntime.py b/test/onnx/test_pytorch_onnx_onnxruntime.py
@@ -7138,18 +7138,6 @@ def forward(self, fp16, fp32):
         fp32 = Tensor([1.5])
         self.run_test(CatModel(), (fp16, fp32))
 
-    @skipIfUnsupportedMinOpsetVersion(9)
-    def test_scalar_type_does_not_trigger_upcast_type_promotion(self):
-        class DoNotUpcastModel(torch.nn.Module):
-            def forward(self, x):
-                scale = x.size()[-1] ** -0.5
-                # 'scale' is exported as onnx float32 rank 0 tensor.
-                # The following 'Mul' should NOT be promoted to float32.
-                return x * scale
-
-        x = torch.ones(2, 3, dtype=torch.float16)
-        self.run_test(DoNotUpcastModel(), x)
-
     @skipIfUnsupportedMinOpsetVersion(9)
     def test_full_like(self):
         class FullLikeModel(torch.nn.Module):
diff --git a/torch/csrc/jit/passes/onnx/scalar_type_analysis.cpp b/torch/csrc/jit/passes/onnx/scalar_type_analysis.cpp
@@ -161,8 +161,7 @@ static c10::optional<c10::ScalarType> InferExpectedScalarType(const Node* n) {
             at::typeMetaToScalarType(at::get_default_dtype());
         switch (scalar_type) {
           case at::kDouble:
-          case at::kFloat:
-            // floating-point numbers wrapped as float32/float64 tensors are
+            // floating-point numbers wrapped as double tensors are
             // considered to have default type, instead of double.
             typesFromScalars.emplace_back(default_scalar_type);
             break;
@@ -203,33 +202,23 @@ static c10::optional<c10::ScalarType> InferExpectedScalarType(const Node* n) {
           } else {
             typesFromTensors.emplace_back(scalar_type);
           }
-        } else if (auto scalar_type = get_scalar_type(input)) {
-          auto tensor_type = input->type()->castRaw<TensorType>();
-          // get_scalar_type returns non-null value already guarantees
-          // that the input has a valid tensor_type.
-          TORCH_INTERNAL_ASSERT(nullptr != tensor_type);
-          // ONNX model track shape related computes that were done in pytorch
-          // by python numbers as tensor computes. This is the only way for ONNX
-          // to track them properly since ONNX only has tensor type, otherwise
-          // the computation result will be tracked statically as constant, and
-          // the model won't work for another input that differs in shape.
-
-          // Now for type promotion logic, scalars should be treated differently
-          // with tensors. More info regarding type promotion logic commented at
-          // `emplace_type_from_scalar`. Here we filter out rank 0 tensors and
-          // run it with `emplace_type_from_scalar` to determine if they are
-          // considered scalars for type promotion.
-
-          // NOTE that this might introduce regression that a REAL 0-rank tensor
-          // is now being recognized as scalar. The downside is the model will
-          // drop in accuracy for these cases as certain computations will
-          // happen in lower precision data types.
-          auto rank = tensor_type->dim();
-          if (rank && rank.value() == 0) {
-            emplace_type_from_scalar(scalar_type.value());
-          } else {
-            typesFromTensors.emplace_back(scalar_type.value());
+        } else if (nkind == prim::Param) {
+          // ONNX doesn't support scalar as graph input. When
+          // seeing a scalar input, we convert its expected type to tensor.
+          if (auto scalar_type = get_scalar_type(input)) {
+            auto tensor_type = input->type()->castRaw<TensorType>();
+            // get_scalar_type returns non-null value already guarantees
+            // that the input has a valid tensor_type.
+            TORCH_INTERNAL_ASSERT(nullptr != tensor_type);
+            auto rank = tensor_type->dim();
+            if (rank && rank.value() == 0) {
+              emplace_type_from_scalar(scalar_type.value());
+            } else {
+              typesFromTensors.emplace_back(scalar_type.value());
+            }
           }
+        } else if (auto scalar_type = get_scalar_type(input)) {
+          typesFromTensors.emplace_back(*scalar_type);
         }
       });
 
