Build a graph containing a Variable op of "variable_op_type" and an Assign op of "assign_op_type", and expect all of the ops to be placed on a device of type "expected_device_type".
| 889 | // Assign op of "assign_op_type", and expect all of the ops to be |
| 890 | // placed on a device of type "expected_device_type". |
| 891 | Status PlacerTest::ReferenceTestHelper(const string& variable_op_type, |
| 892 | const string& assign_op_type, |
| 893 | const DeviceType& expected_device_type) { |
| 894 | Graph g(OpRegistry::Global()); |
| 895 | { // Scope for temporary variables used to construct g. |
| 896 | GraphDefBuilder b(GraphDefBuilder::kFailImmediately); |
| 897 | Node* input = ops::SourceOp("TestInput", b.opts().WithName("in")); |
| 898 | // Build ten variable-and-assignment pairs. |
| 899 | for (int i = 0; i < 10; ++i) { |
| 900 | Node* var = ops::SourceOp(variable_op_type, |
| 901 | b.opts().WithName(strings::StrCat("var_", i))); |
| 902 | ops::BinaryOp(assign_op_type, var, input, |
| 903 | b.opts().WithName(strings::StrCat("assign_", i))); |
| 904 | } |
| 905 | TF_EXPECT_OK(BuildGraph(b, &g)); |
| 906 | } |
| 907 | |
| 908 | TF_RETURN_IF_ERROR(Place(&g)); |
| 909 | |
| 910 | for (int i = 0; i < 10; ++i) { |
| 911 | EXPECT_COLOCATED(g, strings::StrCat("var_", i), |
| 912 | strings::StrCat("assign_", i)); |
| 913 | EXPECT_DEVICE_TYPE(g, strings::StrCat("var_", i), expected_device_type); |
| 914 | EXPECT_DEVICE_TYPE(g, strings::StrCat("assign_", i), expected_device_type); |
| 915 | } |
| 916 | |
| 917 | return Status::OK(); |
| 918 | } |
| 919 | |
| 920 | // Test all 2^3 combinations of Variable and Assignment op types |
| 921 | // (unconstrained, CPU-only, and GPU-only). |