Move scalar_to_tensor_default_dtype out of ScalarOps.h because it's only useful for torch.where. (#50111)

gchanan · facebook-github-bot · commit 0f412aa29392 · 2021-01-11T09:36:29.000-08:00
Summary: Pull Request resolved: #50111 Test Plan: Imported from OSS Reviewed By: mruberry Differential Revision: D25789638 Pulled By: gchanan fbshipit-source-id: 4254e11e08606b64e393433ef2c169889ff2ac07
diff --git a/aten/src/ATen/ScalarOps.h b/aten/src/ATen/ScalarOps.h
@@ -45,30 +45,4 @@ inline at::Tensor scalar_to_tensor(Scalar s, const Device device = at::kCPU) {
   }
 }
 
-// The above function is useful for type promotion
-// in Binary Ops where one argument is `Tensor` and other is `Scalar`.
-// In the above function, we generate wrapped tensor to type with highest
-// range and precision based on scalar's type (to support type promotion).
-// Eg. Floating Point Types -> Double
-//     Complex Types -> Complex Double
-//
-// However for `Scalar-Scalar` Binary Op,we default the type of wrapped tensor
-// to the default type corresponding to scalar's type.
-inline at::Tensor scalar_to_tensor_default_dtype(
-    Scalar s,
-    const Device device = at::kCPU) {
-  if (s.isFloatingPoint()) {
-    return at::scalar_tensor(
-        s, at::device(device).dtype(at::get_default_dtype()));
-  } else if (s.isBoolean()) {
-    return at::scalar_tensor(s, at::device(device).dtype(at::kBool));
-  } else if (s.isComplex()) {
-    return at::scalar_tensor(
-        s, at::device(device).dtype(at::get_default_complex_dtype()));
-  } else {
-    AT_ASSERT(s.isIntegral(false));
-    return at::scalar_tensor(s, at::device(device).dtype(at::kLong));
-  }
-}
-
 }
diff --git a/aten/src/ATen/native/TensorCompare.cpp b/aten/src/ATen/native/TensorCompare.cpp
@@ -195,7 +195,35 @@ bool is_nonzero(const Tensor& self) {
 
 namespace {
 
-static Tensor wrapped_scalar_tensor(
+// DO NOT USE THIS -- it's just an implementation detail of wrapped_scalar tensor below.
+at::Tensor scalar_to_tensor_default_dtype(
+    Scalar s,
+    const Device device = at::kCPU) {
+  if (s.isFloatingPoint()) {
+    return at::scalar_tensor(
+        s, at::device(device).dtype(at::get_default_dtype()));
+  } else if (s.isBoolean()) {
+    return at::scalar_tensor(s, at::device(device).dtype(at::kBool));
+  } else if (s.isComplex()) {
+    return at::scalar_tensor(
+        s, at::device(device).dtype(at::get_default_complex_dtype()));
+  } else {
+    TORCH_INTERNAL_ASSERT(s.isIntegral(false));
+    return at::scalar_tensor(s, at::device(device).dtype(at::kLong));
+  }
+}
+
+// TLDR: Don't call with `use_default_dtype` true -- this is only necessary to support the partial
+// type-promotion that torch.where supports.  Once torch.where fully supports type promotion, we
+// won't need this function.
+//
+// Longer explanation:
+// `use_default_dtype` is a bit of a hack because torch.where doesn't support type promotion, but
+// does support `torch.where(tensor, scalar1, scalar2)` with default scalar types.  The trickiness is we
+// usually convert double scalars to doubles, and `set_wrapped_number` defines type promotion priority
+// as being below tensor types rather than as the default dtype (perhaps we should?).  This wouldn't matter
+// if we just supported type normal type promotion on torch.where, however.
+Tensor wrapped_scalar_tensor(
     Scalar scalar,
     Device device,
     bool use_default_dtype = false) {
