Method InferShapesForFunction

tensorflow/core/common_runtime/shape_refiner.cc:158–235

TODO(cwhipkey): When an inference context inside function has requested_input_tensor(i) or requested_input_tensor_as_partial_shape(i) set when input(i) is an _Arg op, then this request should propagate to context, and vice versa. NOTE: Recursive user-defined functions are not supported. Maybe we won't support recursive functions at all in TF, because of other maintainability issues.

// Maybe we won't support recursive functions at all in TF, because of
// other maintainability issues.
Status ShapeRefiner::InferShapesForFunction(
    const FunctionDef* function_def, AttrSlice attributes,
    bool keep_nested_shapes, ExtendedInferenceContext* outer_context) {
  const Graph* graph;
  auto it = functions_.find(function_def);
  if (it != functions_.end()) {
    graph = it->second.get();
  } else {
    InstantiationResult result;
    TF_RETURN_IF_ERROR(InstantiateFunction(
        *function_def, attributes,
        [this](const string& op, const OpDef** sig) {
          return this->function_library_->LookUpOpDef(op, sig);
        },
        &result));

    Graph* new_graph = new Graph(function_library_);
    GraphConstructorOptions options;
    options.allow_internal_ops = true;
    TF_RETURN_IF_ERROR(
        ConvertNodeDefsToGraph(options, result.nodes, new_graph));
    functions_[function_def].reset(new_graph);
    graph = new_graph;
  }

  std::unordered_set<const Node*> function_nodes;
  Status inference_status = Status::OK();
  {
    auto node_shape_inference_lambda = [this, &outer_context, &function_nodes,
                                        &inference_status](const Node* node) {
      if (!inference_status.ok()) return;
      inference_status = InferShapesForFunctionSubNode(
          node, this, outer_context->get_context());
      function_nodes.insert(node);
    };

    // Calls inference lambda for each node after visiting all predecessors.
    // Ensures that we are adding nodes to ShapeRefiner in the topological
    // order.
    ReverseDFS(*graph, {}, node_shape_inference_lambda);
  }

  if (keep_nested_shapes && inference_status.ok()) {
    // Fill the nested inferences map.
    //
    // The materialized function graph has extra nodes for arguments and
    // return values, which are not explicitly listed in the FunctionDef,
    // we filter out these special nodes here to not expose the implementation
    // details and keep only inferences for the nodes listed in the FunctionDef.
    std::unordered_map<string, const NodeDef*> user_defined_nodes;
    for (const auto& node_def : function_def->node_def()) {
      user_defined_nodes[node_def.name()] = &node_def;
    }

    std::unordered_map<string, std::unique_ptr<ExtendedInferenceContext>>
        nested_inferences;
    for (const Node* node : function_nodes) {
      const string& node_name = node->name();