xref: /aosp_15_r20/external/pytorch/torch/csrc/utils/nested.cpp (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
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