[fix] type promotion atan2 (pytorch#43466)

Summary:
Fixes pytorch#43360

Pull Request resolved: pytorch#43466

Reviewed By: malfet

Differential Revision: D23834928

Pulled By: mruberry

fbshipit-source-id: 2e7e0b4fcf1a846efc171c275d65a6daffd3c631

Author: kshitij12345

aten/src/ATen/native/BinaryOps.cpp

@@ -390,14 +390,16 @@ Tensor rsub(const Tensor& self, const Tensor& other, Scalar alpha) {
 }
 
 Tensor& atan2_out(Tensor& result, const Tensor& self, const Tensor& other) {
-  auto iter = TensorIterator::binary_op(result, self, other);
+  auto iter = TensorIterator::binary_float_op(result, self, other);
   atan2_stub(iter.device_type(), iter);
   return result;
 }
 
 Tensor atan2(const Tensor& self, const Tensor& other) {
-  Tensor result = at::empty({0}, self.options());
-  return native::atan2_out(result, self, other);
+  Tensor result;
+  auto iter = TensorIterator::binary_float_op(result, self, other);
+  atan2_stub(iter.device_type(), iter);
+  return iter.output();
 }
 
 Tensor& atan2_(Tensor& self, const Tensor& other) {

test/test_torch.py

@@ -19672,10 +19672,20 @@ def test_movedim_view(self, device):
     torch.int8, torch.short, torch.int, torch.long
 ]
 
+_integer_types = [
+    torch.uint8, torch.int8, torch.int16,
+    torch.int32, torch.int64
+]
+
 _cpu_types: List[torch.dtype] = []
 
 _unsigned_types = [torch.uint8]
 
+# Binary Float Ops
+# Operators which use TensorIterator::binary_float_op
+# These Ops promote integer inputs to Float.
+binary_float_ops_inplace = ['atan2_', 'div_']
+
 # Helper values and functions for producing tensors and scalars to use in tensor op tests.
 # Tensor dimension sizes (Small, Medium, Large, Giant)
 _S = 5
@@ -19896,7 +19906,7 @@ def inner(self, device, dtype):
         lambda t, d: [_number(0.5, 3, t), _number(0.4, 2, t), _medium_1d(t, d), _medium_1d(t, d)],
         1e-2, 1e-1, 1e-4, _float_types2, _cpu_types, True,
         [_wrap_maybe_warns("This overload of addr_? is deprecated")]),
-    ('atan2', '', _medium_2d, lambda t, d: [_medium_2d(t, d)], 1e-2, 1e-5, 1e-5, _float_types),
+    ('atan2', '', _medium_2d, lambda t, d: [_medium_2d(t, d)], 1e-2, 1e-5, 1e-5, _types, _types_no_half),
     ('angle', '', _small_3d, lambda t, d: [], 0, 0, 0, _types_no_half, [torch.bfloat16], False),
     ('fmod', 'value', _small_3d, lambda t, d: [3], 1e-3),
     ('fmod', 'tensor', _small_3d, lambda t, d: [_small_3d(t, d, has_zeros=False)], 1e-3),
@@ -20188,6 +20198,15 @@ def fn(self, device, dtype) -> None:
                            (isinstance(arg, torch.Tensor) and arg.dtype == torch.float) else arg
                            for arg in device_args]
 
+        # Special case for binary float ops (binary ops that promote int to float)
+        if op_str in binary_float_ops_inplace and \
+                'inplace' in subtest_str and dtype in _integer_types:
+            with self.assertRaisesRegex(RuntimeError, "result type Float can't be cast to "):
+                cpu_result = getattr(cpu_tensor, op_str)(*cpu_args)
+            with self.assertRaisesRegex(RuntimeError, "result type Float can't be cast to "):
+                device_result = getattr(device_tensor, op_str)(*device_args)
+            return  # Nothing more to check
+
         # Runs the tensor op on CPU and device
         cpu_result = getattr(cpu_tensor, op_str)(*cpu_args)
         device_result = getattr(device_tensor, op_str)(*device_args)

test/test_type_promotion.py

@@ -7,7 +7,7 @@
 from torch.testing._internal.common_utils import (TestCase, run_tests, load_tests,
                                                   TEST_NUMPY, torch_to_numpy_dtype_dict)
 from torch.testing._internal.common_device_type import (instantiate_device_type_tests, onlyOnCPUAndCUDA,
-                                                        dtypes, onlyCPU)
+                                                        dtypes, dtypesIfCUDA, onlyCPU)
 
 if TEST_NUMPY:
     import numpy as np
@@ -958,6 +958,37 @@ def test_computation_ignores_out(self, device):
         self.assertEqual(result, a - b, exact_dtype=False)
         self.assertNotEqual(result, a.double() - b, exact_dtype=False)
 
+    @dtypesIfCUDA(*itertools.product(torch.testing.get_all_dtypes(include_bfloat16=False, include_complex=False),
+                                     torch.testing.get_all_dtypes(include_bfloat16=False, include_complex=False)))
+    @dtypes(*itertools.product(torch.testing.get_all_dtypes(include_half=False, include_bfloat16=False,
+                                                            include_complex=False),
+                               torch.testing.get_all_dtypes(include_half=False, include_bfloat16=False,
+                                                            include_complex=False)))
+    def test_atan2_type_promotion(self, device, dtypes):
+        dtype1, dtype2 = dtypes
+        default_float = torch.get_default_dtype()
+
+        def is_int(dtype):
+            return dtype in torch.testing.get_all_int_dtypes() + [torch.bool]
+
+        def is_float(dtype):
+            return dtype in torch.testing.get_all_fp_dtypes(include_half=True, include_bfloat16=False)
+
+        def get_binary_float_result_type(x, y):
+            dtype1 = x.dtype
+            dtype2 = y.dtype
+            if is_float(dtype1) and is_float(dtype2):
+                return torch.result_type(x, y)
+            elif is_float(dtype1) and is_int(dtype2):
+                return dtype1
+            elif is_int(dtype1) and is_float(dtype2):
+                return dtype2
+            elif is_int(dtype1) and is_int(dtype2):
+                return default_float
+
+        x = torch.tensor(1, dtype=dtype1, device=device)
+        y = torch.tensor(2, dtype=dtype2, device=device)
+        self.assertEqual(get_binary_float_result_type(x, y), torch.atan2(x, y).dtype)
 
 instantiate_device_type_tests(TestTypePromotion, globals())