[MPS] Fix bug when value is of complex (pytorch#111937)

When the value of `fill` is of complex, this line `value.toDouble() == 0.0` will error out saying that converting complex to double will cause overflow. So we should firstly handle the complex value and then enter this condition.

Pull Request resolved: pytorch#111937
Approved by: https://github.com/malfet
ghstack dependencies: pytorch#111885

aten/src/ATen/native/mps/operations/ConstantOps.mm

@@ -86,10 +86,6 @@ static bool fill_mps_tensor_(Tensor& self, uint8_t value) {
 }
 
 Tensor& fill_scalar_mps(Tensor& self, const Scalar& value) {
-  // check if it's possible to use fillBuffer() to fill the Tensor's storage
-  if (value.toDouble() == 0.0 && fill_mps_tensor_(self, 0) == true)
-    return self;
-
   if (isComplexType(self.scalar_type())) {
     auto self_as_real = at::view_as_real(self);
     auto self_as_real_real = self_as_real.select(self.dim(), 0);
@@ -104,6 +100,10 @@ static bool fill_mps_tensor_(Tensor& self, uint8_t value) {
     fill_scalar_mps_impl(self_as_real_imag, 0.0f);
     return self;
   }
+  // check if it's possible to use fillBuffer() to fill the Tensor's storage
+  if (value.toDouble() == 0.0 && fill_mps_tensor_(self, 0) == true)
+    return self;
+
   return fill_scalar_mps_impl(self, value);
 }
 

test/test_mps.py

@@ -1407,19 +1407,18 @@ def testNumbersGPU(self):
 
     def test_fill(self):
 
-        def helper(val, shape):
-            tensor = torch.zeros(shape, device='mps')
+        def helper(val, shape, dtype):
+            tensor = torch.zeros(shape, device='mps', dtype=dtype)
             tensor_mps = tensor.fill_(val)
-            tensor_mps = torch.tanh(tensor_mps)
 
-            tensor_0 = torch.zeros(shape, device='cpu')
+            tensor_0 = torch.zeros(shape, device='cpu', dtype=dtype)
             tensor_cpu = tensor_0.fill_(val)
-            tensor_cpu = torch.tanh(tensor_cpu)
 
             self.assertEqual(tensor_mps, tensor_cpu)
 
-        helper(0, [1024])
-        helper(0.2, [2, 3])
+        helper(0, [1024], torch.float32)
+        helper(0.2, [2, 3], torch.float32)
+        helper(0.2 + 0.5j, [2, 3], torch.complex64)
 
     def test_fill_storage_offset(self):
         shape = [2, 10]