p27049618
/
mindspore
forked from mindspore-Ecosystem/mindspore
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

!49941 Correct rpc recv abs. 

Merge pull request !49941 from ZPaC/fix-rpc-recv-infer
This commit is contained in:
i-robot 2023-03-08 02:24:56 +00:00 committed by Gitee
parent 4d14ae79bd 81895a3f14
commit 568adc3562
No known key found for this signature in database
GPG Key ID: 173E9B9CA92EEF8F
1 changed files with 11 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
mindspore/core/ops/rpc_recv.cc
View File

@ -20,6 +20,7 @@
#include "abstract/ops/primitive_infer_map.h" #include "abstract/ops/primitive_infer_map.h"
#include "ir/anf.h" #include "ir/anf.h"
#include "ops/core_ops.h" #include "ops/core_ops.h"
#include "ops/op_name.h"
#include "ops/primitive_c.h" #include "ops/primitive_c.h"
#include "utils/log_adapter.h" #include "utils/log_adapter.h"
#include "mindapi/src/helper.h" #include "mindapi/src/helper.h"
@ -29,9 +30,16 @@ namespace ops {
MIND_API_OPERATOR_IMPL(RpcRecv, BaseOperator); MIND_API_OPERATOR_IMPL(RpcRecv, BaseOperator);
AbstractBasePtr RpcRecvInfer(const abstract::AnalysisEnginePtr &, const PrimitivePtr &, AbstractBasePtr RpcRecvInfer(const abstract::AnalysisEnginePtr &, const PrimitivePtr &,
const std::vector<AbstractBasePtr> &input_args) { const std::vector<AbstractBasePtr> &input_args) {
abstract::AbstractTuplePtr rpc_send_abs = std::make_shared<abstract::AbstractTuple>(input_args); if (input_args.empty()) {
MS_EXCEPTION_IF_NULL(rpc_send_abs); MS_LOG(EXCEPTION) << "The input size of RpcRecv is 0.";
return rpc_send_abs; }
if (input_args.size() == static_cast<size_t>(kDim1)) {
return input_args[kInputIndex0];
} else {
abstract::AbstractTuplePtr rpc_recv_abs = std::make_shared<abstract::AbstractTuple>(input_args);
MS_EXCEPTION_IF_NULL(rpc_recv_abs);
return rpc_recv_abs;
}
} }
REGISTER_PRIMITIVE_EVAL_IMPL(RpcRecv, prim::kPrimRpcRecv, RpcRecvInfer, nullptr, true); REGISTER_PRIMITIVE_EVAL_IMPL(RpcRecv, prim::kPrimRpcRecv, RpcRecvInfer, nullptr, true);
} // namespace ops } // namespace ops