Function GetItemInScalarTensor

oneflow/core/framework/tensor_util.cpp:75–104

}

Maybe<void> GetItemInScalarTensor(const std::shared_ptr<Tensor>& scalar_tensor, void* scalar_ptr,
                                  size_t size) {
  CHECK_EQ_OR_RETURN(GetSizeOfDataType(scalar_tensor->dtype()->data_type()), size)
      << "invalid size";
  CHECK_OR_RETURN(scalar_tensor->is_eager()) << "Only eager scalar tensor support GetItem.";
  CHECK_EQ_OR_RETURN(scalar_tensor->nelement(), 1)
      << "can only convert a tensor of size 1 to a Python scalar";
  std::shared_ptr<LocalTensor> local_tensor;
  {
    auto tensor = scalar_tensor;
    if (tensor->is_global()) {
      Symbol<ParallelDesc> parallel_desc;
      {
        const ParallelConf parallel_conf = GenParallelConfOfCpuOnAllRanks();
        JUST(PhysicalRun(
            [&parallel_desc, &parallel_conf](InstructionsBuilder* builder) -> Maybe<void> {
              parallel_desc = SymbolOf(*JUST(builder->GetParallelDescSymbol(parallel_conf)));
              return Maybe<void>::Ok();
            }));
      }
      const auto& broadcast_sbp = JUST(MakeBroadcastSbpParallel());
      tensor = JUST(functional::ToGlobal(tensor, parallel_desc, {broadcast_sbp}, /*grad_sbp=*/{},
                                         /*check_meta=*/false, /*copy=*/false));
      tensor = JUST(functional::GlobalToLocal(tensor, /*copy=*/false));
    }
    local_tensor = JUST(tensor->AsLocalTensor());
  }
  JUST(SyncReadSmallMem(reinterpret_cast<char*>(scalar_ptr), size, local_tensor));
  return Maybe<void>::Ok();
}

}  // namespace one
}  // namespace oneflow