1 #include <ATen/ATen.h>
2 #include <ATen/NestedTensorImpl.h>
3 #include <c10/core/ScalarType.h>
4 #include <torch/csrc/python_headers.h>
5 #include <torch/csrc/utils/nested.h>
6 #include <torch/csrc/utils/pybind.h>
7 #include <torch/csrc/utils/tensor_new.h>
8 #include <torch/torch.h>
9 #include <stdexcept>
10 #include <vector>
11
12 namespace torch::utils {
13
14 // NB: device_idx here is NOT a DeviceIndex, but index into PythonArgs
typeIdWithDefault(PythonArgs & r,int device_idx,c10::DispatchKey dispatch_key)15 static c10::TensorOptions typeIdWithDefault(
16 PythonArgs& r,
17 int device_idx,
18 c10::DispatchKey dispatch_key) {
19 auto options = dispatchKeyToTensorOptions(dispatch_key);
20 if (!r.isNone(device_idx)) {
21 options = options.device(r.device(device_idx));
22 }
23 return options;
24 }
25
nested_tensor_ctor(c10::DispatchKey dispatch_key,at::ScalarType scalar_type,torch::PythonArgs & r)26 at::Tensor nested_tensor_ctor(
27 c10::DispatchKey dispatch_key,
28 at::ScalarType scalar_type,
29 torch::PythonArgs& r) {
30 TORCH_CHECK(r.idx == 0, "nested_tensor(): invalid arguments");
31
32 PyObject* data = r.pyobject(0);
33 // Check if data is a list: Only List[Tensor] and List[List...[Scalar]] are
34 // accepted for now
35 TORCH_CHECK_TYPE(
36 PyList_Check(data),
37 "Only lists (List[Tensor] and List[List...[Scalar]]) are accepted in nested_tensor");
38
39 auto dtype_val = r.scalartypeWithDefault(1, scalar_type);
40 auto tensor_options = typeIdWithDefault(r, 2, dispatch_key);
41 bool pin_memory = r.toBool(3);
42 bool args_requires_grad = r.toBool(4);
43
44 TORCH_CHECK(
45 PyList_Size(data) >= 0,
46 "Something went really wrong and your list has negative size");
47
48 // Check whether we are dealing with lists of tensors or not
49 std::vector<at::Tensor> new_list(PyList_Size(data));
50 for (const auto i : c10::irange(PyList_Size(data))) {
51 PyObject* elem = PyList_GetItem(data, i);
52 if (THPVariable_Check(elem)) {
53 new_list[i] = THPVariable_Unpack(PyList_GetItem(data, i)).detach();
54 TORCH_CHECK(
55 !new_list[i].is_nested(),
56 "We do not accept nested tensors as input to nested tensors");
57 TORCH_CHECK(
58 new_list[i].layout() == kStrided,
59 "We do not accept non-strided layouts as input to nested tensors");
60 } else {
61 PythonArgs elem_r(r);
62 std::array<PyObject*, 6> elem_args = {
63 elem, // data
64 r.args[1], // dtpye
65 nullptr, // device (cpu)
66 nullptr, // no pinned memory
67 r.args[4], // requires grad
68 nullptr // names
69 };
70 elem_r.args = elem_args.data();
71 new_list[i] = tensor_ctor(dispatch_key, scalar_type, elem_r);
72 }
73 }
74
75 at::ScalarType final_dtype = dtype_val;
76 if (r.isNone(1) && !new_list.empty()) {
77 final_dtype = c10::typeMetaToScalarType(new_list[0].dtype());
78 }
79 at::Device final_device = tensor_options.device();
80 if (r.isNone(2) && !new_list.empty()) {
81 final_device = new_list[0].device();
82 }
83 auto out = at::_nested_tensor_from_tensor_list(
84 new_list, final_dtype, std::nullopt, final_device, pin_memory);
85 out.requires_grad_(args_requires_grad);
86 return out;
87 }
88
89 } // namespace torch::utils
90