| 16 | class DiceFusionTest : public GrapplerTest {}; |
| 17 | |
| 18 | TEST_F(DiceFusionTest, DiceFusionOp) { |
| 19 | using test::function::NDef; |
| 20 | |
| 21 | DiceFusion optimizer; |
| 22 | |
| 23 | const Tensor input = test::AsTensor<float>({1.1, 0.9, 1.2, 0.8}, {2, 2}); |
| 24 | const Tensor sub_1_in = test::AsTensor<float>({1.2, 0.8}); |
| 25 | const Tensor sub_2_in = test::AsTensor<float>({1.0}); |
| 26 | const Tensor mul_1_in = test::AsTensor<float>({1.2, 0.8}); |
| 27 | const Tensor mul_3_in = test::AsTensor<float>({1.2, 0.8}); |
| 28 | |
| 29 | const auto scalar = PartialTensorShape({-1, 2}); |
| 30 | |
| 31 | GrapplerItem item; |
| 32 | item.graph = test::function::GDef({ |
| 33 | |
| 34 | NDef("input", "Const", {}, {{"T", DT_FLOAT}, {"value", input}}), |
| 35 | NDef("sub_1_in", "Const", {}, {{"T", DT_FLOAT}, {"value", sub_1_in}}), |
| 36 | NDef("sub_2_in", "Const", {}, {{"T", DT_FLOAT}, {"value", sub_2_in}}), |
| 37 | NDef("mul_1_in", "Const", {}, {{"T", DT_FLOAT}, {"value", mul_1_in}}), |
| 38 | NDef("mul_3_in", "Const", {}, {{"T", DT_FLOAT}, {"value", mul_3_in}}), |
| 39 | |
| 40 | NDef("sub_1", "Sub", {"input", "sub_1_in"}, |
| 41 | {{"T", DT_FLOAT}, {"_output_shapes", scalar}}), |
| 42 | NDef("mul_1", "Mul", {"sub_1", "mul_1_in"}, |
| 43 | {{"T", DT_FLOAT}, {"_output_shapes", scalar}}), |
| 44 | NDef("sigmoid", "Sigmoid", {"mul_1"}, |
| 45 | {{"T", DT_FLOAT}, {"_output_shapes", scalar}}), |
| 46 | NDef("sub_2", "Sub", {"sigmoid", "sub_2_in"}, |
| 47 | {{"T", DT_FLOAT}, {"_output_shapes", scalar}}), |
| 48 | NDef("mul_2", "Mul", {"sub_2", "input"}, |
| 49 | {{"T", DT_FLOAT}, {"_output_shapes", scalar}}), |
| 50 | NDef("mul_3", "Mul", {"mul_2", "mul_3_in"}, |
| 51 | {{"T", DT_FLOAT}, {"_output_shapes", scalar}}), |
| 52 | NDef("mul_4", "Mul", {"sigmoid", "input"}, |
| 53 | {{"T", DT_FLOAT}, {"_output_shapes", scalar}}), |
| 54 | NDef("add", "Add", {"mul_3", "mul_4"}, |
| 55 | {{"T", DT_FLOAT}, {"_output_shapes", scalar}}), |
| 56 | }); |
| 57 | |
| 58 | // Place all nodes on CPU. |
| 59 | for (int i = 0; i < item.graph.node_size(); ++i) { |
| 60 | item.graph.mutable_node(i)->set_device("/device:CPU:0"); |
| 61 | } |
| 62 | |
| 63 | GraphDef output; |
| 64 | TF_EXPECT_OK(optimizer.Optimize(/*cluster=*/nullptr, item, &output)); |
| 65 | |
| 66 | GraphDef expected = test::function::GDef({ |
| 67 | NDef("input", "Const", {}, {{"T", DT_FLOAT}, {"value", input}}, "/device:CPU:0"), |
| 68 | NDef("sub_1_in", "Const", {}, {{"T", DT_FLOAT}, {"value", sub_1_in}}, "/device:CPU:0"), |
| 69 | NDef("mul_1_in", "Const", {}, {{"T", DT_FLOAT}, {"value", mul_1_in}}, "/device:CPU:0"), |
| 70 | NDef("mul_3_in", "Const", {}, {{"T", DT_FLOAT}, {"value", mul_3_in}}, "/device:CPU:0"), |
| 71 | |
| 72 | // fusing node. |
| 73 | NDef("add", "Dice", |
| 74 | {"input", "sub_1_in", "mul_1_in", "mul_3_in"}, |
| 75 | {{"T", DT_FLOAT}}, "/device:CPU:0"), |
nothing calls this directly
no test coverage detected