| 196 | } // namespace |
| 197 | |
| 198 | Status GraphDefBuilderWrapper::AddDataset( |
| 199 | const DatasetBase* dataset, |
| 200 | const std::vector<std::pair<size_t, Node*>>& inputs, |
| 201 | const std::vector<std::pair<size_t, gtl::ArraySlice<Node*>>>& list_inputs, |
| 202 | const std::vector<std::pair<StringPiece, AttrValue>>& attrs, |
| 203 | Node** output) { |
| 204 | const string& type_string = dataset->type_string(); |
| 205 | std::unique_ptr<const GraphDefBuilder::Options> opts( |
| 206 | new GraphDefBuilder::Options(b_->opts())); |
| 207 | // TODO(srbs|mrry): Not all datasets have output_types and output_shapes |
| 208 | // attributes defined. It will be nice to have a consistent pattern. |
| 209 | bool has_output_types_attr = HasAttr(type_string, "output_types"); |
| 210 | bool has_output_shapes_attr = HasAttr(type_string, "output_shapes"); |
| 211 | if (has_output_shapes_attr) { |
| 212 | opts.reset(new GraphDefBuilder::Options( |
| 213 | opts->WithAttr("output_shapes", dataset->output_shapes()))); |
| 214 | } |
| 215 | if (has_output_types_attr) { |
| 216 | opts.reset(new GraphDefBuilder::Options( |
| 217 | opts->WithAttr("output_types", dataset->output_dtypes()))); |
| 218 | } |
| 219 | for (auto attr : attrs) { |
| 220 | opts.reset( |
| 221 | new GraphDefBuilder::Options(opts->WithAttr(attr.first, attr.second))); |
| 222 | } |
| 223 | if (opts->HaveError()) { |
| 224 | return errors::Internal("AddDataset: Failed to build Options with error ", |
| 225 | opts->StatusToString()); |
| 226 | } |
| 227 | NodeBuilder node_builder(opts->GetNameForOp(type_string), type_string, |
| 228 | opts->op_registry()); |
| 229 | { |
| 230 | size_t total_size = inputs.size() + list_inputs.size(); |
| 231 | auto inputs_iter = inputs.begin(); |
| 232 | auto list_inputs_iter = list_inputs.begin(); |
| 233 | for (int i = 0; i < total_size; i++) { |
| 234 | if (inputs_iter != inputs.end() && inputs_iter->first == i) { |
| 235 | node_builder.Input(NodeBuilder::NodeOut(inputs_iter->second)); |
| 236 | inputs_iter++; |
| 237 | } else if (list_inputs_iter != list_inputs.end() && |
| 238 | list_inputs_iter->first == i) { |
| 239 | std::vector<NodeBuilder::NodeOut> nodeout_inputs; |
| 240 | nodeout_inputs.reserve(list_inputs_iter->second.size()); |
| 241 | for (Node* n : list_inputs_iter->second) { |
| 242 | nodeout_inputs.emplace_back(n); |
| 243 | } |
| 244 | node_builder.Input(nodeout_inputs); |
| 245 | list_inputs_iter++; |
| 246 | } else { |
| 247 | return errors::InvalidArgument("No input found for index ", i); |
| 248 | } |
| 249 | } |
| 250 | } |
| 251 | *output = opts->FinalizeBuilder(&node_builder); |
| 252 | if (*output == nullptr) { |
| 253 | return errors::Internal("AddDataset: Failed to build ", type_string, |
| 254 | " op with error ", opts->StatusToString()); |
| 255 | } |