| 202 | } |
| 203 | |
| 204 | Status OpTestBuilder::BuildGraph(const string& name_prefix, |
| 205 | const string& device, bool use_jit, |
| 206 | GraphDef* graphdef, NodeDef** test_node_def, |
| 207 | std::vector<string>* inputs, |
| 208 | std::vector<string>* outputs) const { |
| 209 | OpRegistryInterface* op_registry = OpRegistry::Global(); |
| 210 | |
| 211 | const OpDef* op_def; |
| 212 | TF_RETURN_IF_ERROR(op_registry->LookUpOpDef(node_def_.op(), &op_def)); |
| 213 | |
| 214 | NodeDef* test_def = graphdef->add_node(); |
| 215 | *test_def = node_def_; |
| 216 | test_def->set_name(absl::StrCat(name_prefix, "_op_under_test")); |
| 217 | test_def->set_device(device); |
| 218 | AddDefaultsToNodeDef(*op_def, test_def); |
| 219 | if (use_jit) { |
| 220 | AddNodeAttr(kXlaCompileAttr, true, test_def); |
| 221 | } |
| 222 | VLOG(1) << "Op under test: " << test_def->DebugString(); |
| 223 | |
| 224 | DataTypeVector input_types, output_types; |
| 225 | TF_RETURN_IF_ERROR( |
| 226 | InOutTypesForNode(*test_def, *op_def, &input_types, &output_types)); |
| 227 | |
| 228 | // Build feed and fetch nodes. |
| 229 | for (int i = 0; i < input_types.size(); ++i) { |
| 230 | NodeDef* def = graphdef->add_node(); |
| 231 | string name = absl::StrCat(name_prefix, "_input_", i); |
| 232 | TF_RETURN_IF_ERROR(NodeDefBuilder(name, "Placeholder") |
| 233 | .Device(device) |
| 234 | .Attr("dtype", input_types[i]) |
| 235 | .Finalize(def)); |
| 236 | inputs->push_back(name); |
| 237 | test_def->add_input(name); |
| 238 | } |
| 239 | |
| 240 | for (int i = 0; i < output_types.size(); ++i) { |
| 241 | NodeDef* def = graphdef->add_node(); |
| 242 | string name = absl::StrCat(name_prefix, "_output_", i); |
| 243 | TF_RETURN_IF_ERROR(NodeDefBuilder(name, "Identity") |
| 244 | .Device(device) |
| 245 | .Attr("T", output_types[i]) |
| 246 | .Input(test_def->name(), i, output_types[i]) |
| 247 | .Finalize(def)); |
| 248 | outputs->push_back(name); |
| 249 | } |
| 250 | |
| 251 | if (test_node_def) { |
| 252 | *test_node_def = test_def; |
| 253 | } |
| 254 | |
| 255 | return Status::OK(); |
| 256 | } |
| 257 | |
| 258 | // Test fixture. The fixture manages the random number generator and its seed, |
| 259 | // and has a number of convenience methods for building random Tensors, shapes, |
no test coverage detected