| 50 | } |
| 51 | |
| 52 | void TensorToNumpy(const user_op::Tensor* tensor, PyObject** arg_ptr) { |
| 53 | if (tensor == nullptr) { |
| 54 | Py_INCREF(Py_None); |
| 55 | *arg_ptr = Py_None; |
| 56 | return; |
| 57 | } |
| 58 | int type_num = CHECK_JUST(numpy::OFDataTypeToNumpyType(tensor->data_type())); |
| 59 | VLOG(3) << "Tensor data type " << DataType_Name(tensor->data_type()) << " Numpy type " |
| 60 | << type_num; |
| 61 | int dim_size = tensor->shape_view().NumAxes(); |
| 62 | npy_intp dims[dim_size]; |
| 63 | FOR_RANGE(size_t, i, 0, dim_size) { dims[i] = tensor->shape_view().At(i); } |
| 64 | |
| 65 | void* data = TensorToMem(tensor); |
| 66 | auto* np_array = |
| 67 | reinterpret_cast<PyArrayObject*>(PyArray_SimpleNewFromData(dim_size, dims, type_num, data)); |
| 68 | // Numpy will not release the data |
| 69 | PyArray_CLEARFLAGS(np_array, NPY_ARRAY_OWNDATA); |
| 70 | *arg_ptr = reinterpret_cast<PyObject*>(np_array); |
| 71 | } |
| 72 | |
| 73 | #define TENSOR_MEM_ASSIGN(dtype) \ |
| 74 | do { \ |
no test coverage detected