[ONNX][dynamo_export] Turn off opmath type promotion for div (pytorch#119112)

Skip opmath promotion for `_prims_common.ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT` as well.
Fixes pytorch#118941

Pull Request resolved: pytorch#119112
Approved by: https://github.com/thiagocrepaldi

Author: BowenBao

test/onnx/test_fx_op_consistency.py

@@ -1927,6 +1927,7 @@ class TestOnnxModelOutputConsistency(onnx_test_common._TestONNXRuntime):
         "linalg.multi_dot": [3e-2, 1e-3],
         "linalg.vecdot": [1e-2, 2e-2],
         "linspace": [2e-2, 2e-3],
+        "masked.var": [2e-2, 1e-3],
         "matmul": [2e-2, 6e-2],
         "nn.functional.batch_norm": [3e-2, 1e-3],
         "nn.functional.binary_cross_entropy": [3e-2, 1e-3],

test/onnx/test_fx_to_onnx.py

@@ -1,13 +1,12 @@
 # Owner(s): ["module: onnx"]
 from __future__ import annotations
 
-import io
-
 import tempfile
 
 from typing import Mapping, Tuple
 
 import onnx
+import onnx.inliner
 import pytorch_test_common
 import torch
 import transformers  # type: ignore[import]
@@ -621,17 +620,22 @@ def forward(self, x):
                 x,
             )
 
-    def test_aten_linalg_vector_norm_with_reducel2(self):
-        class Net(nn.Module):
-            def forward(self, x):
-                x = F.normalize(x)
-                return x
-
-        f = io.BytesIO()
-        torch.onnx.export(Net(), (torch.randn(1, 2, 2),), f)
-        onnx_model = onnx.load_from_string(f.getvalue())
-        onnx_nodes = [n.op_type for n in onnx_model.graph.node]
-        self.assertTrue("ReduceL2" in onnx_nodes)
+    def test_aten_div_no_opmath_type_promotion(self):
+        class Model(torch.nn.Module):
+            def forward(self, input):
+                return input / 2
+
+        model = Model()
+        input = torch.randn(3, 5, requires_grad=True, dtype=torch.float16)
+
+        model_proto = torch.onnx.dynamo_export(model, input).model_proto
+        model_proto = onnx.inliner.inline_local_functions(model_proto)
+        div_node = next(
+            node for node in model_proto.graph.node if node.op_type == "Div"
+        )
+        # The input of Div node should be the input of the model,
+        # with no Cast node in between.
+        self.assertEqual(div_node.input[0], model_proto.graph.input[0].name)
 
     def test_exported_program_as_input_with_model_signature(self):
         class Model(torch.nn.Module):

test/onnx/test_pytorch_onnx_no_runtime.py

@@ -1314,6 +1314,18 @@ def test_aten_device_with_index(self):
                 decoder_attention_mask=ids["attention_mask"],
             )
 
+    def test_aten_linalg_vector_norm_with_reducel2(self):
+        class Net(torch.nn.Module):
+            def forward(self, x):
+                x = F.normalize(x)
+                return x
+
+        f = io.BytesIO()
+        torch.onnx.export(Net(), (torch.randn(1, 2, 2),), f)
+        onnx_model = onnx.load_from_string(f.getvalue())
+        onnx_nodes = [n.op_type for n in onnx_model.graph.node]
+        self.assertIn("ReduceL2", onnx_nodes)
+
 
 if __name__ == "__main__":
     common_utils.run_tests()

torch/onnx/_internal/fx/passes/type_promotion.py

@@ -184,10 +184,9 @@ def _consolidate_input_dtype(
         since there is no way to differentiate between inserted upcasts and model code
         casts. Hence we consolidate the input dtype to the result dtype to avoid this.
         """
-        if (
-            not self._USE_OPMATH
-            and self.promotion_kind
-            == _prims_common.ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT
+        if not self._USE_OPMATH and self.promotion_kind in (
+            _prims_common.ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT,
+            _prims_common.ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT,
         ):
             return result_dtype
         return computed_dtype