Function GetSubLogicalShape

oneflow/core/framework/placement_sbp_util.cpp:512–525

}

Maybe<Shape> GetSubLogicalShape(Symbol<one::GlobalTensorMeta> tensor_meta,
                                Symbol<ParallelDesc> sub_parallel_desc, Symbol<NdSbp> sub_nd_sbp) {
  CHECK_EQ_OR_RETURN(sub_nd_sbp->sbp_parallel_size(), 1);  // NOLINT(maybe-need-error-msg)
  const auto& logical_shape = tensor_meta->shape();
  const auto& physical_shape =
      JUST(GetPhysicalShape(logical_shape, tensor_meta->nd_sbp(), tensor_meta->parallel_desc()));

  std::shared_ptr<Shape> sub_logical_shape = std::make_shared<Shape>(*physical_shape);
  if (sub_nd_sbp->sbp_parallel(0).has_split_parallel()) {
    const int64_t split_axis = sub_nd_sbp->sbp_parallel(0).split_parallel().axis();
    sub_logical_shape->Set(split_axis, logical_shape.At(split_axis));
  }
  return sub_logical_shape;
}

Maybe<Symbol<one::GlobalTensorMeta>> CalcSubGlobalTensorMeta(
    Symbol<one::GlobalTensorMeta> tensor_meta, Symbol<ParallelDesc> sub_parallel_desc,