Function SetKernelDistOutput

paddle/phi/api/lib/api_gen_utils.cc:635–651

}

phi::distributed::DistTensor* SetKernelDistOutput(
    Tensor* out, const phi::distributed::ArgDistAttr& dist_attr) {
  PADDLE_ENFORCE_EQ(
      paddle::holds_alternative<phi::distributed::TensorDistAttr>(dist_attr),
      true,
      common::errors::PreconditionNotMet(
          "Arg must be a single TensorDistAttr"));
  if (out) {
    if (out->impl() == nullptr) {
      auto dist_t = std::make_shared<phi::distributed::DistTensor>(
          phi::DDim(), paddle::get<0>(dist_attr));
      out->set_impl(dist_t);
    }
    return static_cast<phi::distributed::DistTensor*>(out->impl().get());
  }
  return nullptr;
}

std::vector<phi::distributed::DistTensor*> SetKernelDistOutput(
    size_t out_size, std::vector<Tensor>* out) {