| 320 | } |
| 321 | |
| 322 | Status ShapeRefiner::SetShape(const Node* node, int output_port, |
| 323 | ShapeHandle shape) { |
| 324 | auto c = GetContext(node); |
| 325 | if (c == nullptr) { |
| 326 | return errors::Internal("Could not find context for ", node->name()); |
| 327 | } |
| 328 | |
| 329 | if (output_port < 0 || output_port >= node->num_outputs()) { |
| 330 | return errors::InvalidArgument( |
| 331 | "output_port '", output_port, "' is out of range, ", "node '", |
| 332 | node->name(), "' has ", node->num_outputs(), " outputs"); |
| 333 | } |
| 334 | // Note: it's possible, if the node's been updated, that the shape inference |
| 335 | // context doesn't have the right number of outputs. |
| 336 | if (node->num_outputs() > c->num_outputs()) { |
| 337 | TF_RETURN_IF_ERROR(c->ExpandOutputs(node->num_outputs())); |
| 338 | } |
| 339 | |
| 340 | // Check compatibility, and merge the shapes. |
| 341 | ShapeHandle existing_shape = c->output(output_port); |
| 342 | TF_RETURN_IF_ERROR(c->Merge(existing_shape, shape, &shape)); |
| 343 | c->set_output(output_port, shape); |
| 344 | |
| 345 | // TODO(vrv): Do we need to propagate the new shape through all |
| 346 | // consumers that change their outputs? At the moment, python |
| 347 | // does not do this, but this seems like a nice feature. |
| 348 | |
| 349 | // TODO(vrv): We might need to keep track of the fact that the |
| 350 | // existing shape is invalidated, in case we need to propagate |
| 351 | // this information to remote workers. |
| 352 | return Status::OK(); |
| 353 | } |
| 354 | |
| 355 | Status ShapeRefiner::UpdateNode(const Node* node, bool relax, bool* refined) { |
| 356 | auto it = node_to_context_.find(node); |