Caffe2 - C++ API
A deep learning, cross platform ML framework
wrap_outputs.h
1 #pragma once
2 
3 // Wrap tensor operation outputs as PyObject*
4 
5 #include <ATen/ATen.h>
6 #include <torch/csrc/python_headers.h>
7 #include <tuple>
8 
9 #include <torch/csrc/Dtype.h>
10 #include <torch/csrc/Layout.h>
11 #include <torch/csrc/autograd/python_variable.h>
12 #include <torch/csrc/autograd/variable.h>
13 #include <torch/csrc/utils/python_numbers.h>
14 
15 namespace torch { namespace autograd { namespace utils {
16 
17 inline PyObject* wrap(bool value) {
18  if (value) {
19  Py_RETURN_TRUE;
20  } else {
21  Py_RETURN_FALSE;
22  }
23 }
24 
25 inline PyObject* wrap(int64_t value) {
26  return THPUtils_packInt64(value);
27 }
28 
29 inline PyObject* wrap(double value) {
30  return PyFloat_FromDouble(value);
31 }
32 
33 inline PyObject* wrap(std::complex<double> value) {
34  // I could probably also use FromComplex with a reinterpret cast,
35  // but... eh.
36  return PyComplex_FromDoubles(value.real(), value.imag());
37 }
38 
39 inline PyObject* wrap(void* value) {
40  return THPUtils_packInt64(reinterpret_cast<intptr_t>(value));
41 }
42 
43 inline PyObject* wrap(THPDtype *dtype) {
44  Py_INCREF(dtype);
45  return (PyObject*)dtype;
46 }
47 
48 inline PyObject* wrap(THPLayout *layout) {
49  Py_INCREF(layout);
50  return (PyObject*)layout;
51 }
52 
53 inline PyObject* wrap(at::Tensor tensor) {
54  return THPVariable_Wrap(Variable(std::move(tensor)));
55 }
56 
57 inline PyObject* wrap(at::Scalar scalar) {
58  return wrap(scalar_to_tensor(scalar));
59 }
60 
61 inline PyObject* wrap(std::tuple<at::Tensor, at::Tensor> tensors) {
62  auto r = THPObjectPtr{PyTuple_New(2)};
63  if (!r) throw python_error();
64  PyTuple_SET_ITEM(r.get(), 0, wrap(std::get<0>(tensors)));
65  PyTuple_SET_ITEM(r.get(), 1, wrap(std::get<1>(tensors)));
66  return r.release();
67 }
68 
69 inline PyObject* wrap(PyTypeObject *type, std::tuple<at::Tensor, at::Tensor> tensors) {
70  auto r = THPObjectPtr{PyStructSequence_New(type)};
71  if (!r) throw python_error();
72  PyStructSequence_SET_ITEM(r.get(), 0, wrap(std::get<0>(tensors)));
73  PyStructSequence_SET_ITEM(r.get(), 1, wrap(std::get<1>(tensors)));
74  return r.release();
75 }
76 
77 inline PyObject* wrap(std::tuple<at::Tensor, at::Tensor, at::Tensor> tensors) {
78  auto r = THPObjectPtr{PyTuple_New(3)};
79  if (!r) throw python_error();
80  PyTuple_SET_ITEM(r.get(), 0, wrap(std::move(std::get<0>(tensors))));
81  PyTuple_SET_ITEM(r.get(), 1, wrap(std::move(std::get<1>(tensors))));
82  PyTuple_SET_ITEM(r.get(), 2, wrap(std::move(std::get<2>(tensors))));
83  return r.release();
84 }
85 
86 inline PyObject* wrap(PyTypeObject *type, std::tuple<at::Tensor, at::Tensor, at::Tensor> tensors) {
87  auto r = THPObjectPtr{PyStructSequence_New(type)};
88  if (!r) throw python_error();
89  PyStructSequence_SET_ITEM(r.get(), 0, wrap(std::get<0>(tensors)));
90  PyStructSequence_SET_ITEM(r.get(), 1, wrap(std::get<1>(tensors)));
91  PyStructSequence_SET_ITEM(r.get(), 2, wrap(std::get<2>(tensors)));
92  return r.release();
93 }
94 
95 inline PyObject* wrap(std::tuple<at::Tensor, at::Tensor, at::Tensor, int64_t> tensors) {
96  auto r = THPObjectPtr{PyTuple_New(4)};
97  if (!r) throw python_error();
98  PyTuple_SET_ITEM(r.get(), 0, wrap(std::move(std::get<0>(tensors))));
99  PyTuple_SET_ITEM(r.get(), 1, wrap(std::move(std::get<1>(tensors))));
100  PyTuple_SET_ITEM(r.get(), 2, wrap(std::move(std::get<2>(tensors))));
101  PyTuple_SET_ITEM(r.get(), 3, wrap(std::get<3>(tensors)));
102  return r.release();
103 }
104 
105 inline PyObject* wrap(std::tuple<at::Tensor, at::Tensor, float, int64_t> tensors) {
106  auto r = THPObjectPtr{PyTuple_New(4)};
107  if (!r) throw python_error();
108  PyTuple_SET_ITEM(r.get(), 0, wrap(std::move(std::get<0>(tensors))));
109  PyTuple_SET_ITEM(r.get(), 1, wrap(std::move(std::get<1>(tensors))));
110  PyTuple_SET_ITEM(r.get(), 2, wrap(std::move(std::get<2>(tensors))));
111  PyTuple_SET_ITEM(r.get(), 3, wrap(std::move(std::get<3>(tensors))));
112  return r.release();
113 }
114 
115 inline PyObject* wrap(std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> tensors) {
116  auto r = THPObjectPtr{PyTuple_New(4)};
117  if (!r) throw python_error();
118  PyTuple_SET_ITEM(r.get(), 0, wrap(std::move(std::get<0>(tensors))));
119  PyTuple_SET_ITEM(r.get(), 1, wrap(std::move(std::get<1>(tensors))));
120  PyTuple_SET_ITEM(r.get(), 2, wrap(std::move(std::get<2>(tensors))));
121  PyTuple_SET_ITEM(r.get(), 3, wrap(std::move(std::get<3>(tensors))));
122  return r.release();
123 }
124 
125 inline PyObject* wrap(std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor> tensors) {
126  auto r = THPObjectPtr{PyTuple_New(5)};
127  if (!r) throw python_error();
128  PyTuple_SET_ITEM(r.get(), 0, wrap(std::move(std::get<0>(tensors))));
129  PyTuple_SET_ITEM(r.get(), 1, wrap(std::move(std::get<1>(tensors))));
130  PyTuple_SET_ITEM(r.get(), 2, wrap(std::move(std::get<2>(tensors))));
131  PyTuple_SET_ITEM(r.get(), 3, wrap(std::move(std::get<3>(tensors))));
132  PyTuple_SET_ITEM(r.get(), 4, wrap(std::move(std::get<4>(tensors))));
133  return r.release();
134 }
135 
136 inline PyObject* wrap(at::TensorList tl) {
137  auto r = THPObjectPtr{PyTuple_New(tl.size())};
138  if (!r) throw python_error();
139  for (size_t i = 0; i < tl.size(); ++i) {
140  PyTuple_SET_ITEM(r.get(), i, wrap(tl[i]));
141  }
142  return r.release();
143 }
144 
145 }}} // namespace torch::autograd::utils
Scalar represents a 0-dimensional tensor which contains a single element.
Definition: Scalar.h:22
constexpr size_t size() const
size - Get the array size.
Definition: ArrayRef.h:138
Definition: Dtype.h:9
Definition: jit_type.h:17
ArrayRef - Represent a constant reference to an array (0 or more elements consecutively in memory)...
Definition: ArrayRef.h:41