| 371 | } |
| 372 | |
| 373 | Maybe<void> Operator::InferOutBlobDescs( |
| 374 | const std::function<BlobDesc*(const std::string&)>& GetBlobDesc4BnInOp, |
| 375 | const ParallelContext* parallel_ctx) const { |
| 376 | if (parallel_ctx->parallel_num() == 1) { |
| 377 | JUST(InferLogicalOutBlobDescs(GetBlobDesc4BnInOp, *JUST(GetOpParallelDesc()))); |
| 378 | } else { |
| 379 | const auto& nd_sbp_signature = JUST(this->nd_sbp_signature()); |
| 380 | const auto& parallel_desc = JUST(this->GetOpParallelDesc()); |
| 381 | for (const auto& bn : input_bns()) { |
| 382 | const auto& nd_sbp = nd_sbp_signature->bn_in_op2nd_sbp().at(bn); |
| 383 | std::shared_ptr<const BlobDesc> in_logical = JUST(GetLogicalBlobDesc4Ibn(bn)); |
| 384 | CHECK_OR_RETURN( |
| 385 | *JUST(GetPhysicalShape(in_logical->shape(), nd_sbp, *parallel_desc, *parallel_ctx)) |
| 386 | == GetBlobDesc4BnInOp(bn)->shape()); |
| 387 | } |
| 388 | for (const auto& bn : output_bns()) { |
| 389 | BlobDesc* desc = GetBlobDesc4BnInOp(bn); |
| 390 | *desc = *JUST(GetLogicalBlobDesc4Obn(bn)); |
| 391 | const auto& nd_sbp = nd_sbp_signature->bn_in_op2nd_sbp().at(bn); |
| 392 | desc->set_shape( |
| 393 | *JUST(GetPhysicalShape(desc->shape(), nd_sbp, *parallel_desc, *parallel_ctx))); |
| 394 | } |
| 395 | } |
| 396 | return Maybe<void>::Ok(); |
| 397 | } |
| 398 | |
| 399 | Maybe<void> Operator::InferInternalBlobDescsIf( |
| 400 | const std::function<BlobDesc*(const std::string&)>& GetBlobDesc4BnInOp, |
nothing calls this directly
no test coverage detected