Function toDLPack

oneflow/api/python/dlpack/converter.cpp:172–206

void deleter(DLManagedTensor* arg) { delete static_cast<ATenDLMTensor*>(arg->manager_ctx); }

Maybe<DLManagedTensor*> toDLPack(const std::shared_ptr<one::Tensor>& src) {
  auto shape = *src->shape();
  auto strides = *JUST(src->stride());
  // create a new tensor with possibly normalized strides
  // Reference:
  // https://github.com/pytorch/pytorch/issues/83069
  // https://github.com/pytorch/pytorch/issues/82610
  for (int i = 0; i < src->ndim(); i++) {
    if (shape[i] <= 1) { strides[i] = 1; }
  }

  ATenDLMTensor* atDLMTensor(new ATenDLMTensor);
  atDLMTensor->handle = src;
  atDLMTensor->tensor.manager_ctx = atDLMTensor;
  atDLMTensor->tensor.deleter = &deleter;
  JUST(one::SyncAccessTensorWithTimeOut(
      src,
      [&](ep::Stream*, const std::shared_ptr<vm::EagerBlobObject>& tensor) {
        atDLMTensor->tensor.dl_tensor.data = tensor->mut_raw_dptr();
      },
      "const"));
  auto dldevice = JUST(ToDLDevice(JUST(src->device())));
  auto dldtype = JUST(ToDLDataType(src->dtype()->data_type()));
  atDLMTensor->tensor.dl_tensor.device = *dldevice;
  atDLMTensor->tensor.dl_tensor.ndim = src->ndim();
  atDLMTensor->tensor.dl_tensor.dtype = *dldtype;
  atDLMTensor->tensor.dl_tensor.shape =
      // NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast)
      const_cast<int64_t*>(src->shape()->data());
  atDLMTensor->tensor.dl_tensor.strides =
      // NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast)
      const_cast<int64_t*>(JUST(src->stride())->data());
  atDLMTensor->tensor.dl_tensor.byte_offset = 0;
  return &(atDLMTensor->tensor);
}

// This function is mostly copied from PyTorch
void DLPack_Capsule_Destructor(PyObject* data) {