1
#include <c10/util/irange.h>
2
#include <torch/csrc/jit/python/python_arg_flatten.h>
3
#include <torch/csrc/utils/python_strings.h>
4
#include <torch/csrc/utils/six.h>
6
#include <torch/csrc/autograd/grad_mode.h>
8
namespace torch::jit::python {
10
using namespace torch::autograd;
13
// Alphabet used to describe structure of inputs/outputs (D for desc)
15
static constexpr char DictOpen = '<';
16
static constexpr char DictClose = '>';
17
static constexpr char ListOpen = '[';
18
static constexpr char ListClose = ']';
19
static constexpr char TupleOpen = '(';
20
static constexpr char TupleClose = ')';
21
static constexpr char Variable = 'v';
22
static constexpr char Bool = 'b';
23
static constexpr char Long = 'l';
24
static constexpr char Double = 'd';
25
static constexpr char String = 's';
26
static constexpr char NoneType = 'n';
31
inline bool PyNone_Check(PyObject* o) {
36
py::object cast_handle_sequence(std::vector<py::handle> objs) {
37
auto num_objs = objs.size();
39
for (const auto i : c10::irange(num_objs)) {
40
sequence[i] = py::reinterpret_borrow<py::object>(objs[i]);
45
void flatten_rec(PyObject* obj, ParsedArgs& args) {
46
auto& structure = args.desc.structure;
47
if (six::isTuple(obj)) {
48
structure.push_back(D::TupleOpen);
49
for (auto item : py::reinterpret_borrow<py::tuple>(obj))
50
flatten_rec(item.ptr(), args);
51
structure.push_back(D::TupleClose);
52
} else if (PyList_Check(obj)) {
53
structure.push_back(D::ListOpen);
54
for (auto item : py::reinterpret_borrow<py::list>(obj))
55
flatten_rec(item.ptr(), args);
56
structure.push_back(D::ListClose);
57
} else if (PyDict_Check(obj)) {
58
auto* dict_items = PyDict_Items(obj);
59
structure.push_back(D::DictOpen);
60
for (auto item : py::reinterpret_borrow<py::list>(dict_items)) {
61
flatten_rec(item.ptr(), args);
63
structure.push_back(D::DictClose);
64
Py_DECREF(dict_items);
65
} else if (THPUtils_checkString(obj)) {
66
string str = THPUtils_unpackString(obj);
67
args.desc.strings.emplace_back(str);
68
args.desc.structure.push_back(D::String);
69
} else if (THPVariable_Check(obj)) {
70
auto& var = THPVariable_Unpack(obj);
71
args.vars.push_back(var);
72
args.desc.metadata.emplace_back(var);
73
args.desc.structure.push_back(D::Variable);
74
} else if (PyNone_Check(obj)) {
75
args.desc.structure.push_back(D::NoneType);
76
} else if (PyBool_Check(obj)) { // Wrap bools in Bool tensors
77
at::Tensor var = scalar_to_tensor(at::Scalar(THPUtils_unpackBool(obj)));
78
args.vars.push_back(var);
79
args.desc.metadata.emplace_back(var);
80
args.desc.structure.push_back(D::Bool);
81
} else if (PyLong_Check(obj)) { // Wrap longs in Long tensors
82
at::Tensor var = scalar_to_tensor(
83
at::Scalar(static_cast<int64_t>(THPUtils_unpackLong(obj))));
84
args.vars.push_back(var);
85
args.desc.metadata.emplace_back(var);
86
args.desc.structure.push_back(D::Long);
87
} else if (PyFloat_Check(obj)) { // Wrap floats in Double tensors
88
at::Tensor var = scalar_to_tensor(THPUtils_unpackDouble(obj));
89
args.vars.push_back(var);
90
args.desc.metadata.emplace_back(var);
91
args.desc.structure.push_back(D::Double);
94
"Only tuples, lists and Variables are supported as JIT inputs/outputs. "
95
"Dictionaries and strings are also accepted, but their usage is not "
96
"recommended. Here, received an input of unsupported type: ";
97
msg += THPUtils_typename(obj);
98
throw std::runtime_error(msg);
102
} // anonymous namespace
104
ParsedArgs flatten(py::handle obj) {
106
args.desc.grad_enabled = autograd::GradMode::is_enabled();
107
flatten_rec(obj.ptr(), args);
114
py::object cast_sequence(std::vector<py::object> objs) {
115
auto num_objs = objs.size();
116
T sequence{num_objs};
117
for (const auto i : c10::irange(num_objs)) {
118
sequence[i] = std::move(objs[i]);
120
return std::move(sequence);
123
py::object cast_dict(std::vector<py::object> objs) {
124
auto num_objs = objs.size();
125
py::dict sequence = {};
126
for (const auto i : c10::irange(num_objs)) {
127
py::tuple obj = py::reinterpret_borrow<py::tuple>(objs[i]);
128
sequence[obj[0]] = obj[1];
130
return std::move(sequence);
133
py::object unflatten_rec(
134
ArrayRef<Variable>::iterator& var_it,
135
ArrayRef<Variable>::iterator& var_it_end,
136
std::string::const_iterator& desc_it,
137
std::vector<string>::const_iterator& str_it,
138
std::vector<string>::const_iterator& str_it_end) {
139
char type = *desc_it++;
140
if (type == D::TupleOpen) {
141
std::vector<py::object> objs;
142
while (*desc_it != D::TupleClose)
144
unflatten_rec(var_it, var_it_end, desc_it, str_it, str_it_end));
146
return cast_sequence<py::tuple>(objs);
147
} else if (type == D::ListOpen) {
148
std::vector<py::object> objs;
149
while (*desc_it != D::ListClose)
151
unflatten_rec(var_it, var_it_end, desc_it, str_it, str_it_end));
153
return cast_sequence<py::list>(objs);
154
} else if (type == D::DictOpen) {
155
std::vector<py::object> objs;
156
while (*desc_it != D::DictClose) {
158
unflatten_rec(var_it, var_it_end, desc_it, str_it, str_it_end));
161
return cast_dict(objs);
162
} else if (type == D::String) {
163
if (str_it == str_it_end)
164
throw std::runtime_error("Not enough Variables given to unflatten");
165
auto str = *str_it++;
166
return py::reinterpret_borrow<py::object>(THPUtils_packString(str));
167
} else if (type == D::NoneType) {
168
return py::reinterpret_borrow<py::object>(py::none());
170
// if (type == D::Long || type == D::Double || type == D::Bool ||
171
// D::Variable) unwrap variables (D::Variable), or unwrap primitive types
172
// (Long, Double, Bool) as variables for tracer.
173
if (var_it == var_it_end)
174
throw std::runtime_error("Not enough Variables given to unflatten");
175
auto var = *var_it++;
176
return py::reinterpret_steal<py::object>(THPVariable_Wrap(var));
180
} // anonymous namespace
182
PyObject* unflatten(ArrayRef<Variable> vars, const IODescriptor& desc) {
183
// NB: We don't do correctness checking on descriptor.
184
// It has to be a correct bytes object produced by unflatten.
185
auto vars_it = vars.begin();
186
auto vars_it_end = vars.end();
187
auto desc_it = desc.structure.begin();
188
std::vector<std::string>::const_iterator str_it = desc.strings.begin();
189
std::vector<std::string>::const_iterator str_end = desc.strings.end();
190
auto output = unflatten_rec(vars_it, vars_it_end, desc_it, str_it, str_end);
191
if (vars_it != vars_it_end)
192
throw std::runtime_error("Too many Variables given to unflatten");
193
return output.release().ptr();
196
} // namespace torch::jit::python