From 0028bf843d8846bd16b25bf5447b1649fde10fb7 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Tue, 8 May 2018 16:16:57 -0700 Subject: [PATCH] add test for pruning useless function lib in graph. PiperOrigin-RevId: 195893756 --- .../grappler/optimizers/function_optimizer_test.cc | 25 ++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/tensorflow/core/grappler/optimizers/function_optimizer_test.cc b/tensorflow/core/grappler/optimizers/function_optimizer_test.cc index a2dbab3..0aaf57e 100644 --- a/tensorflow/core/grappler/optimizers/function_optimizer_test.cc +++ b/tensorflow/core/grappler/optimizers/function_optimizer_test.cc @@ -835,5 +835,30 @@ TEST_F(FunctionOptimizerTest, SpecializeFunction_OncePerUniqueContext) { test::ExpectTensorEqual(tensors_expected[5], tensors[5]); } +TEST_F(FunctionOptimizerTest, PruningUselessLibraryFunctions) { + using test::function::NDef; + FunctionOptimizer optimizer(RewriterConfig::DEFAULT); + DisableFunctionSpecialization(&optimizer); + auto func = test::function::XTimesTwo(); + (*func.mutable_attr())["_noinline"].set_b(true); + GrapplerItem item; + item.graph = test::function::GDef( + {NDef("x", "Placeholder", {}, {{"dtype", DT_FLOAT}}, "/device:CPU:0"), + NDef("y", "XTimesTwo", {"x"}, {{"T", DT_FLOAT}}, "/device:CPU:0"), + NDef("z", "Identity", {"y"}, {{"T", DT_FLOAT}}, "/device:CPU:0")}, + // FunctionLib + { + func, + test::function::XTimesTwoInt32(), + test::function::XTimes16(), + }); + GraphDef output; + Status status = optimizer.Optimize(nullptr, item, &output); + TF_EXPECT_OK(status); + + EXPECT_EQ(output.library().function().size(), 1); + EXPECT_EQ(output.library().function(0).signature().name(), "XTimesTwo"); +} + } // namespace grappler } // namespace tensorflow -- 2.7.4