| 33 | #define TENSOR_MEM_CAST(dtype) static_cast<void*>(const_cast<dtype*>(tensor->dptr<dtype>())) |
| 34 | |
| 35 | void* TensorToMem(const user_op::Tensor* tensor) { |
| 36 | switch (tensor->data_type()) { |
| 37 | case DataType::kFloat: return TENSOR_MEM_CAST(float); |
| 38 | case DataType::kDouble: return TENSOR_MEM_CAST(double); |
| 39 | case DataType::kBool: return TENSOR_MEM_CAST(bool); |
| 40 | case DataType::kInt8: return TENSOR_MEM_CAST(int8_t); |
| 41 | case DataType::kInt32: return TENSOR_MEM_CAST(int32_t); |
| 42 | case DataType::kInt64: return TENSOR_MEM_CAST(int64_t); |
| 43 | case DataType::kUInt8: return TENSOR_MEM_CAST(uint8_t); |
| 44 | case DataType::kFloat16: return TENSOR_MEM_CAST(float16); |
| 45 | default: |
| 46 | LOG(FATAL) << "OneFlow data type " << DataType_Name(tensor->data_type()) |
| 47 | << " is not supported yet."; |
| 48 | return nullptr; |
| 49 | } |
| 50 | } |
| 51 | |
| 52 | void TensorToNumpy(const user_op::Tensor* tensor, PyObject** arg_ptr) { |
| 53 | if (tensor == nullptr) { |
no test coverage detected