Function PropagateConstIntoIfNode

tensorflow/compiler/tf2xla/tf2xla_util.cc:217–243

For an "If" node in graph `g`, if it has Const node inputs, rewrite its then/else branch function to replace _Arg nodes with those Const inputs.

// For an "If" node in graph `g`, if it has Const node inputs, rewrite its
// then/else branch function to replace _Arg nodes with those Const inputs.
Status PropagateConstIntoIfNode(Graph* g, Node* if_node,
                                const FunctionLibraryDefinition* lookup_fld,
                                FunctionLibraryDefinition* fld) {
  // Notice that first input for If node is predicate; other inputs are function
  // inputs.
  std::unordered_map<int, const Node*> const_input_index_to_node;
  for (int i = 1; i < if_node->num_inputs(); i++) {
    const Node* input_node;
    TF_RETURN_IF_ERROR(if_node->input_node(i, &input_node));
    if (input_node->type_string() == "Const") {
      const_input_index_to_node[i - 1] = input_node;
    }
  }
  if (const_input_index_to_node.empty()) {
    return Status::OK();
  }

  // Rewrite "then_branch" and "else_branch" function, replace usage of those
  // _Arg nodes with corresponding const node.
  for (const auto& attr_name :
       std::vector<string>{"then_branch", "else_branch"}) {
    TF_RETURN_IF_ERROR(PropagateConstIntoFuncAttr(
        if_node, attr_name, const_input_index_to_node, lookup_fld, fld));
  }

  return Status::OK();
}

// For a "While" node in graph `g`, if it has Const node inputs, rewrite its
// cond/body function to replace _Arg nodes with those Const inputs.