Method InferNdSbp

oneflow/user/ops/image_preprocess_ops.cpp:183–214

static */

}

/* static */ Maybe<void> CoinFlipOp::InferNdSbp(user_op::InferNdSbpFnContext* ctx) {
  const Shape& hierarchy = ctx->parallel_hierarchy();
  NdSbp* output_dist = ctx->NdSbp4ArgNameAndIndex("out", 0);
  // the input may be produced by tick which should be broadcast parallel dist
  std::vector<NdSbp*> inputs_dist;
  for (const auto& arg_pair : ctx->inputs()) {
    inputs_dist.emplace_back(ctx->NdSbp4ArgNameAndIndex(arg_pair.first, arg_pair.second));
  }
  const auto& dist_conf = ctx->user_op_conf().attr<std::vector<std::string>>("nd_sbp");
  if (dist_conf.size() == 0) {
    FOR_RANGE(int, i, 0, hierarchy.NumAxes()) {
      output_dist->add_sbp_parallel()->mutable_split_parallel()->set_axis(0);
      for (auto* input_dist : inputs_dist) {
        input_dist->add_sbp_parallel()->mutable_broadcast_parallel();
      }
    }
  } else {
    CHECK_EQ_OR_RETURN(dist_conf.size(), hierarchy.NumAxes());
    for (const std::string& sbp_str : dist_conf) {
      SbpParallel sbp_parallel;
      CHECK_OR_RETURN(ParseSbpParallelFromString(sbp_str, &sbp_parallel));
      CHECK_OR_RETURN(
          (sbp_parallel.has_split_parallel() && sbp_parallel.split_parallel().axis() == 0)
          || sbp_parallel.has_broadcast_parallel());
      *output_dist->add_sbp_parallel() = sbp_parallel;
      for (auto* input_dist : inputs_dist) {
        input_dist->add_sbp_parallel()->mutable_broadcast_parallel();
      }
    }
  }
  return Maybe<void>::Ok();
}

/* static */ Maybe<void> CoinFlipOp::InferDataType(user_op::InferContext* ctx) {
  user_op::TensorDesc* out_tensor = ctx->MutOutputTensorDesc("out", 0);