NumPy Scalar to PyTorch Scalar (#9225)

vishwakftw · facebook-github-bot · commit ea3c36b822a7 · 2018-07-30T14:39:40.000-07:00
Summary: Fixes #4985 . Pull Request resolved: #9225 Differential Revision: D8769317 Pulled By: ezyang fbshipit-source-id: eeaeaf0749c9dc9e372634da68b4bd23e6e3ad28
diff --git a/test/test_torch.py b/test/test_torch.py
@@ -7903,6 +7903,20 @@ def test_ctor_with_numpy_array(self):
                 for i in range(len(array)):
                     self.assertEqual(tensor[i], array[i])
 
+    @unittest.skipIf(not TEST_NUMPY, "Numpy not found")
+    def test_ctor_with_numpy_scalar_ctor(self):
+        dtypes = [
+            np.double,
+            np.float,
+            np.float16,
+            np.int64,
+            np.int32,
+            np.int16,
+            np.uint8
+        ]
+        for dtype in dtypes:
+            self.assertEqual(dtype(42), torch.tensor(dtype(42)).item())
+
     @unittest.skipIf(not TEST_NUMPY, "Numpy not found")
     def test_numpy_index(self):
         i = np.int32([0, 1, 2])
diff --git a/torch/csrc/utils/tensor_new.cpp b/torch/csrc/utils/tensor_new.cpp
@@ -139,8 +139,10 @@ ScalarType infer_scalar_type(PyObject *obj) {
   }
 #ifdef USE_NUMPY
   if (PyArray_Check(obj)) {
-    auto array = (PyArrayObject*)obj;
-    return numpy_dtype_to_aten(PyArray_TYPE(array));
+    return numpy_dtype_to_aten(PyArray_TYPE((PyArrayObject*)obj));
+  }
+  if (PyArray_CheckScalar(obj)) {
+    return numpy_dtype_to_aten(PyArray_TYPE((PyArrayObject*)(PyArray_FromScalar(obj, NULL))));
   }
 #endif
   if (PySequence_Check(obj)) {

Original file line number	Diff line number	Diff line change
`@@ -139,8 +139,10 @@ ScalarType infer_scalar_type(PyObject *obj) {`
`139`	`139`	`}`
`140`	`140`	`#ifdef USE_NUMPY`
`141`	`141`	`if (PyArray_Check(obj)) {`
`142`		`- auto array = (PyArrayObject*)obj;`
`143`		`- return numpy_dtype_to_aten(PyArray_TYPE(array));`
	`142`	`+ return numpy_dtype_to_aten(PyArray_TYPE((PyArrayObject*)obj));`
	`143`	`+ }`
	`144`	`+ if (PyArray_CheckScalar(obj)) {`
	`145`	`+ return numpy_dtype_to_aten(PyArray_TYPE((PyArrayObject*)(PyArray_FromScalar(obj, NULL))));`
`144`	`146`	`}`
`145`	`147`	`#endif`
`146`	`148`	`if (PySequence_Check(obj)) {`