[primtorch] aten.normal decomp has wrong return type due to elementwise_type_promotion_wrapper (#112467)

jon-chuang · pytorchmergebot · commit 53acdb66f7ed · 2023-10-31T20:57:09.000Z
Fixes #112449 elementwise_type_promotion_wrapper will promote `aten.normal` to the dtypes of `mean`, `std` args. But this is incorrect if we provide the dtype param. Hence, we allow overriding the result_dtype if a specified dtype arg is available. This problem is unique to `aten.normal`, all other ops decorated do not have a dtype param. Pull Request resolved: #112467 Approved by: https://github.com/lezcano
diff --git a/test/inductor/test_cpu_repro.py b/test/inductor/test_cpu_repro.py
@@ -2534,6 +2534,22 @@ def forward(self, x):
             )
             assert metrics.generated_kernel_count == 0
 
+    @config.patch(implicit_fallbacks=True)
+    def test_aten_normal_dtype(self):
+        for dtype in [torch.float64, torch.float16, None]:
+
+            def fn():
+                return torch.normal(2, 3, (10, 10), dtype=dtype, device="cpu")
+
+            self.assertEqual(
+                torch.compile(fn, backend="aot_eager_decomp_partition")().dtype,
+                dtype if dtype else torch.float32,
+            )
+            self.assertEqual(
+                torch.compile(fn, backend="inductor")().dtype,
+                dtype if dtype else torch.float32,
+            )
+
     def test_group_norm_vec(self):
         class M(torch.nn.Module):
             def __init__(self):
diff --git a/torch/_prims_common/wrappers.py b/torch/_prims_common/wrappers.py
@@ -93,6 +93,9 @@ class elementwise_type_promotion_wrapper:
     type_promotion_kind must be one of the kinds specified by ELEMENTWISE_TYPE_PROMOTION_KIND.
     See its documentation for details.
 
+    The return_dtype will be coerced to the wrapped function's dtype arg if it is available and
+    not None.
+
     Other type promotion behavior, like validating the Python type of scalar arguments, must
     be handled separately.
     """
@@ -133,6 +136,12 @@ def _fn(*args, **kwargs):
 
             result = fn(**bound.arguments)
 
+            # Override the return_dtype if a dtype arg is present and not None
+            if "dtype" in bound.arguments:
+                maybe_dtype = bound.arguments["dtype"]
+                if maybe_dtype:  # dtype cannot be None
+                    result_dtype = maybe_dtype
+
             if isinstance(result, TensorLike):
                 return _maybe_convert_to_dtype(result, result_dtype)
             if isinstance(result, Sequence):
