Function fromDLPack

oneflow/api/python/dlpack/converter.cpp:83–130

}

Maybe<one::Tensor> fromDLPack(const DLManagedTensor* src) {
  using namespace one;
  const auto& dl_tensor = src->dl_tensor;

  Symbol<Device> device = JUST(ToOneFlowDevice(dl_tensor.device));
  DataType dtype = JUST(ToOneFlowDataType(dl_tensor.dtype));

  // Build TensorMeta
  const Shape shape(dl_tensor.shape, dl_tensor.shape + dl_tensor.ndim);
  Symbol<LocalTensorMeta> tensor_meta;
  if (dl_tensor.strides) {
    const auto stride = Stride(dl_tensor.strides, dl_tensor.strides + dl_tensor.ndim);
    tensor_meta =
        SymbolOf(LocalTensorMeta(shape, stride, dtype, MemoryFormat::kContiguous, device));
  } else {
    tensor_meta = SymbolOf(LocalTensorMeta(shape, dtype, MemoryFormat::kContiguous, device));
  }

  // Build TensorBuffer
  const auto& Free = [src](char* dptr) {
    if (src->deleter) {
      // NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast)
      src->deleter(const_cast<DLManagedTensor*>(src));
    }
  };

  size_t array_size_in_bytes = shape.elem_cnt() * GetSizeOfDataType(dtype);
  auto tensor_data = std::make_shared<vm::TensorStorage>(false, device);
  tensor_data->set_blob_dptr(
      std::unique_ptr<char, std::function<void(char*)>>(static_cast<char*>(dl_tensor.data), Free),
      array_size_in_bytes);

  // Build TensorStorage: decrease ndarray reference count before releasing
  auto tensor_storage = std::make_shared<TensorStorage>(tensor_data);

  // Build Tensor
  auto tensor_impl = std::make_shared<EagerLocalTensorImpl>(tensor_storage,
                                                            /*requires_grad=*/false,
                                                            /*ls_leaf=*/true);

  // Init blob
  JUST(tensor_impl->InitEagerBlobObject(tensor_meta, NewLocalDepObject()));
  const auto& stream = JUST(GetDefaultStreamByDevice(device));
  const auto& eager_blob_object = JUST(tensor_impl->eager_blob_object());
  JUST(eager_blob_object->init_producer_stream(stream));
  eager_blob_object->set_last_used_stream(stream);
  return std::static_pointer_cast<Tensor>(std::make_shared<LocalTensor>(tensor_impl));
}

Maybe<DLDevice> ToDLDevice(Symbol<Device> ofdevice) {
  DLDevice ctx;