Handling of scalars in torch.Size (#5676)

zou3519 · web-flow · commit 769f5f7cfe64 · 2018-05-30T17:50:32.000-04:00
* Handling of scalars in torch.Size

torch.Size() constructor uses python_arg_parser

IntList in python_arg_parser can take iter/range

Have IntList take python iterables and ranges.

Address comments: don't use python_arg_parser and instead call __index__ in THPSize_pynew

Address comments

Address comments

* Rebased

* Address nit
diff --git a/test/test_torch.py b/test/test_torch.py
@@ -7100,6 +7100,19 @@ def test_Size(self):
         self.assertIsInstance(x[:-1], torch.Size)
         self.assertIsInstance(x + x, torch.Size)
 
+    def test_Size_scalar(self):
+        three = torch.tensor(3)
+        two = torch.tensor(2)
+        x = torch.Size([0, 1, two, three, 4])
+        for i in range(1, 5):
+            self.assertEqual(x[i], i)
+
+    def test_Size_iter(self):
+        for sizes in [iter([1, 2, 3, 4, 5]), range(1, 6)]:
+            x = torch.Size(sizes)
+            for i in range(0, 5):
+                self.assertEqual(x[i], i + 1)
+
     def test_t_not_2d_error(self):
         self.assertRaises(RuntimeError, lambda: torch.randn(2, 3, 4).t())
         self.assertRaises(RuntimeError, lambda: torch.randn(2, 3, 4).t_())
diff --git a/torch/csrc/Size.cpp b/torch/csrc/Size.cpp
@@ -50,10 +50,25 @@ static PyObject * THPSize_pynew(PyTypeObject *type, PyObject *args, PyObject *kw
   if (self) {
     for (Py_ssize_t i = 0; i < PyTuple_Size(self); ++i) {
       PyObject *item = PyTuple_GET_ITEM(self.get(), i);
-      if (!THPUtils_checkLong(item) && !isTracedVar(item)) {
-        return PyErr_Format(PyExc_TypeError, "torch.Size() takes an iterable of 'int' (item %zd is '%s')",
-            i, Py_TYPE(item)->tp_name);
+      if (isTracedVar(item)) {
+        continue;
       }
+      if (THPUtils_checkLong(item)) {
+        continue;
+      }
+      // item.__index__() works with 0-dim tensors and tensors with one element
+      THPObjectPtr number(PyNumber_Index(item));
+      if (number && THPUtils_checkLong(number.get())) {
+        Py_INCREF(number.get());
+        auto status = PyTuple_SetItem(self, i, number.get());
+        if (status != 0) {
+          throw python_error();
+        }
+        continue;
+      }
+      return PyErr_Format(PyExc_TypeError,
+                          "torch.Size() takes an iterable of 'int' (item %zd is '%s')",
+                          i, Py_TYPE(item)->tp_name);
     }
   }
   return self.release();
