Function MakeLocalTensorFromData

oneflow/api/python/utils/tensor_utils.cpp:129–194

}

Maybe<Tensor> MakeLocalTensorFromData(PyObject* data, const Optional<Symbol<DType>>& dtype,
                                      const Optional<Symbol<Device>>& device,
                                      const bool requires_grad, const bool pin_memory) {
  bool is_bfloat16_dtype = dtype ? JUST(dtype)->data_type() == DataType::kBFloat16 : false;
  bool is_cuda_device = device ? JUST(device)->enum_type() == DeviceType::kCUDA : false;
  if (is_bfloat16_dtype && is_cuda_device) {
#if CUDA_VERSION < 11000
    return Error::RuntimeError()
           << "Cannot create a bfloat16 tensor on gpu under cuda version: 11000";
#endif  // CUDA_VERSION >= 11000
  }
  PyArray_Descr* np_dtype =
      dtype.has_value() && !is_bfloat16_dtype
          ? PyArray_DescrFromType(JUST(numpy::OFDataTypeToNumpyType(JUST(dtype)->data_type())))
          : nullptr;
  // NPY_ARRAY_DEFAULT is NPY_ARRAY_C_CONTIGUOUS | NPY_ARRAY_BEHAVED, so the
  // array with NPY_ARRAY_DEFAULT flag is C-style contiguous.
  // NPY_ARRAY_FORCECAST is needed otherwise there will a segfault.
  //
  // Even though PyArray_FromAny can cast the input array to the desired dtype
  // if `dtype` argument is set, it fails to handle the following case:
  // >> x = [flow.tensor([1, 2])] * 3 <-- x is a list of flow.Tensor
  // >> y = flow.tensor(x, dtype=flow.float32) <-- returns nullptr
  // However, the following case without `dtype` argument works well:
  // >> x = [flow.tensor([1, 2])] * 3
  // >> y = flow.tensor(x)
  // So we cast the input array to the desired dtype manually.
  PyArrayObject* _array = reinterpret_cast<PyArrayObject*>(
      PyArray_FromAny(data, nullptr, 0, 0,
                      NPY_ARRAY_DEFAULT | NPY_ARRAY_ENSURECOPY | NPY_ARRAY_FORCECAST, nullptr));
  if (!_array) {
    return Error::RuntimeError() << "Can not convert input data to a new numpy array.";
  }
  // PyArray_FromArray steals a reference to np_dtype object, so no need to decref it.
  PyObject* array = PyArray_FromArray(
      _array, np_dtype, NPY_ARRAY_DEFAULT | NPY_ARRAY_ENSURECOPY | NPY_ARRAY_FORCECAST);
  Py_DECREF(_array);
  auto* np_arr = reinterpret_cast<PyArrayObject*>(array);
  const npy_intp* dims_ptr = PyArray_SHAPE(np_arr);
  const Shape shape(DimVector(dims_ptr, dims_ptr + PyArray_NDIM(np_arr)));
  DataType np_data_type = JUST(numpy::GetOFDataTypeFromNpArray(np_arr));

  Symbol<Device> device_;
  if (device) {
    device_ = JUST(device);
  } else {
    device_ = JUST(Device::New("cpu"));
  }
  std::shared_ptr<Tensor> tensor =
      JUST(functional::Empty(shape, JUST(DType::Get(np_data_type)), device_,
                             /*requires_grad=*/false, /*pin_memory=*/pin_memory));
  if (device_->enum_type() != DeviceType::kMeta) {
    JUST(CopyLocalTensorFromUntypedArray(tensor, array));
  }

  Py_DECREF(array);
  if (dtype && JUST(dtype)->data_type() != np_data_type) {
    tensor = JUST(functional::To(tensor, JUST(dtype), false));