| 137 | } |
| 138 | |
| 139 | void GetPyOutputs(const UserOpDef& op_def, user_op::KernelComputeContext* ctx, |
| 140 | PyObject* py_outputs) { |
| 141 | const size_t kernel_out_num = ctx->outputs().size(); |
| 142 | const size_t def_out_num = op_def.output_size(); |
| 143 | CHECK_EQ(kernel_out_num, def_out_num) << "kernel output num " << kernel_out_num |
| 144 | << " not equal to definition output num " << def_out_num; |
| 145 | if (PyList_Check(py_outputs)) { |
| 146 | FOR_RANGE(size_t, i, 0, def_out_num) { |
| 147 | const std::string& arg_name = op_def.output(i).name(); |
| 148 | VLOG(3) << "output arg_name " << arg_name; |
| 149 | int32_t index = 0; |
| 150 | NumpyToTensor(PyList_GetItem(py_outputs, i), ctx->Tensor4ArgNameAndIndex(arg_name, index)); |
| 151 | } |
| 152 | } else if (PyArray_Check(py_outputs)) { |
| 153 | const std::string& arg_name = ctx->outputs().at(0).first; |
| 154 | VLOG(3) << "output arg_name " << arg_name; |
| 155 | int32_t index = 0; |
| 156 | NumpyToTensor(py_outputs, ctx->Tensor4ArgNameAndIndex(arg_name, index)); |
| 157 | } else { |
| 158 | LOG(FATAL) << "Unexpeted PyObject was returned: " << Py_TYPE(py_outputs)->tp_name; |
| 159 | } |
| 160 | } |
| 161 | |
| 162 | } // namespace |
| 163 |
no test coverage detected