pytorch · gchanan · Apr 27, 2018 · Apr 26, 2018 · Apr 26, 2018 · Apr 27, 2018
diff --git a/test/test_sparse.py b/test/test_sparse.py
@@ -891,6 +891,19 @@ def test_factory_type_inference(self):
         t = torch.sparse_coo_tensor(torch.tensor(([0], [2])), torch.tensor([1]))
         self.assertEqual(torch.int64, t.dtype)
 
+    @cuda_only
+    def test_factory_device_type_inference(self):
+        # both indices/values are CUDA
+        shape = (1, 3)
+        for indices_device in ['cuda', 'cpu']:
+            for values_device in ['cuda', 'cpu']:
+                for sparse_device in ['cuda', 'cpu', None]:
+                    t = torch.sparse_coo_tensor(torch.tensor(([0], [2]), device=indices_device),
+                                                torch.tensor([1.], device=values_device),
+                                                (1, 3), device=sparse_device)
+                    should_be_cuda = sparse_device == 'cuda' or (sparse_device is None and values_device == 'cuda')
+                    self.assertEqual(should_be_cuda, t.is_cuda)
+
     @cpu_only
     def test_factory_copy(self):
         # both correct

diff --git a/test/test_torch.py b/test/test_torch.py
@@ -1745,6 +1745,10 @@ def assertEqual(device_str, fn):
         # NOTE: 'cpu' is the canonical representation of 'cpu:0', but 'cuda:X' is the canonical
         # representation of cuda devices.
         assertEqual('cpu', lambda: torch.ones((2, 3), dtype=torch.float32, device='cpu:0'))
+        assertEqual('cpu', lambda: torch.tensor(torch.ones((2, 3), dtype=torch.float32), device='cpu:0'))
+        if TEST_NUMPY:
+            assertEqual('cpu', lambda: torch.tensor(np.random.randn(2, 3), device='cpu'))
+
         if torch.cuda.is_available():
             assertEqual('cuda:0', lambda: torch.tensor(5).cuda(0))
             assertEqual('cuda:0', lambda: torch.tensor(5).cuda('cuda:0'))
@@ -1754,12 +1758,18 @@ def assertEqual(device_str, fn):
             assertEqual('cuda:0', lambda: torch.tensor(5, dtype=torch.int64, device='cuda:0'))
             assertEqual('cuda:' + str(torch.cuda.current_device()),
                         lambda: torch.tensor(5, dtype=torch.int64, device='cuda'))
+            assertEqual('cuda:0', lambda: torch.tensor(torch.ones((2, 3), dtype=torch.float32), device='cuda:0'))
+            if TEST_NUMPY:
+                assertEqual('cuda:0', lambda: torch.tensor(np.random.randn(2, 3), device='cuda:0'))
 
             if torch.cuda.device_count() > 1:
                 assertEqual('cuda:1', lambda: torch.tensor(5).cuda(1))
                 assertEqual('cuda:1', lambda: torch.tensor(5).cuda('cuda:1'))
                 assertEqual('cuda:1', lambda: torch.tensor(5, dtype=torch.int64, device=1))
                 assertEqual('cuda:1', lambda: torch.tensor(5, dtype=torch.int64, device='cuda:1'))
+                assertEqual('cuda:1', lambda: torch.tensor(torch.ones((2, 3), dtype=torch.float32), device='cuda:1'))
+                if TEST_NUMPY:
+                    assertEqual('cuda:1', lambda: torch.tensor(np.random.randn(2, 3), device='cuda:1'))
 
     def test_to(self):
         a = torch.tensor(5)

