MCPcopy Create free account
hub / github.com/Oneflow-Inc/oneflow / UnpackTensorTuple

Function UnpackTensorTuple

oneflow/api/python/autograd/autograd_function.cpp:35–59  ·  view source on GitHub ↗

Transform input to TensorTuple

Source from the content-addressed store, hash-verified

33
34// Transform input to TensorTuple
35Maybe<one::TensorTuple> UnpackTensorTuple(const py::object& input) {
36 one::TensorTuple tp;
37 if (one::PyTensor_Check(input.ptr())) {
38 tp.emplace_back(input.cast<std::shared_ptr<one::Tensor>>());
39 } else if (py::isinstance<py::tuple>(input)) {
40 auto tuple = input.cast<py::tuple>();
41 tp.resize(tuple.size());
42 for (int i = 0; i < tuple.size(); ++i) {
43 PyObject* obj = tuple[i].ptr();
44 if (obj == Py_None) {
45 // do nothing
46 } else if (one::PyTensor_Check(obj)) {
47 tp[i] = one::PyTensor_Unpack(obj);
48 } else {
49 return Error::RuntimeError()
50 << "expected Tensor or None as element " << i << ", but got "
51 << one::functional::PyStringAsString(PyObject_Str((PyObject*)Py_TYPE(obj)));
52 }
53 }
54 } else {
55 return Error::RuntimeError()
56 << "autograd.Function's output only support tensor or list of tensors";
57 }
58 return tp;
59}
60
61// Return single Tensor when TensorTuple's size is one, otherwise py::tuple
62py::object PackTensorTuple(const one::TensorTuple& tp) {

Callers 2

PackPyFunctionToFTypeFunction · 0.85

Calls 6

PyTensor_CheckFunction · 0.85
PyStringAsStringFunction · 0.85
ptrMethod · 0.45
emplace_backMethod · 0.45
resizeMethod · 0.45
sizeMethod · 0.45

Tested by

no test coverage detected