已合并
fix(op_host): 10 个算子补齐缺失的 InferDataType 注册 #9616
zl_hw创建于 8月31日
fix(op_host): 10 个算子补齐缺失的 InferDataType 注册 #9616
已合并
共 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 ops | 42 | } // 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 ops | 94 | } // 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 ops | 103 | } // 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 ops | 94 | } // namespace ops |
Moptim/lamb_apply_optimizer_assign/tests/ut/op_host/test_lamb_apply_optimizer_assign_infershape.cpp+48-0
| @@ -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 ops | 65 | } // 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 ops | 64 | } // 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 ops | 66 | } // 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 ops | 59 | } // 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 ops | 57 | } // 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 ops | 57 | } // 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 | +} | ||