allow passing Python built-in types as dtypes (#21215)

Brennan Vincent · facebook-github-bot · commit 0a3fb45d3d2c · 2019-06-06T13:17:23.000-07:00
Summary: Another simple bit of syntax that NumPy supports and we don't. Support int, float, and bool. ```python >>> torch.randn((2,3), dtype=float) tensor([[-0.1752, -0.3240, -0.6148], [ 0.1861, 1.6472, 0.1687]], dtype=torch.float64) ``` A bit confusingly, Python's "float" actually means double, but nothing we can do about that. Pull Request resolved: #21215 Differential Revision: D15697012 Pulled By: umanwizard fbshipit-source-id: 9a38d960a610b8e67023486b0c9265edd3c22246
diff --git a/test/test_torch.py b/test/test_torch.py
@@ -11789,6 +11789,19 @@ def test_T(self):
         scalar = torch.tensor(5)
         self.assertEqual(scalar, scalar.T)
 
+    def test_python_types(self):
+        a1 = torch.randn((1, 2), dtype=torch.float64)
+        a2 = torch.randn((1, 2), dtype=float)
+        self.assertEqual(a1.dtype, a2.dtype)
+
+        b1 = torch.arange(10, 20, dtype=torch.int64)
+        b2 = torch.arange(10, 20, dtype=int)
+        self.assertEqual(b1.dtype, b2.dtype)
+
+        c1 = torch.tensor([True, False], dtype=torch.bool)
+        c2 = torch.tensor([True, False], dtype=bool)
+        self.assertEqual(c1.dtype, c2.dtype)
+
 # Functions to test negative dimension wrapping
 METHOD = 1
 INPLACE_METHOD = 2
diff --git a/torch/csrc/Dtype.h b/torch/csrc/Dtype.h
@@ -18,6 +18,15 @@ inline bool THPDtype_Check(PyObject *obj) {
   return Py_TYPE(obj) == &THPDtypeType;
 }
 
+inline bool THPPythonScalarType_Check(PyObject *obj) {
+  return obj == (PyObject*)(&PyFloat_Type) ||
+    obj == (PyObject*)(&PyBool_Type) ||
+#if PY_MAJOR_VERSION == 2
+    obj == (PyObject*)(&PyInt_Type) ||
+#endif
+    obj == (PyObject*)(&PyLong_Type);
+}
+
 PyObject * THPDtype_New(at::ScalarType scalar_type, const std::string& name);
 
 void THPDtype_init(PyObject *module);
diff --git a/torch/csrc/utils/python_arg_parser.cpp b/torch/csrc/utils/python_arg_parser.cpp
@@ -169,7 +169,7 @@ bool FunctionParameter::check(PyObject* obj) {
     case ParameterType::BOOL: return PyBool_Check(obj);
     case ParameterType::STORAGE: return isStorage(obj);
     case ParameterType::PYOBJECT: return true;
-    case ParameterType::SCALARTYPE: return THPDtype_Check(obj);
+    case ParameterType::SCALARTYPE: return THPDtype_Check(obj) || THPPythonScalarType_Check(obj);
     case ParameterType::LAYOUT: return THPLayout_Check(obj);
     case ParameterType::MEMORY_FORMAT: return THPMemoryFormat_Check(obj);
     case ParameterType::DEVICE:
diff --git a/torch/csrc/utils/python_arg_parser.h b/torch/csrc/utils/python_arg_parser.h
@@ -337,7 +337,21 @@ inline at::ScalarType PythonArgs::scalartype(int i) {
     return (scalartype == at::ScalarType::Undefined) ?
             torch::tensors::get_default_scalar_type() : scalartype;
   }
-  return reinterpret_cast<THPDtype*>(args[i])->scalar_type;
+  PyObject *obj = args[i];
+  if (obj == (PyObject*)&PyFloat_Type) {
+    return at::ScalarType::Double;
+  }
+  if (obj == (PyObject*)&PyBool_Type) {
+    return at::ScalarType::Bool;
+  }
+  if (obj == (PyObject*)&PyLong_Type
+#if PY_MAJOR_VERSION == 2
+      || obj == (PyObject*)&PyInt_Type
+#endif
+  ) {
+    return at::ScalarType::Long;
+  }
+  return reinterpret_cast<THPDtype*>(obj)->scalar_type;
 }
 
 inline c10::optional<at::ScalarType> PythonArgs::scalartypeOptional(int i) {
