A helper function to infer blob's physical shape with ND SBP.
| 105 | |
| 106 | // A helper function to infer blob's physical shape with ND SBP. |
| 107 | Maybe<void> InferPhysicalBlobDesc( |
| 108 | const Operator& op, const PbRpf<std::string>& bns, |
| 109 | const std::function<Maybe<const BlobDesc>(const std::string&)>& GetLogicalBlobDesc, |
| 110 | const NdSbpSignature* nd_sbp_signature, const ParallelContext* parallel_ctx, |
| 111 | const std::function<BlobDesc*(const std::string&)>& GetPhysicalBlobDesc) { |
| 112 | const std::shared_ptr<const ParallelDesc> op_parallel_desc = JUST(op.GetOpParallelDesc()); |
| 113 | for (const auto& bn : bns) { |
| 114 | BlobDesc* physical_blob_desc = GetPhysicalBlobDesc(bn); |
| 115 | const auto& logical_blob_desc = *JUST(GetLogicalBlobDesc(bn)); |
| 116 | CHECK_NOTNULL_OR_RETURN(physical_blob_desc) |
| 117 | << "physical_blob_desc should not be nullptr. op location: " << op.op_loc(); |
| 118 | *physical_blob_desc = logical_blob_desc; |
| 119 | const auto& physical_shape = JUST_MSG( |
| 120 | GetPhysicalShape(logical_blob_desc.shape(), nd_sbp_signature->bn_in_op2nd_sbp().at(bn), |
| 121 | *op_parallel_desc, *parallel_ctx), |
| 122 | std::stringstream() << " check physical shape failed, op name " << op.op_loc()); |
| 123 | physical_blob_desc->set_shape(*physical_shape); |
| 124 | } |
| 125 | return Maybe<void>::Ok(); |
| 126 | } |
| 127 | |
| 128 | } // namespace |
| 129 |
no test coverage detected