experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

[CostModel] Add TF_ASSERT_OK

Browse Source
This commit is contained in:
Xinan Jiang 2021-02-02 11:56:46 +08:00
parent f71df7a434
commit 175e49a0df
1 changed files with 5 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
tensorflow/core/graph/costmodel_test.cc
View File

@ -24,6 +24,7 @@ limitations under the License.
#include "tensorflow/core/framework/cost_graph.pb.h" #include "tensorflow/core/framework/cost_graph.pb.h"
#include "tensorflow/core/framework/step_stats.pb.h" #include "tensorflow/core/framework/step_stats.pb.h"
#include "tensorflow/core/graph/graph.h" #include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/protobuf.h" #include "tensorflow/core/platform/protobuf.h"
#include "tensorflow/core/platform/test.h" #include "tensorflow/core/platform/test.h"
#include "tensorflow/core/util/dump_graph.h" #include "tensorflow/core/util/dump_graph.h"
@ -83,8 +84,10 @@ TEST(CostModelTest, GlobalId) {
CostModelManager cost_model_manager; CostModelManager cost_model_manager;
collector.BuildCostModel(&cost_model_manager, device_map); collector.BuildCostModel(&cost_model_manager, device_map);
CostGraphDef cost_graph_def; CostGraphDef cost_graph_def;
cost_model_manager.AddToCostGraphDef(graph1.get(), &cost_graph_def); TF_ASSERT_OK(
cost_model_manager.AddToCostGraphDef(graph2.get(), &cost_graph_def); cost_model_manager.AddToCostGraphDef(graph1.get(), &cost_graph_def));
TF_ASSERT_OK(
cost_model_manager.AddToCostGraphDef(graph2.get(), &cost_graph_def));
ASSERT_EQ(cost_graph_def.node_size(), 12); ASSERT_EQ(cost_graph_def.node_size(), 12);
absl::flat_hash_map<int32, const CostGraphDef::Node> ids; absl::flat_hash_map<int32, const CostGraphDef::Node> ids;
for (auto node : cost_graph_def.node()) { for (auto node : cost_graph_def.node()) {