| 208 | } // namespace |
| 209 | |
| 210 | Maybe<Tensor> MakeGlobalTensorFromData(PyObject* data, const Optional<Symbol<DType>>& dtype, |
| 211 | Symbol<ParallelDesc> placement, |
| 212 | const std::vector<Symbol<SbpParallel>>& sbp_tuple, |
| 213 | const bool requires_grad) { |
| 214 | PyObject* array = NULL; |
| 215 | if (PyArray_Check(data)) { |
| 216 | // Only NPY_CORDER is supported, and returns a new C-style contiguous array. |
| 217 | array = PyArray_NewCopy((PyArrayObject*)data, NPY_CORDER); |
| 218 | } else { |
| 219 | // NPY_ARRAY_DEFAULT is NPY_ARRAY_C_CONTIGUOUS | NPY_ARRAY_BEHAVED, so the |
| 220 | // array with NPY_ARRAY_DEFAULT flag is C-style contiguous. |
| 221 | array = PyArray_FromAny(data, nullptr, 0, 0, NPY_ARRAY_DEFAULT | NPY_ARRAY_ENSURECOPY, nullptr); |
| 222 | if (!array) { return Error::RuntimeError() << "Can not convert input data to a numpy array."; } |
| 223 | } |
| 224 | auto* np_arr = reinterpret_cast<PyArrayObject*>(array); |
| 225 | const npy_intp* dims_ptr = PyArray_SHAPE(np_arr); |
| 226 | const Shape shape(DimVector(dims_ptr, dims_ptr + PyArray_NDIM(np_arr))); |
| 227 | DataType data_type = JUST(numpy::GetOFDataTypeFromNpArray(np_arr)); |
| 228 | |
| 229 | if (placement->parallel_num() > 1) { |
| 230 | const void* buf_ptr = PyArray_DATA(np_arr); |
| 231 | size_t array_size = PyArray_SIZE(np_arr); |
| 232 | CHECK_EQ_OR_RETURN(array_size, shape.elem_cnt()); |
| 233 | size_t byte_size = array_size * GetSizeOfDataType(data_type); |
| 234 | JUST(DataConsistencyCheck(buf_ptr, byte_size, placement)); |
| 235 | } |
| 236 | |
| 237 | Symbol<Device> device = JUST(Device::New(placement->device_tag())); |
| 238 | std::shared_ptr<Tensor> local_tensor; |
| 239 | { |
| 240 | GlobalMode::Guard guard(/* disable global mode */ false); |
| 241 | local_tensor = |
| 242 | JUST(functional::Empty(shape, JUST(DType::Get(data_type)), device, /*requires_grad=*/false, |
| 243 | /*pin_memory=*/false)); |
| 244 | } |
| 245 | if (device->enum_type() != DeviceType::kMeta) { |
| 246 | JUST(CopyLocalTensorFromUntypedArray(local_tensor, array)); |
| 247 | } |
| 248 | |
| 249 | Py_DECREF(array); |
| 250 | // Cast to float if data is double sequence, rather than numpy array. |
| 251 | Symbol<DType> dtype_; |
| 252 | if (dtype) { |
| 253 | dtype_ = JUST(dtype); |
| 254 | } else if (!dtype && data_type == DataType::kDouble && !PyArray_Check(data)) { |
| 255 | dtype_ = DType::Float(); |
| 256 | } |
| 257 | if (dtype_) { local_tensor = JUST(functional::Cast(local_tensor, dtype_, /*pin_memory=*/false)); } |
| 258 | |
| 259 | size_t sbp_dims = sbp_tuple.size(); |
| 260 | Symbol<NdSbp> broadcast_nd_sbp = JUST(CachedGetAllBroadcastNdSbp(sbp_dims)); |
| 261 | |
| 262 | std::shared_ptr<Tensor> broadcast_tensor = JUST( |
| 263 | functional::LocalToGlobal(local_tensor, placement, *JUST(GetSbpList(broadcast_nd_sbp)), shape, |
| 264 | local_tensor->dtype(), /* sync_data */ true, /*copy=*/false)); |
| 265 | |
| 266 | std::vector<Symbol<SbpParallel>> grad_sbp_tuple; |
| 267 | auto global_tensor = |
no test coverage detected