diff --git a/torch/csrc/autograd/python_variable_indexing.cpp b/torch/csrc/autograd/python_variable_indexing.cpp
@@ -102,7 +102,7 @@ static Variable applySelect(const Variable& self, int64_t dim, int64_t index) {
 
 static Variable sequenceToVariable(const Type& type, PyObject* seq) {
   auto& idx_type = type.toScalarType(kLong);
-  return torch::utils::legacy_new_from_data(idx_type, -1, seq);
+  return torch::utils::legacy_new_from_data(idx_type, at::nullopt, seq);
 }
 
 static Variable valueToTensor(const Type & type, PyObject* value) {

diff --git a/torch/csrc/utils/python_arg_parser.h b/torch/csrc/utils/python_arg_parser.h
@@ -101,6 +101,7 @@ struct PythonArgs {
   inline Device device(int i);
   inline Device deviceWithDefault(int i, const Device& default_device);
   inline int64_t deviceInt64(int i);
+  inline at::optional<Device> deviceOptional(int i);
   inline std::string string(int i);
   inline PyObject* pyobject(int i);
   inline int64_t toInt64(int i);
@@ -332,6 +333,11 @@ inline int64_t PythonArgs::deviceInt64(int i) {
   return dev.deviceInt64();
 }
 
+inline at::optional<Device> PythonArgs::deviceOptional(int i) {
+  if (!args[i]) return at::nullopt;
+  return device(i);
+}
+
 inline std::string PythonArgs::string(int i) {
   if (!args[i]) return "";
   return THPUtils_unpackString(args[i]);

diff --git a/torch/csrc/utils/tensor_new.cpp b/torch/csrc/utils/tensor_new.cpp
@@ -171,24 +171,32 @@ static void recursive_store(char* data, IntList sizes, IntList strides, int64_t
   }
 }
 
-static Tensor internal_new_from_data(const Type & type, int device, PyObject* data,
+static Tensor internal_new_from_data(const Type & type, at::optional<Device> device_opt, PyObject* data,
                                      bool copy_variables, bool copy_numpy,
                                      bool type_inference) {
+  int64_t device = device_opt.has_value() ? device_opt.value().deviceInt64() : -1;
   if (THPUtils_checkString(data)) {
     throw TypeError("new(): invalid data type '%s'", Py_TYPE(data)->tp_name);
   }
 
   if (THPVariable_Check(data)) {
       auto var = reinterpret_cast<THPVariable*>(data)->cdata;
-      const auto& type_to_use = type_inference ? var.type() : type;
+      auto type_inference_device_type = device_opt.has_value() ? device_opt.value().type
+                                                               : torch::getDeviceType(var.type());
+      // infer the scalar type and device type; it's not expected to infer the layout since these constructors
+      // are defined per-layout-type (e.g. tensor vs sparse_coo_tensor).
+      const auto& type_inference_type = torch::getType(var.type().scalarType(),
+                                                       *torch::getLayout(type.backend()),
+                                                       type_inference_device_type);
+      const auto& type_to_use = type_inference ? type_inference_type : type;
       return copy_variables ? new_with_tensor_copy(type_to_use, var, device) :
                               new_with_type_conversion(type_to_use, var, device);
   }
 
 #ifdef WITH_NUMPY
   if (PyArray_Check(data)) {
     auto tensor = autograd::make_variable(tensor_from_numpy(data), /*requires_grad=*/false);
-    const auto& type_to_use = type_inference ? tensor.type() : type;
+    const auto& type_to_use = type_inference ? type.toScalarType(tensor.type().scalarType()) : type;
     return copy_numpy ? new_with_tensor_copy(type_to_use, tensor, device) :
                         new_with_type_conversion(type_to_use, tensor, device);
   }
@@ -204,15 +212,15 @@ static Tensor internal_new_from_data(const Type & type, int device, PyObject* da
   return new_with_type_conversion(type_to_use, tensor, device);
 }
 
-Tensor legacy_new_from_data(const Type & type, int device, PyObject *data) {
+Tensor legacy_new_from_data(const Type & type, at::optional<Device> device, PyObject *data) {
   return internal_new_from_data(type, device, data, false, false, false);
 }
 
-static Tensor new_from_data_copy(const Type & type, int device, PyObject *data) {
+static Tensor new_from_data_copy(const Type & type, at::optional<Device> device, PyObject *data) {
   return internal_new_from_data(type, device, data, true, true, false);
 }
 
-static Tensor legacy_new_from_sequence(const Type & type, int device, PyObject* data) {
+static Tensor legacy_new_from_sequence(const Type & type, at::optional<Device> device, PyObject* data) {
   if (!PySequence_Check(data)) {
     throw TypeError("new(): data must be a sequence (got %s)", Py_TYPE(data)->tp_name);
   }
@@ -246,7 +254,7 @@ static Tensor legacy_sparse_tensor_ctor(const Type& type, PyObject* args, PyObje
     if (!THPSize_Check(arg) && PyTuple_GET_SIZE(args) >= 1 && arg == PyTuple_GET_ITEM(args, 0)) {
       // new(sequence) binds to this signature but should be treated differently
       // unless the sequences is a torch.Size
-      return legacy_new_from_sequence(type, r.deviceInt64(1), r.pyobject(0));
+      return legacy_new_from_sequence(type, r.deviceOptional(1), r.pyobject(0));
     }
     return new_with_sizes(type, r.deviceInt64(1), r.intlist(0));
   }
@@ -284,11 +292,11 @@ Tensor legacy_tensor_ctor(const Type& type, PyObject* args, PyObject* kwargs) {
     if (!THPSize_Check(arg) && PyTuple_GET_SIZE(args) >= 1 && arg == PyTuple_GET_ITEM(args, 0)) {
       // new(sequence) binds to this signature but should be treated differently
       // unless the sequences is a torch.Size
-      return legacy_new_from_sequence(type, r.deviceInt64(1), r.pyobject(0));
+      return legacy_new_from_sequence(type, r.deviceOptional(1), r.pyobject(0));
     }
     return new_with_sizes(type, r.deviceInt64(1), r.intlist(0));
   } else if (r.idx == 5) {
-    return legacy_new_from_sequence(type, r.deviceInt64(1), r.pyobject(0));
+    return legacy_new_from_sequence(type, r.deviceOptional(1), r.pyobject(0));
   }
   throw std::runtime_error("new(): invalid arguments");
 }
@@ -324,7 +332,7 @@ static Tensor legacy_sparse_tensor_new(const Type& type, PyObject* args, PyObjec
     if (!THPSize_Check(arg) && PyTuple_GET_SIZE(args) >= 1 && arg == PyTuple_GET_ITEM(args, 0)) {
       // new(sequence) binds to this signature but should be treated differently
       // unless the sequences is a torch.Size
-      return legacy_new_from_sequence(type, r.deviceInt64(1), r.pyobject(0));
+      return legacy_new_from_sequence(type, r.deviceOptional(1), r.pyobject(0));
     }
     return new_with_sizes(type, r.deviceInt64(1), r.intlist(0));
   }
@@ -362,11 +370,11 @@ Tensor legacy_tensor_new(const Type& type, PyObject* args, PyObject* kwargs) {
     if (!THPSize_Check(arg) && PyTuple_GET_SIZE(args) >= 1 && arg == PyTuple_GET_ITEM(args, 0)) {
       // new(sequence) binds to this signature but should be treated differently
       // unless the sequences is a torch.Size
-      return legacy_new_from_sequence(type, r.deviceInt64(1), r.pyobject(0));
+      return legacy_new_from_sequence(type, r.deviceOptional(1), r.pyobject(0));
     }
     return new_with_sizes(type, r.deviceInt64(1), r.intlist(0));
   } else if (r.idx == 5) {
-    return legacy_new_from_sequence(type, r.deviceInt64(1), r.pyobject(0));
+    return legacy_new_from_sequence(type, r.deviceOptional(1), r.pyobject(0));
   }
   throw std::runtime_error("new(): invalid arguments");
 }
@@ -398,22 +406,21 @@ Tensor sparse_coo_tensor_ctor(const Type& type, PyObject* args, PyObject* kwargs
     bool type_inference = r.isNone(2);
     const auto& sparse_type = typeWithDefault(r, 2, 3, default_sparse_type);
     const auto& dense_type = sparse_type.toBackend(sparse_type.is_cuda() ? kCUDA : kCPU);
-    const auto& index_type = dense_type.toScalarType(kLong);
     AutoGPU autogpu(r.deviceInt64(3));
-    // explanation of booleans: allow variables, do type conversion of them, copy numpy data
-    Tensor indices = internal_new_from_data(index_type, -1, r.pyobject(0), false, true, false);
-    Tensor values = internal_new_from_data(dense_type, -1, r.pyobject(1), false, true, type_inference);
+    Tensor values = internal_new_from_data(dense_type, r.deviceOptional(3), r.pyobject(1), false, true, type_inference);
+    // if no dtype provided, infer type based on value type.
+    const auto& index_type = values.type().toScalarType(kLong);
+    Tensor indices = internal_new_from_data(index_type, r.deviceOptional(3), r.pyobject(0), false, true, false);
     const auto& sparse_type_to_use = values.type().toBackend(values.type().is_cuda() ? kSparseCUDA : kSparseCPU);
     return set_requires_grad(sparse_type_to_use.sparse_coo_tensor(indices, values), r.toBool(4));
   } else if (r.idx == 1) {
     bool type_inference = r.isNone(3);
     const auto& sparse_type = typeWithDefault(r, 3, 4, default_sparse_type);
     const auto& dense_type = sparse_type.toBackend(sparse_type.is_cuda() ? kCUDA : kCPU);
-    const auto& index_type = dense_type.toScalarType(kLong);
     AutoGPU autogpu(r.deviceInt64(4));
-    // explanation of booleans: allow variables, do type conversion of them, copy numpy data
-    Tensor indices = internal_new_from_data(index_type, -1, r.pyobject(0), false, true, false);
-    Tensor values = internal_new_from_data(dense_type, -1, r.pyobject(1), false, true, type_inference);
+    Tensor values = internal_new_from_data(dense_type, r.deviceOptional(4), r.pyobject(1), false, true, type_inference);
+    const auto& index_type = values.type().toScalarType(kLong);
+    Tensor indices = internal_new_from_data(index_type, r.deviceOptional(4), r.pyobject(0), false, true, false);
     const auto& sparse_type_to_use = values.type().toBackend(values.type().is_cuda() ? kSparseCUDA : kSparseCPU);
     return set_requires_grad(sparse_type_to_use.sparse_coo_tensor(indices, values, r.intlist(2)), r.toBool(5));
   }
@@ -430,7 +437,7 @@ Tensor tensor_ctor(const Type& type, PyObject* args, PyObject* kwargs) {
   if (r.idx == 0) {
     bool type_inference = r.isNone(1);
     return set_requires_grad(internal_new_from_data(
-        typeWithDefault(r, 1, 2, type), r.deviceInt64(2), r.pyobject(0), true, true, type_inference), r.toBool(3));
+        typeWithDefault(r, 1, 2, type), r.deviceOptional(2), r.pyobject(0), true, true, type_inference), r.toBool(3));
   }
   throw std::runtime_error("tensor(): invalid arguments");
 }
@@ -445,7 +452,7 @@ Tensor new_tensor(const Type& type, PyObject* args, PyObject* kwargs) {
   auto r = parser.parse(args, kwargs, parsed_args);
   if (r.idx == 0) {
     return set_requires_grad(new_from_data_copy(
-        typeWithDefault(r, 1, 2, type), r.deviceInt64(2), r.pyobject(0)), r.toBool(3));
+        typeWithDefault(r, 1, 2, type), r.deviceOptional(2), r.pyobject(0)), r.toBool(3));
   }
   throw std::runtime_error("new_tensor(): invalid arguments");
 }

diff --git a/torch/csrc/utils/tensor_new.h b/torch/csrc/utils/tensor_new.h
@@ -1,13 +1,14 @@
 #pragma once
 
 #include "torch/csrc/python_headers.h"
+#include "torch/csrc/utils/device.h"
 #include <ATen/ATen.h>
 
 namespace torch { namespace utils {
 
 at::Tensor legacy_tensor_ctor(const at::Type& type, PyObject* args, PyObject* kwargs);
 at::Tensor legacy_tensor_new(const at::Type& type, PyObject* args, PyObject* kwargs);
-at::Tensor legacy_new_from_data(const at::Type& type, int device, PyObject *data);
+at::Tensor legacy_new_from_data(const at::Type& type, at::optional<Device> device, PyObject *data);
 at::Tensor sparse_coo_tensor_ctor(const at::Type& type, PyObject* args, PyObject* kwargs);
 at::Tensor tensor_ctor(const at::Type& type, PyObject* args, PyObject* kwargs);
 at::Tensor new_tensor(const at::Type& type, PyObject* args, PyObject* kwargs);