| 353 | } |
| 354 | |
| 355 | Status ShapeRefiner::UpdateNode(const Node* node, bool relax, bool* refined) { |
| 356 | auto it = node_to_context_.find(node); |
| 357 | if (it == node_to_context_.end()) { |
| 358 | *refined = true; |
| 359 | return AddNode(node); |
| 360 | } |
| 361 | ExtendedInferenceContext* node_ext_context = it->second.get(); |
| 362 | InferenceContext* node_context = node_ext_context->get_context(); |
| 363 | |
| 364 | // Give up if the context wasn't successfully built by the AddNode() method. |
| 365 | TF_RETURN_IF_ERROR(node_context->construction_status()); |
| 366 | |
| 367 | // Check if the shapes of the nodes in the fan-in of this node have changed, |
| 368 | // and if they have update the node input shapes. |
| 369 | for (const Edge* e : node->in_edges()) { |
| 370 | if (e->IsControlEdge()) continue; |
| 371 | |
| 372 | int dst_input = e->dst_input(); |
| 373 | int src_output = e->src_output(); |
| 374 | |
| 375 | Node* input = e->src(); |
| 376 | auto iter = node_to_context_.find(input); |
| 377 | if (iter == node_to_context_.end()) { |
| 378 | return errors::FailedPrecondition( |
| 379 | "Input ", dst_input, " ('", input->name(), "') for '", node->name(), |
| 380 | "' was not previously added to ShapeRefiner."); |
| 381 | } |
| 382 | |
| 383 | InferenceContext* c = iter->second->get_context(); |
| 384 | DCHECK_GE(dst_input, 0); |
| 385 | ShapeHandle existing_input = node_context->input(dst_input); |
| 386 | if (!relax) { |
| 387 | if (node_context->MergeInput(dst_input, c->output(src_output))) { |
| 388 | if (!SameDefinedShape(node_context, node_context->input(dst_input), |
| 389 | existing_input)) { |
| 390 | *refined = true; |
| 391 | } |
| 392 | } |
| 393 | } else { |
| 394 | if (node_context->RelaxInput(dst_input, c->output(src_output))) { |
| 395 | if (!SameDefinedShape(node_context, node_context->input(dst_input), |
| 396 | existing_input)) { |
| 397 | *refined = true; |
| 398 | } |
| 399 | } |
| 400 | } |
| 401 | if (node_context->requested_input_tensor_as_partial_shape(dst_input)) { |
| 402 | // The input value may have changed. Since we have no way to know if |
| 403 | // that's indeed the case, err on the safe side. |
| 404 | *refined = true; |
| 405 | } |
| 406 | |
| 407 | // Also propagate handle shape and dtype of edges which are carrying |
| 408 | // resource handles. |
| 409 | if (e->src()->output_type(src_output) == DT_RESOURCE) { |
| 410 | auto* outputs = c->output_handle_shapes_and_types(src_output); |
| 411 | if (!outputs) continue; |
| 412 | |