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