Method CompileFunctionalNode

tensorflow/compiler/tf2xla/graph_compiler.cc:229–334

}  // namespace

Status GraphCompiler::CompileFunctionalNode(Node* n,
                                            OpKernelContext* op_context) {
  TF_RET_CHECK(IsFunctionCall(*flib_->GetFunctionLibraryDefinition(), *n));
  // For functional nodes, compile them using compiler from the context and call
  // into the functions.
  XlaOpKernelContext xla_op_context(op_context);

  XlaContext& context = XlaContext::Get(op_context);
  auto* b = context.builder();

  XlaCompiler* compiler = xla_op_context.compiler();

  NameAttrList func;
  TF_RETURN_IF_ERROR(GetFunctionNameAndAttr(*flib_, *n, &func));

  std::vector<const XlaExpression*> expressions;

  for (auto tensor : tensor_inputs_) {
    auto expression =
        reinterpret_cast<const XlaExpression*>(tensor->tensor_data().data());
    expressions.push_back(expression);
  }

  // Prepare the arguments and compile the function.
  std::vector<XlaCompiler::Argument> arguments;
  const FunctionBody* fbody;
  TF_RETURN_IF_ERROR(compiler->FindFunctionBody(func, &fbody));

  auto graph = compiler->GetGraph(fbody);

  TF_RETURN_IF_ERROR(
      PrepareArguments(&xla_op_context, graph.get(), expressions, &arguments));

  bool add_token_input_output =
      func.attr().find(kXlaTokenInputNodesAttrName) != func.attr().end();

  XlaCompiler::CompileOptions compile_options;
  compile_options.is_entry_computation = false;
  compile_options.add_token_input_output = add_token_input_output;
  XlaCompiler::CompilationResult result;
  TF_RETURN_IF_ERROR(
      compiler->CompileFunction(compile_options, func, arguments, &result));

  TF_RET_CHECK(arguments.size() == expressions.size());

  std::vector<xla::XlaOp> handles;
  for (int64 i = 0; i < expressions.size(); ++i) {
    if (arguments[i].kind == XlaCompiler::Argument::kConstant) {
      continue;
    }
    if (arguments[i].kind == XlaCompiler::Argument::kResource) {
      handles.push_back(expressions[i]->resource()->value());
    } else {
      handles.push_back(expressions[i]->handle());
    }
  }
  if (add_token_input_output) {
    std::vector<string> token_input_nodes;