[fix] torch.{lin, log}space(): properly examine passed dtype (#53685)

kshitij12345 · facebook-github-bot · commit c90260905f38 · 2021-06-10T11:59:54.000-07:00
Summary: Fixes #53171 Pull Request resolved: #53685 Reviewed By: jbschlosser Differential Revision: D28331863 Pulled By: anjali411 fbshipit-source-id: e89359b607d058158cfa1c9a82389d9a4a71185b
diff --git a/aten/src/ATen/native/TensorFactories.cpp b/aten/src/ATen/native/TensorFactories.cpp
@@ -517,25 +517,20 @@ namespace {
 TensorOptions linspace_logspace_infer_options(
     const Scalar& start,
     const Scalar& end,
-    const TensorOptions& options) {
-  auto result_options = options;
+    const TensorOptions& options,
+    const char* fn_name) {
   if (start.isComplex() || end.isComplex()) {
-    // Since result_options.has_dtype() returns true (dtype is default type),
-    // even if the user hasn't specified the dtype.
-    // We just check to see if either `start` or `end` is complex,
-    // and if the `result_dtype` is not complex (be it default float type or
-    // user provided), we cast it to default complex dtype with a Warning!.
-    auto result_dtype = c10::typeMetaToScalarType(options.dtype());
-    if (!at::isComplexType(result_dtype)) {
-      TORCH_WARN(
-          "As either `start` or `stop` is complex, return type will be the complex dtype corresponding to default dtype.",
-          "In future, this may throw an error when a non-complex dtype arg is passed as input along ",
-          "with complex valued start or end value.");
-      result_options = result_options.dtype(c10::get_default_complex_dtype());
+    const auto default_complex_dtype = c10::get_default_complex_dtype();
+    if (options.has_dtype()) {
+      auto dtype = c10::typeMetaToScalarType(options.dtype());
+      TORCH_CHECK(at::isComplexType(dtype),
+          fn_name, ": inferred dtype ", default_complex_dtype, " can't be safely cast to passed dtype ", dtype);
+    } else {
+      return options.dtype(default_complex_dtype);
     }
   }
 
-  return result_options;
+  return options.has_dtype() ? options : options.dtype(c10::get_default_dtype());
 }
 } // anonymous namespace
 
@@ -554,7 +549,7 @@ Tensor linspace(
 
   const auto steps_ = steps.value_or(100);
   TORCH_CHECK(steps_ >= 0, "number of steps must be non-negative");
-  auto result_options = linspace_logspace_infer_options(start, end, options);
+  auto result_options = linspace_logspace_infer_options(start, end, options, "torch.linspace()");
   Tensor result = at::empty({steps_}, result_options);
   return at::linspace_out(result, start, end, steps);
 }
@@ -575,7 +570,7 @@ Tensor logspace(
 
   const auto steps_ = steps.value_or(100);
   TORCH_CHECK(steps_ >= 0, "number of steps must be non-negative");
-  auto result_options = linspace_logspace_infer_options(start, end, options);
+  auto result_options = linspace_logspace_infer_options(start, end, options, "torch.logspace()");
   Tensor result = at::empty({steps_}, result_options);
   return at::logspace_out(result, start, end, steps, base);
 }
diff --git a/test/test_tensor_creation_ops.py b/test/test_tensor_creation_ops.py
@@ -2754,6 +2754,15 @@ def test_linspace(self, device, dtype):
         # steps = 0
         self.assertEqual(torch.linspace(0, 1, 0, device=device, dtype=dtype).numel(), 0, atol=0, rtol=0)
 
+        if dtype == torch.float:
+            # passed dtype can't be safely casted to inferred dtype
+            with self.assertRaisesRegex(RuntimeError, r"torch.linspace\(\): inferred dtype"):
+                torch.linspace(0, 1j, 5, device=device, dtype=dtype)
+            with self.assertRaisesRegex(RuntimeError, r"torch.linspace\(\): inferred dtype"):
+                torch.linspace(0j, 1, 5, device=device, dtype=dtype)
+            with self.assertRaisesRegex(RuntimeError, r"torch.linspace\(\): inferred dtype"):
+                torch.linspace(0j, 1j, 5, device=device, dtype=dtype)
+
         # Check linspace for generating the correct output for each dtype.
         start = 0 if dtype == torch.uint8 else -100
         expected_lin = torch.tensor([start + .5 * i for i in range(401)], device=device, dtype=torch.double)
@@ -2797,14 +2806,7 @@ def _test_linspace_logspace_deduction_helper(self, fn, device):
             if isinstance(start, complex) or isinstance(end, complex):
                 dtype = torch.cfloat
 
-            if dtype == torch.cfloat:
-                # TODO(kshitij12345): Fix unnecessary warning
-                # Reference: https://github.com/pytorch/pytorch/issues/53171
-                with self.assertWarnsRegex(UserWarning,
-                                           "As either `start` or `stop` is complex"):
-                    self.assertEqual(fn(start, end, steps=100, device=device).dtype, dtype)
-            else:
-                self.assertEqual(fn(start, end, steps=100, device=device).dtype, dtype)
+            self.assertEqual(fn(start, end, steps=100, device=device).dtype, dtype)
 
     def test_linspace_deduction(self, device):
         # Test deduction from input parameters.
@@ -2896,6 +2898,15 @@ def test_logspace(self, device, dtype):
         self.assertEqual(torch.logspace(0, 1, 1, device=device, dtype=dtype),
                          torch.ones(1, device=device, dtype=dtype), atol=0, rtol=0)
 
+        if dtype == torch.float:
+            # passed dtype can't be safely casted to inferred dtype
+            with self.assertRaisesRegex(RuntimeError, r"torch.logspace\(\): inferred dtype"):
+                torch.logspace(0, 1j, 5, device=device, dtype=dtype)
+            with self.assertRaisesRegex(RuntimeError, r"torch.logspace\(\): inferred dtype"):
+                torch.logspace(0j, 1, 5, device=device, dtype=dtype)
+            with self.assertRaisesRegex(RuntimeError, r"torch.logspace\(\): inferred dtype"):
+                torch.logspace(0j, 1j, 5, device=device, dtype=dtype)
+
         # Check precision - start, stop and base are chosen to avoid overflow
         # steps is chosen so that step size is not subject to rounding error
         # a tolerance is needed for gpu tests due to differences in computation
diff --git a/tools/autograd/gen_python_functions.py b/tools/autograd/gen_python_functions.py
@@ -80,7 +80,7 @@
     '.*_backward', '.*_backward_(out|input|weight|bias)', '.*_forward',
     '.*_forward_out', '_unsafe_view', 'tensor', '_?sparse_coo_tensor.*',
     '_?sparse_csr_tensor.*',
-    '_arange.*', '_range.*', '_linspace.*', '_logspace.*',
+    '_arange.*', '_range.*', 'linspace.*', 'logspace.*',
     '_sparse_add_out', '_sparse_div.*', '_sparse_mul.*', '_sparse_sub.*', '_sparse_dense_add_out',
     'index', 'unique_dim_consecutive',
     '_cumsum.*', '_cumprod.*', '_sum.*', '_prod.*',
diff --git a/tools/autograd/templates/python_torch_functions.cpp b/tools/autograd/templates/python_torch_functions.cpp
@@ -467,6 +467,102 @@ static PyObject * THPVariable_get_device(PyObject* self_, PyObject* args, PyObje
 
 static PyObject * THPVariable_numel(PyObject* self_, PyObject* args, PyObject* kwargs);
 
+// linspace
+static PyObject * THPVariable_linspace(PyObject* self_, PyObject* args, PyObject* kwargs)
+{
+  HANDLE_TH_ERRORS
+  static PythonArgParser parser({
+    "linspace(Scalar start, Scalar end, int64_t? steps=None, *, Tensor out=None, ScalarType dtype=None, Layout layout=torch.strided, Device device=None, bool pin_memory=False, bool requires_grad=False)",
+  }, /*traceable=*/true);
+
+  ParsedArgs<9> parsed_args;
+  auto _r = parser.parse(nullptr, args, kwargs, parsed_args);
+  if(_r.has_torch_function()) {
+    return handle_torch_function(_r, nullptr, args, kwargs, THPVariableFunctionsModule, "torch");
+  }
+  if (_r.isNone(3)) {
+    // aten::linspace(Scalar start, Scalar end, int? steps=None, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor
+
+    // NOTE: r.scalartype(X) gives the default dtype if r.isNone(X)
+    // This leads to problem in the operator argument checks,
+    // when either `start` or `end` is complex and dtype is None
+    const auto options = TensorOptions()
+        .dtype(_r.scalartypeOptional(4))
+        .device(_r.device(6))
+        .layout(_r.layoutOptional(5))
+        .requires_grad(_r.toBool(8))
+        .pinned_memory(_r.toBool(7));
+    torch::utils::maybe_initialize_cuda(options);
+
+    auto dispatch_linspace = [](Scalar start, Scalar end, c10::optional<int64_t> steps, TensorOptions options) -> Tensor {
+      pybind11::gil_scoped_release no_gil;
+      return torch::linspace(start, end, steps, options);
+    };
+    return wrap(dispatch_linspace(_r.scalar(0), _r.scalar(1), _r.toInt64Optional(2), options));
+  } else {
+    // aten::linspace.out(Scalar start, Scalar end, int? steps=None, *, Tensor(a!) out) -> Tensor(a!)
+    check_out_type_matches(_r.tensor(3), _r.scalartype(4),
+                           _r.isNone(4), _r.layoutOptional(5),
+                           _r.device(6), _r.isNone(6));
+
+    auto dispatch_linspace_out = [](Tensor out, Scalar start, Scalar end, c10::optional<int64_t> steps) -> Tensor {
+      pybind11::gil_scoped_release no_gil;
+      return at::linspace_out(out, start, end, steps);
+    };
+    return wrap(dispatch_linspace_out(_r.tensor(3), _r.scalar(0), _r.scalar(1), _r.toInt64Optional(2)).set_requires_grad(_r.toBool(8)));
+  }
+  Py_RETURN_NONE;
+  END_HANDLE_TH_ERRORS
+}
+
+// logspace
+static PyObject * THPVariable_logspace(PyObject* self_, PyObject* args, PyObject* kwargs)
+{
+  HANDLE_TH_ERRORS
+  static PythonArgParser parser({
+    "logspace(Scalar start, Scalar end, int64_t? steps=None, double base=10.0, *, Tensor out=None, ScalarType dtype=None, Layout layout=torch.strided, Device device=None, bool pin_memory=False, bool requires_grad=False)",
+  }, /*traceable=*/true);
+
+  ParsedArgs<10> parsed_args;
+  auto _r = parser.parse(nullptr, args, kwargs, parsed_args);
+  if(_r.has_torch_function()) {
+    return handle_torch_function(_r, nullptr, args, kwargs, THPVariableFunctionsModule, "torch");
+  }
+  if (_r.isNone(4)) {
+    // aten::logspace(Scalar start, Scalar end, int? steps=None, float base=10.0, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor
+
+    // NOTE: r.scalartype(X) gives the default dtype if r.isNone(X)
+    // This leads to problem in the operator argument checks,
+    // when either `start` or `end` is complex and dtype is None
+    const auto options = TensorOptions()
+        .dtype(_r.scalartypeOptional(5))
+        .device(_r.device(7))
+        .layout(_r.layoutOptional(6))
+        .requires_grad(_r.toBool(9))
+        .pinned_memory(_r.toBool(8));
+    torch::utils::maybe_initialize_cuda(options);
+
+    auto dispatch_logspace = [](Scalar start, Scalar end, c10::optional<int64_t> steps, double base, TensorOptions options) -> Tensor {
+      pybind11::gil_scoped_release no_gil;
+      return torch::logspace(start, end, steps, base, options);
+    };
+    return wrap(dispatch_logspace(_r.scalar(0), _r.scalar(1), _r.toInt64Optional(2), _r.toDouble(3), options));
+  } else {
+    // aten::logspace.out(Scalar start, Scalar end, int? steps=None, float base=10.0, *, Tensor(a!) out) -> Tensor(a!)
+    check_out_type_matches(_r.tensor(4), _r.scalartype(5),
+                           _r.isNone(5), _r.layoutOptional(6),
+                           _r.device(7), _r.isNone(7));
+
+    auto dispatch_logspace_out = [](Tensor out, Scalar start, Scalar end, c10::optional<int64_t> steps, double base) -> Tensor {
+      pybind11::gil_scoped_release no_gil;
+      return at::logspace_out(out, start, end, steps, base);
+    };
+    return wrap(dispatch_logspace_out(_r.tensor(4), _r.scalar(0), _r.scalar(1), _r.toInt64Optional(2), _r.toDouble(3)).set_requires_grad(_r.toBool(9)));
+  }
+  Py_RETURN_NONE;
+  END_HANDLE_TH_ERRORS
+}
+
 // generated forward declarations start here
 
 ${py_forwards}
@@ -496,6 +592,8 @@ static PyMethodDef torch_functions[] = {
   {"from_numpy", THPVariable_from_numpy, METH_STATIC | METH_O, NULL},
   {"full", castPyCFunctionWithKeywords(THPVariable_full), METH_VARARGS | METH_KEYWORDS | METH_STATIC, NULL},
   {"hsmm", castPyCFunctionWithKeywords(THPVariable_hspmm), METH_VARARGS | METH_KEYWORDS | METH_STATIC, NULL},
+  {"linspace", castPyCFunctionWithKeywords(THPVariable_linspace), METH_VARARGS | METH_KEYWORDS | METH_STATIC, NULL},
+  {"logspace", castPyCFunctionWithKeywords(THPVariable_logspace), METH_VARARGS | METH_KEYWORDS | METH_STATIC, NULL},
   {"nonzero", castPyCFunctionWithKeywords(THPVariable_nonzero), METH_VARARGS | METH_KEYWORDS | METH_STATIC, NULL},
   {"randint", castPyCFunctionWithKeywords(THPVariable_randint), METH_VARARGS | METH_KEYWORDS | METH_STATIC, NULL},
   {"range", castPyCFunctionWithKeywords(THPVariable_range), METH_VARARGS | METH_KEYWORDS | METH_STATIC, NULL},
diff --git a/tools/pyi/gen_pyi.py b/tools/pyi/gen_pyi.py
@@ -314,6 +314,10 @@ def gen_pyi(native_yaml_path: str, deprecated_yaml_path: str, fm: FileManager) -
                    .format(FACTORY_PARAMS),
                    'def arange(end: Number, *, out: Optional[Tensor]=None, {}) -> Tensor: ...'
                    .format(FACTORY_PARAMS)],
+        'linspace': ['def linspace(start: Number, end: Number, steps: Optional[_int]=None, *,'
+                     ' out: Optional[Tensor]=None, {}) -> Tensor: ...'.format(FACTORY_PARAMS)],
+        'logspace': ['def logspace(start: Number, end: Number, steps: Optional[_int]=None, base: _float=10.0, *,'
+                     ' out: Optional[Tensor]=None, {}) -> Tensor: ...'.format(FACTORY_PARAMS)],
         'randint': ['def randint(low: _int, high: _int, size: _size, *,'
                     ' generator: Optional[Generator]=None, {}) -> Tensor: ...'
                     .format(FACTORY_PARAMS),
