| 106 | } |
| 107 | |
| 108 | Result<DLDevice> ExportDevice(const std::shared_ptr<Array>& arr) { |
| 109 | // Check if array is supported by the DLPack protocol. |
| 110 | if (arr->null_count() > 0) { |
| 111 | return Status::TypeError("Can only use DLPack on arrays with no nulls."); |
| 112 | } |
| 113 | const DataType& type = *arr->type(); |
| 114 | if (type.id() == Type::BOOL) { |
| 115 | return Status::TypeError("Bit-packed boolean data type not supported by DLPack."); |
| 116 | } |
| 117 | if (!is_integer(type.id()) && !is_floating(type.id())) { |
| 118 | return Status::TypeError("DataType is not compatible with DLPack spec: ", |
| 119 | type.ToString()); |
| 120 | } |
| 121 | |
| 122 | // Define DLDevice struct |
| 123 | DLDevice device; |
| 124 | if (arr->data()->buffers[1]->device_type() == DeviceAllocationType::kCPU) { |
| 125 | device.device_id = 0; |
| 126 | device.device_type = DLDeviceType::kDLCPU; |
| 127 | return device; |
| 128 | } else { |
| 129 | return Status::NotImplemented( |
| 130 | "DLPack support is implemented only for buffers on CPU device."); |
| 131 | } |
| 132 | } |
| 133 | |
| 134 | struct TensorManagerCtx { |
| 135 | std::shared_ptr<Tensor> t; |