better support set_default_dtype(torch.float16), update doc (pytorch#121730)

1. Fixes pytorch#121300
2. Previously, calling `torch.tensor([2j])` after `torch.set_default_dtype(torch.float16)` will cause a runtime error. This PR also fixes it and enables test.

Pull Request resolved: pytorch#121730
Approved by: https://github.com/peterbell10

Author: LamForest

test/test_complex.py

@@ -18,12 +18,17 @@ def test_to_list(self, device, dtype):
         # there's no garbage value in the resultant list
         self.assertEqual(torch.zeros((2, 2), device=device, dtype=dtype).tolist(), [[0j, 0j], [0j, 0j]])
 
-    @dtypes(torch.float32, torch.float64)
+    @dtypes(torch.float32, torch.float64, torch.float16)
     def test_dtype_inference(self, device, dtype):
         # issue: https://github.com/pytorch/pytorch/issues/36834
         with set_default_dtype(dtype):
             x = torch.tensor([3., 3. + 5.j], device=device)
-        self.assertEqual(x.dtype, torch.cdouble if dtype == torch.float64 else torch.cfloat)
+        if dtype == torch.float16:
+            self.assertEqual(x.dtype, torch.chalf)
+        elif dtype == torch.float32:
+            self.assertEqual(x.dtype, torch.cfloat)
+        else:
+            self.assertEqual(x.dtype, torch.cdouble)
 
     @dtypes(*complex_types())
     def test_conj_copy(self, device, dtype):

torch/__init__.py

@@ -750,17 +750,17 @@ def set_default_tensor_type(t):
 def set_default_dtype(d):
     r"""
 
-    Sets the default floating point dtype to :attr:`d`. Supports torch.float32
-    and torch.float64 as inputs. Other dtypes may be accepted without complaint
-    but are not supported and are unlikely to work as expected.
+    Sets the default floating point dtype to :attr:`d`. Supports floating point dtype
+    as inputs. Other dtypes will cause torch to raise an exception.
 
     When PyTorch is initialized its default floating point dtype is torch.float32,
     and the intent of set_default_dtype(torch.float64) is to facilitate NumPy-like
     type inference. The default floating point dtype is used to:
 
-    1. Implicitly determine the default complex dtype. When the default floating point
-       type is float32 the default complex dtype is complex64, and when the default
-       floating point type is float64 the default complex type is complex128.
+    1. Implicitly determine the default complex dtype. When the default floating type is float16,
+       the default complex dtype is complex32. For float32, the default complex dtype is complex64.
+       For float64, it is complex128. For bfloat16, an exception will be raised because
+       there is no corresponding complex type for bfloat16.
     2. Infer the dtype for tensors constructed using Python floats or complex Python
        numbers. See examples below.
     3. Determine the result of type promotion between bool and integer tensors and
@@ -782,14 +782,21 @@ def set_default_dtype(d):
         torch.complex64
 
         >>> torch.set_default_dtype(torch.float64)
-
         >>> # Python floats are now interpreted as float64
         >>> torch.tensor([1.2, 3]).dtype    # a new floating point tensor
         torch.float64
         >>> # Complex Python numbers are now interpreted as complex128
         >>> torch.tensor([1.2, 3j]).dtype   # a new complex tensor
         torch.complex128
 
+        >>> torch.set_default_dtype(torch.float16)
+        >>> # Python floats are now interpreted as float16
+        >>> torch.tensor([1.2, 3]).dtype    # a new floating point tensor
+        torch.float16
+        >>> # Complex Python numbers are now interpreted as complex128
+        >>> torch.tensor([1.2, 3j]).dtype   # a new complex tensor
+        torch.complex32
+
     """
     _C._set_default_dtype(d)
 

torch/_refs/__init__.py

@@ -6296,6 +6296,8 @@ def _infer_scalar_type(obj):
             return torch.cfloat
         elif default_dtype is torch.double:
             return torch.cdouble
+        elif default_dtype is torch.half:
+            return torch.chalf
         else:
             raise RuntimeError("invalid default scalar type for complex")
     if isinstance(obj, torch.Tensor):

torch/csrc/utils/tensor_new.cpp

@@ -154,6 +154,8 @@ ScalarType infer_scalar_type(PyObject* obj) {
         return ScalarType::ComplexFloat;
       case ScalarType::Double:
         return ScalarType::ComplexDouble;
+      case ScalarType::Half:
+        return ScalarType::ComplexHalf;
       default:
         TORCH_CHECK(false, "invalid default scalar type for complex");
     }