已合并
fix(op_host): 10 个算子补齐缺失的 InferDataType 注册 #9616
fix(op_host): 10 个算子补齐缺失的 InferDataType 注册 #9616
已合并
zl_hw创建于 8月31日
共 20 个文件变更+519-10
@@ -25,6 +25,18 @@ static ge::graphStatus InferShape4Relu6Grad(gert::InferShapeContext* context)
25 return Ops::Base::InferShape4Broadcast(context);25 return Ops::Base::InferShape4Broadcast(context);
26}26}
27 27 
28-IMPL_OP_INFERSHAPE(Relu6Grad).InferShape(InferShape4Relu6Grad);28+static ge::graphStatus InferDataType4Relu6Grad(gert::InferDataTypeContext* context)
29+{
30+ if (context == nullptr) {
31+ return ge::GRAPH_FAILED;
32+ }
33+ OP_LOGD(context->GetNodeName(), "InferDataType4Relu6Grad enter");
34+ // backprops 与 gradients 同类型
35+ context->SetOutputDataType(0, context->GetInputDataType(0));
36+ OP_LOGD(context->GetNodeName(), "InferDataType4Relu6Grad end");
37+ return ge::GRAPH_SUCCESS;
38+}
39+ 
40+IMPL_OP_INFERSHAPE(Relu6Grad).InferShape(InferShape4Relu6Grad).InferDataType(InferDataType4Relu6Grad);
29 41 
30} // namespace ops42} // namespace ops
@@ -337,3 +337,28 @@ TEST_F(relu6grad, relu6grad_infer_shape_dynamic)
337 337 
338 ASSERT_EQ(inferShapeFunc(holder.GetContext<gert::InferShapeContext>()), ge::GRAPH_SUCCESS);338 ASSERT_EQ(inferShapeFunc(holder.GetContext<gert::InferShapeContext>()), ge::GRAPH_SUCCESS);
339}339}
340+ 
341+TEST_F(relu6grad, relu6grad_infer_datatype)
342+{
343+ ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl("Relu6Grad"), nullptr);
344+ auto dataTypeFunc = gert::OpImplRegistry::GetInstance().GetOpImpl("Relu6Grad")->infer_datatype;
345+ ASSERT_NE(dataTypeFunc, nullptr);
346+ // backprops 跟随 gradients,覆盖 def 中声明的 fp16/fp32/bf16
347+ for (auto dtype : {ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16}) {
348+ ge::DataType in0 = dtype;
349+ ge::DataType in1 = dtype;
350+ ge::DataType expOut0 = ge::DT_UNDEFINED;
351+ auto contextHolder = gert::InferDataTypeContextFaker()
352+ .NodeIoNum(2, 1)
353+ .NodeInputTd(0, dtype, ge::FORMAT_ND, ge::FORMAT_ND)
354+ .NodeInputTd(1, dtype, ge::FORMAT_ND, ge::FORMAT_ND)
355+ .NodeOutputTd(0, dtype, ge::FORMAT_ND, ge::FORMAT_ND)
356+ .InputDataTypes({&in0, &in1})
357+ .OutputDataTypes({&expOut0})
358+ .Build();
359+ auto context = contextHolder.GetContext<gert::InferDataTypeContext>();
360+ ASSERT_NE(context, nullptr);
361+ EXPECT_EQ(dataTypeFunc(context), ge::GRAPH_SUCCESS);
362+ EXPECT_EQ(context->GetOutputDataType(0), dtype);
363+ }
364+}
@@ -74,5 +74,21 @@ static ge::graphStatus InferShapeMultilabelMarginLoss(gert::InferShapeContext* c
74 return GRAPH_SUCCESS;74 return GRAPH_SUCCESS;
75}75}
76 76 
77-IMPL_OP_INFERSHAPE(MultilabelMarginLoss).InferShape(InferShapeMultilabelMarginLoss);77+static ge::graphStatus InferDataType4MultilabelMarginLoss(gert::InferDataTypeContext* context)
78+{
79+ if (context == nullptr) {
80+ return ge::GRAPH_FAILED;
81+ }
82+ OP_LOGD(context->GetNodeName(), "InferDataType4MultilabelMarginLoss enter");
83+ // y 与 x 同类型
84+ context->SetOutputDataType(0, context->GetInputDataType(0));
85+ // is_target 为 INT32,与 target 同类型(不能跟随 x)
86+ context->SetOutputDataType(1, context->GetInputDataType(1));
87+ OP_LOGD(context->GetNodeName(), "InferDataType4MultilabelMarginLoss end");
88+ return ge::GRAPH_SUCCESS;
89+}
90+ 
91+IMPL_OP_INFERSHAPE(MultilabelMarginLoss)
92+ .InferShape(InferShapeMultilabelMarginLoss)
93+ .InferDataType(InferDataType4MultilabelMarginLoss);
78} // namespace ops94} // namespace ops
@@ -115,3 +115,28 @@ TEST_F(MultilabelMarginLossInferShapeTest, unknown_rank_minus2)
115 EXPECT_EQ(ShapeToVec(y), (std::vector<int64_t>{-2}));115 EXPECT_EQ(ShapeToVec(y), (std::vector<int64_t>{-2}));
116 EXPECT_EQ(ShapeToVec(it), (std::vector<int64_t>{-2}));116 EXPECT_EQ(ShapeToVec(it), (std::vector<int64_t>{-2}));
117}117}
118+ 
119+TEST_F(MultilabelMarginLossInferShapeTest, multilabelmarginloss_infer_datatype)
120+{
121+ ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl("MultilabelMarginLoss"), nullptr);
122+ auto dataTypeFunc = gert::OpImplRegistry::GetInstance().GetOpImpl("MultilabelMarginLoss")->infer_datatype;
123+ ASSERT_NE(dataTypeFunc, nullptr);
124+ ge::DataType in0 = ge::DT_FLOAT;
125+ ge::DataType in1 = ge::DT_INT32;
126+ ge::DataType expOut0 = ge::DT_FLOAT;
127+ ge::DataType expOut1 = ge::DT_INT32;
128+ auto contextHolder = gert::InferDataTypeContextFaker()
129+ .NodeIoNum(2, 2)
130+ .NodeInputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
131+ .NodeInputTd(1, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND)
132+ .NodeOutputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
133+ .NodeOutputTd(1, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND)
134+ .InputDataTypes({&in0, &in1})
135+ .OutputDataTypes({&expOut0, &expOut1})
136+ .Build();
137+ auto context = contextHolder.GetContext<gert::InferDataTypeContext>();
138+ ASSERT_NE(context, nullptr);
139+ EXPECT_EQ(dataTypeFunc(context), ge::GRAPH_SUCCESS);
140+ EXPECT_EQ(context->GetOutputDataType(0), ge::DT_FLOAT);
141+ EXPECT_EQ(context->GetOutputDataType(1), ge::DT_INT32);
142+}
@@ -87,5 +87,17 @@ static ge::graphStatus InferShapePoissonNllLoss(gert::InferShapeContext* context
87 return GRAPH_SUCCESS;87 return GRAPH_SUCCESS;
88}88}
89 89 
90-IMPL_OP_INFERSHAPE(PoissonNllLoss).InferShape(InferShapePoissonNllLoss);90+static ge::graphStatus InferDataType4PoissonNllLoss(gert::InferDataTypeContext* context)
91+{
92+ if (context == nullptr) {
93+ return ge::GRAPH_FAILED;
94+ }
95+ OP_LOGD(context->GetNodeName(), "InferDataType4PoissonNllLoss enter");
96+ // loss 与 input_x 同类型
97+ context->SetOutputDataType(0, context->GetInputDataType(0));
98+ OP_LOGD(context->GetNodeName(), "InferDataType4PoissonNllLoss end");
99+ return ge::GRAPH_SUCCESS;
100+}
101+ 
102+IMPL_OP_INFERSHAPE(PoissonNllLoss).InferShape(InferShapePoissonNllLoss).InferDataType(InferDataType4PoissonNllLoss);
91} // namespace ops103} // namespace ops
@@ -151,3 +151,25 @@ TEST_F(PoissonNllLossInfershapeTest, unknown_rank_sum_is_scalar)
151 EXPECT_EQ(InferShapeTest(op, param), ge::GRAPH_SUCCESS);151 EXPECT_EQ(InferShapeTest(op, param), ge::GRAPH_SUCCESS);
152 EXPECT_EQ(op.GetOutputDesc(0).GetShape().GetDimNum(), 0);152 EXPECT_EQ(op.GetOutputDesc(0).GetShape().GetDimNum(), 0);
153}153}
154+ 
155+TEST_F(PoissonNllLossInfershapeTest, poissonnllloss_infer_datatype)
156+{
157+ ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl("PoissonNllLoss"), nullptr);
158+ auto dataTypeFunc = gert::OpImplRegistry::GetInstance().GetOpImpl("PoissonNllLoss")->infer_datatype;
159+ ASSERT_NE(dataTypeFunc, nullptr);
160+ ge::DataType in0 = ge::DT_FLOAT;
161+ ge::DataType in1 = ge::DT_FLOAT;
162+ ge::DataType expOut0 = ge::DT_FLOAT;
163+ auto contextHolder = gert::InferDataTypeContextFaker()
164+ .NodeIoNum(2, 1)
165+ .NodeInputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
166+ .NodeInputTd(1, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
167+ .NodeOutputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
168+ .InputDataTypes({&in0, &in1})
169+ .OutputDataTypes({&expOut0})
170+ .Build();
171+ auto context = contextHolder.GetContext<gert::InferDataTypeContext>();
172+ ASSERT_NE(context, nullptr);
173+ EXPECT_EQ(dataTypeFunc(context), ge::GRAPH_SUCCESS);
174+ EXPECT_EQ(context->GetOutputDataType(0), ge::DT_FLOAT);
175+}
@@ -72,5 +72,23 @@ static ge::graphStatus InferShape4LambApplyOptimizerAssign(gert::InferShapeConte
72 72 
73 return GRAPH_SUCCESS;73 return GRAPH_SUCCESS;
74}74}
75-IMPL_OP_INFERSHAPE(LambApplyOptimizerAssign).InferShape(InferShape4LambApplyOptimizerAssign);75+static ge::graphStatus InferDataType4LambApplyOptimizerAssign(gert::InferDataTypeContext* context)
76+{
77+ if (context == nullptr) {
78+ return ge::GRAPH_FAILED;
79+ }
80+ OP_LOGD(context->GetNodeName(), "InferDataType4LambApplyOptimizerAssign enter");
81+ // output0 与 grad 同类型
82+ context->SetOutputDataType(0, context->GetInputDataType(0));
83+ // inputv 为同名 ref 输出
84+ context->SetOutputDataType(1, context->GetInputDataType(1));
85+ // inputm 为同名 ref 输出
86+ context->SetOutputDataType(2, context->GetInputDataType(2));
87+ OP_LOGD(context->GetNodeName(), "InferDataType4LambApplyOptimizerAssign end");
88+ return ge::GRAPH_SUCCESS;
89+}
90+ 
91+IMPL_OP_INFERSHAPE(LambApplyOptimizerAssign)
92+ .InferShape(InferShape4LambApplyOptimizerAssign)
93+ .InferDataType(InferDataType4LambApplyOptimizerAssign);
76} // namespace ops94} // namespace ops
@@ -179,3 +179,51 @@ TEST_F(LambApplyOptimizerAssignProtoTest, inputv_inputm_mismatch_is_rejected)
179 gert::Shape o0 = {}, o1 = {}, o2 = {};179 gert::Shape o0 = {}, o1 = {}, o2 = {};
180 ASSERT_EQ(RunInferShape({512, 1024}, {512, 1024}, {1, 1024}, {512, 1024}, &o0, &o1, &o2), ge::GRAPH_FAILED);180 ASSERT_EQ(RunInferShape({512, 1024}, {512, 1024}, {1, 1024}, {512, 1024}, &o0, &o1, &o2), ge::GRAPH_FAILED);
181}181}
182+ 
183+TEST_F(LambApplyOptimizerAssignProtoTest, lambapplyoptimizerassign_infer_datatype)
184+{
185+ ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl("LambApplyOptimizerAssign"), nullptr);
186+ auto dataTypeFunc = gert::OpImplRegistry::GetInstance().GetOpImpl("LambApplyOptimizerAssign")->infer_datatype;
187+ ASSERT_NE(dataTypeFunc, nullptr);
188+ ge::DataType in0 = ge::DT_FLOAT;
189+ ge::DataType in1 = ge::DT_FLOAT;
190+ ge::DataType in2 = ge::DT_FLOAT;
191+ ge::DataType in3 = ge::DT_FLOAT;
192+ ge::DataType in4 = ge::DT_FLOAT;
193+ ge::DataType in5 = ge::DT_FLOAT;
194+ ge::DataType in6 = ge::DT_FLOAT;
195+ ge::DataType in7 = ge::DT_FLOAT;
196+ ge::DataType in8 = ge::DT_FLOAT;
197+ ge::DataType in9 = ge::DT_FLOAT;
198+ ge::DataType in10 = ge::DT_FLOAT;
199+ ge::DataType in11 = ge::DT_FLOAT;
200+ ge::DataType expOut0 = ge::DT_FLOAT;
201+ ge::DataType expOut1 = ge::DT_FLOAT;
202+ ge::DataType expOut2 = ge::DT_FLOAT;
203+ auto contextHolder = gert::InferDataTypeContextFaker()
204+ .NodeIoNum(12, 3)
205+ .NodeInputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
206+ .NodeInputTd(1, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
207+ .NodeInputTd(2, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
208+ .NodeInputTd(3, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
209+ .NodeInputTd(4, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
210+ .NodeInputTd(5, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
211+ .NodeInputTd(6, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
212+ .NodeInputTd(7, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
213+ .NodeInputTd(8, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
214+ .NodeInputTd(9, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
215+ .NodeInputTd(10, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
216+ .NodeInputTd(11, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
217+ .NodeOutputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
218+ .NodeOutputTd(1, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
219+ .NodeOutputTd(2, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
220+ .InputDataTypes({&in0, &in1, &in2, &in3, &in4, &in5, &in6, &in7, &in8, &in9, &in10, &in11})
221+ .OutputDataTypes({&expOut0, &expOut1, &expOut2})
222+ .Build();
223+ auto context = contextHolder.GetContext<gert::InferDataTypeContext>();
224+ ASSERT_NE(context, nullptr);
225+ EXPECT_EQ(dataTypeFunc(context), ge::GRAPH_SUCCESS);
226+ EXPECT_EQ(context->GetOutputDataType(0), ge::DT_FLOAT);
227+ EXPECT_EQ(context->GetOutputDataType(1), ge::DT_FLOAT);
228+ EXPECT_EQ(context->GetOutputDataType(2), ge::DT_FLOAT);
229+}
@@ -47,5 +47,19 @@ static ge::graphStatus InferShape4LambApplyWeightAssign(gert::InferShapeContext*
47 47 
48 return GRAPH_SUCCESS;48 return GRAPH_SUCCESS;
49}49}
50-IMPL_OP_INFERSHAPE(LambApplyWeightAssign).InferShape(InferShape4LambApplyWeightAssign);50+static ge::graphStatus InferDataType4LambApplyWeightAssign(gert::InferDataTypeContext* context)
51+{
52+ if (context == nullptr) {
53+ return ge::GRAPH_FAILED;
54+ }
55+ OP_LOGD(context->GetNodeName(), "InferDataType4LambApplyWeightAssign enter");
56+ // input_param 为同名 ref 输出,取第 5 个输入
57+ context->SetOutputDataType(0, context->GetInputDataType(4));
58+ OP_LOGD(context->GetNodeName(), "InferDataType4LambApplyWeightAssign end");
59+ return ge::GRAPH_SUCCESS;
60+}
61+ 
62+IMPL_OP_INFERSHAPE(LambApplyWeightAssign)
63+ .InferShape(InferShape4LambApplyWeightAssign)
64+ .InferDataType(InferDataType4LambApplyWeightAssign);
51} // namespace ops65} // namespace ops
@@ -74,3 +74,31 @@ TEST_F(LambApplyWeightAssignProtoTest, lamb_apply_weight_assign_case_fp16_4d)
74 auto od0 = holder.GetContext<gert::InferShapeContext>()->GetOutputShape(0);74 auto od0 = holder.GetContext<gert::InferShapeContext>()->GetOutputShape(0);
75 ASSERT_EQ(Ops::Base::ToString(*od0), Ops::Base::ToString(expShape));75 ASSERT_EQ(Ops::Base::ToString(*od0), Ops::Base::ToString(expShape));
76}76}
77+ 
78+TEST_F(LambApplyWeightAssignProtoTest, lambapplyweightassign_infer_datatype)
79+{
80+ ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl("LambApplyWeightAssign"), nullptr);
81+ auto dataTypeFunc = gert::OpImplRegistry::GetInstance().GetOpImpl("LambApplyWeightAssign")->infer_datatype;
82+ ASSERT_NE(dataTypeFunc, nullptr);
83+ ge::DataType in0 = ge::DT_FLOAT;
84+ ge::DataType in1 = ge::DT_FLOAT;
85+ ge::DataType in2 = ge::DT_FLOAT;
86+ ge::DataType in3 = ge::DT_FLOAT;
87+ ge::DataType in4 = ge::DT_FLOAT;
88+ ge::DataType expOut0 = ge::DT_FLOAT;
89+ auto contextHolder = gert::InferDataTypeContextFaker()
90+ .NodeIoNum(5, 1)
91+ .NodeInputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
92+ .NodeInputTd(1, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
93+ .NodeInputTd(2, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
94+ .NodeInputTd(3, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
95+ .NodeInputTd(4, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
96+ .NodeOutputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
97+ .InputDataTypes({&in0, &in1, &in2, &in3, &in4})
98+ .OutputDataTypes({&expOut0})
99+ .Build();
100+ auto context = contextHolder.GetContext<gert::InferDataTypeContext>();
101+ ASSERT_NE(context, nullptr);
102+ EXPECT_EQ(dataTypeFunc(context), ge::GRAPH_SUCCESS);
103+ EXPECT_EQ(context->GetOutputDataType(0), ge::DT_FLOAT);
104+}
@@ -42,5 +42,23 @@ static ge::graphStatus InferShape4LambNextMV(gert::InferShapeContext* context)
42 }42 }
43 return GRAPH_SUCCESS;43 return GRAPH_SUCCESS;
44}44}
45-IMPL_OP_INFERSHAPE(LambNextMV).InferShape(InferShape4LambNextMV);45+static ge::graphStatus InferDataType4LambNextMV(gert::InferDataTypeContext* context)
46+{
47+ if (context == nullptr) {
48+ return ge::GRAPH_FAILED;
49+ }
50+ OP_LOGD(context->GetNodeName(), "InferDataType4LambNextMV enter");
51+ // y1 与 input_mul3 同类型
52+ context->SetOutputDataType(0, context->GetInputDataType(0));
53+ // y2 与 input_mul3 同类型
54+ context->SetOutputDataType(1, context->GetInputDataType(0));
55+ // y3 与 input_mul3 同类型
56+ context->SetOutputDataType(2, context->GetInputDataType(0));
57+ // y4 与 input_mul3 同类型
58+ context->SetOutputDataType(3, context->GetInputDataType(0));
59+ OP_LOGD(context->GetNodeName(), "InferDataType4LambNextMV end");
60+ return ge::GRAPH_SUCCESS;
61+}
62+ 
63+IMPL_OP_INFERSHAPE(LambNextMV).InferShape(InferShape4LambNextMV).InferDataType(InferDataType4LambNextMV);
46} // namespace ops64} // namespace ops
@@ -110,3 +110,57 @@ TEST_F(LambNextMVProtoTest, lamb_next_mv_case_fp16_4d)
110 auto od3 = holder.GetContext<gert::InferShapeContext>()->GetOutputShape(3);110 auto od3 = holder.GetContext<gert::InferShapeContext>()->GetOutputShape(3);
111 ASSERT_EQ(Ops::Base::ToString(*od3), Ops::Base::ToString(expShape));111 ASSERT_EQ(Ops::Base::ToString(*od3), Ops::Base::ToString(expShape));
112}112}
113+ 
114+TEST_F(LambNextMVProtoTest, lambnextmv_infer_datatype)
115+{
116+ ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl("LambNextMV"), nullptr);
117+ auto dataTypeFunc = gert::OpImplRegistry::GetInstance().GetOpImpl("LambNextMV")->infer_datatype;
118+ ASSERT_NE(dataTypeFunc, nullptr);
119+ ge::DataType in0 = ge::DT_FLOAT;
120+ ge::DataType in1 = ge::DT_FLOAT;
121+ ge::DataType in2 = ge::DT_FLOAT;
122+ ge::DataType in3 = ge::DT_FLOAT;
123+ ge::DataType in4 = ge::DT_FLOAT;
124+ ge::DataType in5 = ge::DT_FLOAT;
125+ ge::DataType in6 = ge::DT_FLOAT;
126+ ge::DataType in7 = ge::DT_FLOAT;
127+ ge::DataType in8 = ge::DT_FLOAT;
128+ ge::DataType in9 = ge::DT_FLOAT;
129+ ge::DataType in10 = ge::DT_FLOAT;
130+ ge::DataType in11 = ge::DT_FLOAT;
131+ ge::DataType in12 = ge::DT_FLOAT;
132+ ge::DataType expOut0 = ge::DT_FLOAT;
133+ ge::DataType expOut1 = ge::DT_FLOAT;
134+ ge::DataType expOut2 = ge::DT_FLOAT;
135+ ge::DataType expOut3 = ge::DT_FLOAT;
136+ auto contextHolder = gert::InferDataTypeContextFaker()
137+ .NodeIoNum(13, 4)
138+ .NodeInputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
139+ .NodeInputTd(1, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
140+ .NodeInputTd(2, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
141+ .NodeInputTd(3, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
142+ .NodeInputTd(4, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
143+ .NodeInputTd(5, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
144+ .NodeInputTd(6, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
145+ .NodeInputTd(7, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
146+ .NodeInputTd(8, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
147+ .NodeInputTd(9, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
148+ .NodeInputTd(10, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
149+ .NodeInputTd(11, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
150+ .NodeInputTd(12, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
151+ .NodeOutputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
152+ .NodeOutputTd(1, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
153+ .NodeOutputTd(2, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
154+ .NodeOutputTd(3, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
155+ .InputDataTypes(
156+ {&in0, &in1, &in2, &in3, &in4, &in5, &in6, &in7, &in8, &in9, &in10, &in11, &in12})
157+ .OutputDataTypes({&expOut0, &expOut1, &expOut2, &expOut3})
158+ .Build();
159+ auto context = contextHolder.GetContext<gert::InferDataTypeContext>();
160+ ASSERT_NE(context, nullptr);
161+ EXPECT_EQ(dataTypeFunc(context), ge::GRAPH_SUCCESS);
162+ EXPECT_EQ(context->GetOutputDataType(0), ge::DT_FLOAT);
163+ EXPECT_EQ(context->GetOutputDataType(1), ge::DT_FLOAT);
164+ EXPECT_EQ(context->GetOutputDataType(2), ge::DT_FLOAT);
165+ EXPECT_EQ(context->GetOutputDataType(3), ge::DT_FLOAT);
166+}
@@ -42,5 +42,25 @@ static ge::graphStatus InferShape4LambNextMVWithDecay(gert::InferShapeContext* c
42 }42 }
43 return GRAPH_SUCCESS;43 return GRAPH_SUCCESS;
44}44}
45-IMPL_OP_INFERSHAPE(LambNextMVWithDecay).InferShape(InferShape4LambNextMVWithDecay);45+static ge::graphStatus InferDataType4LambNextMVWithDecay(gert::InferDataTypeContext* context)
46+{
47+ if (context == nullptr) {
48+ return ge::GRAPH_FAILED;
49+ }
50+ OP_LOGD(context->GetNodeName(), "InferDataType4LambNextMVWithDecay enter");
51+ // y1 与 input_mul3 同类型
52+ context->SetOutputDataType(0, context->GetInputDataType(0));
53+ // y2 与 input_mul3 同类型
54+ context->SetOutputDataType(1, context->GetInputDataType(0));
55+ // y3 与 input_mul3 同类型
56+ context->SetOutputDataType(2, context->GetInputDataType(0));
57+ // y4 与 input_mul3 同类型
58+ context->SetOutputDataType(3, context->GetInputDataType(0));
59+ OP_LOGD(context->GetNodeName(), "InferDataType4LambNextMVWithDecay end");
60+ return ge::GRAPH_SUCCESS;
61+}
62+ 
63+IMPL_OP_INFERSHAPE(LambNextMVWithDecay)
64+ .InferShape(InferShape4LambNextMVWithDecay)
65+ .InferDataType(InferDataType4LambNextMVWithDecay);
46} // namespace ops66} // namespace ops
@@ -110,3 +110,57 @@ TEST_F(LambNextMVWithDecayProtoTest, lamb_next_mv_with_decay_case_fp16_4d)
110 auto od3 = holder.GetContext<gert::InferShapeContext>()->GetOutputShape(3);110 auto od3 = holder.GetContext<gert::InferShapeContext>()->GetOutputShape(3);
111 ASSERT_EQ(Ops::Base::ToString(*od3), Ops::Base::ToString(expShape));111 ASSERT_EQ(Ops::Base::ToString(*od3), Ops::Base::ToString(expShape));
112}112}
113+ 
114+TEST_F(LambNextMVWithDecayProtoTest, lambnextmvwithdecay_infer_datatype)
115+{
116+ ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl("LambNextMVWithDecay"), nullptr);
117+ auto dataTypeFunc = gert::OpImplRegistry::GetInstance().GetOpImpl("LambNextMVWithDecay")->infer_datatype;
118+ ASSERT_NE(dataTypeFunc, nullptr);
119+ ge::DataType in0 = ge::DT_FLOAT;
120+ ge::DataType in1 = ge::DT_FLOAT;
121+ ge::DataType in2 = ge::DT_FLOAT;
122+ ge::DataType in3 = ge::DT_FLOAT;
123+ ge::DataType in4 = ge::DT_FLOAT;
124+ ge::DataType in5 = ge::DT_FLOAT;
125+ ge::DataType in6 = ge::DT_FLOAT;
126+ ge::DataType in7 = ge::DT_FLOAT;
127+ ge::DataType in8 = ge::DT_FLOAT;
128+ ge::DataType in9 = ge::DT_FLOAT;
129+ ge::DataType in10 = ge::DT_FLOAT;
130+ ge::DataType in11 = ge::DT_FLOAT;
131+ ge::DataType in12 = ge::DT_FLOAT;
132+ ge::DataType expOut0 = ge::DT_FLOAT;
133+ ge::DataType expOut1 = ge::DT_FLOAT;
134+ ge::DataType expOut2 = ge::DT_FLOAT;
135+ ge::DataType expOut3 = ge::DT_FLOAT;
136+ auto contextHolder = gert::InferDataTypeContextFaker()
137+ .NodeIoNum(13, 4)
138+ .NodeInputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
139+ .NodeInputTd(1, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
140+ .NodeInputTd(2, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
141+ .NodeInputTd(3, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
142+ .NodeInputTd(4, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
143+ .NodeInputTd(5, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
144+ .NodeInputTd(6, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
145+ .NodeInputTd(7, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
146+ .NodeInputTd(8, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
147+ .NodeInputTd(9, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
148+ .NodeInputTd(10, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
149+ .NodeInputTd(11, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
150+ .NodeInputTd(12, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
151+ .NodeOutputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
152+ .NodeOutputTd(1, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
153+ .NodeOutputTd(2, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
154+ .NodeOutputTd(3, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
155+ .InputDataTypes(
156+ {&in0, &in1, &in2, &in3, &in4, &in5, &in6, &in7, &in8, &in9, &in10, &in11, &in12})
157+ .OutputDataTypes({&expOut0, &expOut1, &expOut2, &expOut3})
158+ .Build();
159+ auto context = contextHolder.GetContext<gert::InferDataTypeContext>();
160+ ASSERT_NE(context, nullptr);
161+ EXPECT_EQ(dataTypeFunc(context), ge::GRAPH_SUCCESS);
162+ EXPECT_EQ(context->GetOutputDataType(0), ge::DT_FLOAT);
163+ EXPECT_EQ(context->GetOutputDataType(1), ge::DT_FLOAT);
164+ EXPECT_EQ(context->GetOutputDataType(2), ge::DT_FLOAT);
165+ EXPECT_EQ(context->GetOutputDataType(3), ge::DT_FLOAT);
166+}
@@ -41,5 +41,19 @@ static ge::graphStatus InferShape4LambNextRight(gert::InferShapeContext* context
41 }41 }
42 return GRAPH_SUCCESS;42 return GRAPH_SUCCESS;
43}43}
44-IMPL_OP_INFERSHAPE(LambNextRight).InferShape(InferShape4LambNextRight);44+static ge::graphStatus InferDataType4LambNextRight(gert::InferDataTypeContext* context)
45+{
46+ if (context == nullptr) {
47+ return ge::GRAPH_FAILED;
48+ }
49+ OP_LOGD(context->GetNodeName(), "InferDataType4LambNextRight enter");
50+ // y1 与 input_square 同类型
51+ context->SetOutputDataType(0, context->GetInputDataType(0));
52+ // y2 与 input_square 同类型
53+ context->SetOutputDataType(1, context->GetInputDataType(0));
54+ OP_LOGD(context->GetNodeName(), "InferDataType4LambNextRight end");
55+ return ge::GRAPH_SUCCESS;
56+}
57+ 
58+IMPL_OP_INFERSHAPE(LambNextRight).InferShape(InferShape4LambNextRight).InferDataType(InferDataType4LambNextRight);
45} // namespace ops59} // namespace ops
@@ -82,3 +82,36 @@ TEST_F(LambNextRightProtoTest, lamb_next_right_case_fp16_4d)
82 auto od1 = holder.GetContext<gert::InferShapeContext>()->GetOutputShape(1);82 auto od1 = holder.GetContext<gert::InferShapeContext>()->GetOutputShape(1);
83 ASSERT_EQ(Ops::Base::ToString(*od1), Ops::Base::ToString(expShape));83 ASSERT_EQ(Ops::Base::ToString(*od1), Ops::Base::ToString(expShape));
84}84}
85+ 
86+TEST_F(LambNextRightProtoTest, lambnextright_infer_datatype)
87+{
88+ ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl("LambNextRight"), nullptr);
89+ auto dataTypeFunc = gert::OpImplRegistry::GetInstance().GetOpImpl("LambNextRight")->infer_datatype;
90+ ASSERT_NE(dataTypeFunc, nullptr);
91+ ge::DataType in0 = ge::DT_FLOAT;
92+ ge::DataType in1 = ge::DT_FLOAT;
93+ ge::DataType in2 = ge::DT_FLOAT;
94+ ge::DataType in3 = ge::DT_FLOAT;
95+ ge::DataType in4 = ge::DT_FLOAT;
96+ ge::DataType in5 = ge::DT_FLOAT;
97+ ge::DataType expOut0 = ge::DT_FLOAT;
98+ ge::DataType expOut1 = ge::DT_FLOAT;
99+ auto contextHolder = gert::InferDataTypeContextFaker()
100+ .NodeIoNum(6, 2)
101+ .NodeInputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
102+ .NodeInputTd(1, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
103+ .NodeInputTd(2, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
104+ .NodeInputTd(3, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
105+ .NodeInputTd(4, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
106+ .NodeInputTd(5, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
107+ .NodeOutputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
108+ .NodeOutputTd(1, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
109+ .InputDataTypes({&in0, &in1, &in2, &in3, &in4, &in5})
110+ .OutputDataTypes({&expOut0, &expOut1})
111+ .Build();
112+ auto context = contextHolder.GetContext<gert::InferDataTypeContext>();
113+ ASSERT_NE(context, nullptr);
114+ EXPECT_EQ(dataTypeFunc(context), ge::GRAPH_SUCCESS);
115+ EXPECT_EQ(context->GetOutputDataType(0), ge::DT_FLOAT);
116+ EXPECT_EQ(context->GetOutputDataType(1), ge::DT_FLOAT);
117+}
@@ -39,5 +39,19 @@ static ge::graphStatus InferShape4LambUpdateWithLr(gert::InferShapeContext* cont
39 return ge::GRAPH_FAILED);39 return ge::GRAPH_FAILED);
40 return GRAPH_SUCCESS;40 return GRAPH_SUCCESS;
41}41}
42-IMPL_OP_INFERSHAPE(LambUpdateWithLr).InferShape(InferShape4LambUpdateWithLr);42+static ge::graphStatus InferDataType4LambUpdateWithLr(gert::InferDataTypeContext* context)
43+{
44+ if (context == nullptr) {
45+ return ge::GRAPH_FAILED;
46+ }
47+ OP_LOGD(context->GetNodeName(), "InferDataType4LambUpdateWithLr enter");
48+ // y 与 input_greater1 同类型
49+ context->SetOutputDataType(0, context->GetInputDataType(0));
50+ OP_LOGD(context->GetNodeName(), "InferDataType4LambUpdateWithLr end");
51+ return ge::GRAPH_SUCCESS;
52+}
53+ 
54+IMPL_OP_INFERSHAPE(LambUpdateWithLr)
55+ .InferShape(InferShape4LambUpdateWithLr)
56+ .InferDataType(InferDataType4LambUpdateWithLr);
43} // namespace ops57} // namespace ops
@@ -84,3 +84,39 @@ TEST_F(LambUpdateWithLrProtoTest, lamb_update_with_lr_case_fp16_4d)
84 auto od0 = holder.GetContext<gert::InferShapeContext>()->GetOutputShape(0);84 auto od0 = holder.GetContext<gert::InferShapeContext>()->GetOutputShape(0);
85 ASSERT_EQ(Ops::Base::ToString(*od0), Ops::Base::ToString(expShape));85 ASSERT_EQ(Ops::Base::ToString(*od0), Ops::Base::ToString(expShape));
86}86}
87+ 
88+TEST_F(LambUpdateWithLrProtoTest, lambupdatewithlr_infer_datatype)
89+{
90+ ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl("LambUpdateWithLr"), nullptr);
91+ auto dataTypeFunc = gert::OpImplRegistry::GetInstance().GetOpImpl("LambUpdateWithLr")->infer_datatype;
92+ ASSERT_NE(dataTypeFunc, nullptr);
93+ ge::DataType in0 = ge::DT_FLOAT;
94+ ge::DataType in1 = ge::DT_FLOAT;
95+ ge::DataType in2 = ge::DT_FLOAT;
96+ ge::DataType in3 = ge::DT_FLOAT;
97+ ge::DataType in4 = ge::DT_FLOAT;
98+ ge::DataType in5 = ge::DT_FLOAT;
99+ ge::DataType in6 = ge::DT_FLOAT;
100+ ge::DataType in7 = ge::DT_FLOAT;
101+ ge::DataType in8 = ge::DT_FLOAT;
102+ ge::DataType expOut0 = ge::DT_FLOAT;
103+ auto contextHolder = gert::InferDataTypeContextFaker()
104+ .NodeIoNum(9, 1)
105+ .NodeInputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
106+ .NodeInputTd(1, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
107+ .NodeInputTd(2, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
108+ .NodeInputTd(3, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
109+ .NodeInputTd(4, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
110+ .NodeInputTd(5, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
111+ .NodeInputTd(6, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
112+ .NodeInputTd(7, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
113+ .NodeInputTd(8, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
114+ .NodeOutputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
115+ .InputDataTypes({&in0, &in1, &in2, &in3, &in4, &in5, &in6, &in7, &in8})
116+ .OutputDataTypes({&expOut0})
117+ .Build();
118+ auto context = contextHolder.GetContext<gert::InferDataTypeContext>();
119+ ASSERT_NE(context, nullptr);
120+ EXPECT_EQ(dataTypeFunc(context), ge::GRAPH_SUCCESS);
121+ EXPECT_EQ(context->GetOutputDataType(0), ge::DT_FLOAT);
122+}
@@ -39,5 +39,19 @@ static ge::graphStatus InferShape4LambUpdateWithLrV2(gert::InferShapeContext* co
39 return ge::GRAPH_FAILED);39 return ge::GRAPH_FAILED);
40 return GRAPH_SUCCESS;40 return GRAPH_SUCCESS;
41}41}
42-IMPL_OP_INFERSHAPE(LambUpdateWithLrV2).InferShape(InferShape4LambUpdateWithLrV2);42+static ge::graphStatus InferDataType4LambUpdateWithLrV2(gert::InferDataTypeContext* context)
43+{
44+ if (context == nullptr) {
45+ return ge::GRAPH_FAILED;
46+ }
47+ OP_LOGD(context->GetNodeName(), "InferDataType4LambUpdateWithLrV2 enter");
48+ // y 与 x1 同类型
49+ context->SetOutputDataType(0, context->GetInputDataType(0));
50+ OP_LOGD(context->GetNodeName(), "InferDataType4LambUpdateWithLrV2 end");
51+ return ge::GRAPH_SUCCESS;
52+}
53+ 
54+IMPL_OP_INFERSHAPE(LambUpdateWithLrV2)
55+ .InferShape(InferShape4LambUpdateWithLrV2)
56+ .InferDataType(InferDataType4LambUpdateWithLrV2);
43} // namespace ops57} // namespace ops
@@ -78,3 +78,35 @@ TEST_F(LambUpdateWithLrV2ProtoTest, lamb_update_with_lr_v2_case_fp16_4d)
78 auto od0 = holder.GetContext<gert::InferShapeContext>()->GetOutputShape(0);78 auto od0 = holder.GetContext<gert::InferShapeContext>()->GetOutputShape(0);
79 ASSERT_EQ(Ops::Base::ToString(*od0), Ops::Base::ToString(expShape));79 ASSERT_EQ(Ops::Base::ToString(*od0), Ops::Base::ToString(expShape));
80}80}
81+ 
82+TEST_F(LambUpdateWithLrV2ProtoTest, lambupdatewithlrv2_infer_datatype)
83+{
84+ ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl("LambUpdateWithLrV2"), nullptr);
85+ auto dataTypeFunc = gert::OpImplRegistry::GetInstance().GetOpImpl("LambUpdateWithLrV2")->infer_datatype;
86+ ASSERT_NE(dataTypeFunc, nullptr);
87+ ge::DataType in0 = ge::DT_FLOAT;
88+ ge::DataType in1 = ge::DT_FLOAT;
89+ ge::DataType in2 = ge::DT_FLOAT;
90+ ge::DataType in3 = ge::DT_FLOAT;
91+ ge::DataType in4 = ge::DT_FLOAT;
92+ ge::DataType in5 = ge::DT_FLOAT;
93+ ge::DataType in6 = ge::DT_FLOAT;
94+ ge::DataType expOut0 = ge::DT_FLOAT;
95+ auto contextHolder = gert::InferDataTypeContextFaker()
96+ .NodeIoNum(7, 1)
97+ .NodeInputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
98+ .NodeInputTd(1, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
99+ .NodeInputTd(2, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
100+ .NodeInputTd(3, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
101+ .NodeInputTd(4, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
102+ .NodeInputTd(5, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
103+ .NodeInputTd(6, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
104+ .NodeOutputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
105+ .InputDataTypes({&in0, &in1, &in2, &in3, &in4, &in5, &in6})
106+ .OutputDataTypes({&expOut0})
107+ .Build();
108+ auto context = contextHolder.GetContext<gert::InferDataTypeContext>();
109+ ASSERT_NE(context, nullptr);
110+ EXPECT_EQ(dataTypeFunc(context), ge::GRAPH_SUCCESS);
111+ EXPECT_EQ(context->GetOutputDataType(0), ge::DT_FLOAT);
112+}