}

// Builds XlaRecvAtHost node, and replaces all _Arg nodes with it.
absl::StatusOr<Node*> ReplaceArgNodesWithRecvAtHostNode(
    Graph* g, const string& oc_cluster_name,
    std::vector<DataType>* recv_at_host_dtypes, Node* key_placeholder) {
  // TODO(b/77601805): use out nodes for source node, instead of traversing all
  // nodes.
  std::vector<Node*> arg_nodes = GatherNodesWithType(*g, "_Arg");

// lifted out of If/While/function node. Attribute value will always be boolean
// value "true".
extern const char kXlaIsLiftedArgAttrName[];

// Attribute indicating that this node is a Placeholder node for an _Arg node
// lifted out of If/While/function node. Attribute value will be a string, which
// is the outside compilation cluster name sending the lifted arg node to host.
extern const char kXlaLiftedArgOutsideCompilationAttrName[];

    int GetResultIndexForEdge(const Edge* edge) const;

    // Creates an _Arg node for the src node of edge, and add its index to
    // args_by_src_, if none exists yet. Also adds its index to args_by_dst_,
    // and adds the edge within the subgraph from the _Arg node to the image of
    // the dst node.
    Status RecordArg(const Edge* edge,

// TODO(phawkins): this code does not infer accurate shapes for cyclic graphs.
// `arg_shapes`: user given map from the `index` to shapes of this
// node, where `index` is the `index` attribute of `_Arg` op or `_index`
// attribute of `Placeholder` op.
Status InferShapes(Graph* graph, const std::map<int, InferredShape>& arg_shapes,
                   const tensorflow::FunctionLibraryDefinition* fnlib_def,

  graph.SetConstructionContext(ConstructionContext::kEagerRuntime);
  tensorflow::set_tf2_execution(true);

  ConfigProto config = ConfigProto();
  Scope root = Scope::NewRootScope().ExitOnError();

  Output a = ops::_Arg(root.WithOpName("A"), DT_RESOURCE, 0);
  std::vector<NodeBuilder::NodeOut> inputs({NodeBuilder::NodeOut(a.node())});

  Node* call;
  NameAttrList f_name_attr;
  f_name_attr.set_name(fd.signature().name());
  TF_ASSERT_OK(

                << context->DebugString(handle);
      }
    }

    int index = -1;
    if (n->type_string() == "_Arg") {
      // NOTE: during runtime, Placeholder ops will be replaced as `_Arg` ops.
      // And Args must have `index` attribute.
      TF_RETURN_IF_ERROR(GetNodeAttr(n->attrs(), "index", &index));
    } else if (n->type_string() == "Placeholder") {

  %graph = tf_executor.graph {
    tf_executor.fetch %arg0 : tensor<10xi32>
  }
  func.return %graph : tensor<10xi32>
}

// CHECK:      node {
// CHECK-NEXT:   name: "input0"
// CHECK-NEXT:   op: "_Arg"
// CHECK:          key: "T"
// CHECK-NEXT:     value {
// CHECK-NEXT:       type: DT_INT32
// CHECK-NEXT:     }
// CHECK:          key: "_output_shapes"
// CHECK-NEXT:     value {
// CHECK-NEXT:       list {

  TF_EXPECT_OK(identity_node->input_node(0, &identity_input_node));
  EXPECT_EQ(identity_input_node->name(), "_arg0");

  const Node* arg_node = node_name_index.at("_arg0");
  EXPECT_EQ(arg_node->op_def().name(), "_Arg");
  EXPECT_EQ(arg_node->attrs().FindByString("T")->type(), DT_FLOAT);

  const Node* retval_node = node_name_index.at("_retval0");
  EXPECT_EQ(retval_node->op_def().name(), "_Retval");

      count: 5
    }
    op_histogram {
      op: "Mul"
      count: 4
    }
  }
  unclustered_op_histogram {
    op: "NoOp"
    count: 2
  }
  unclustered_op_histogram {
    op: "_Arg"
    count: 1
  }
  unclustered_op_histogram {
    op: "_Retval"
    count: 1
  }
}
)",
      &expected_auto_clustering_activity);
  EXPECT_EQ(listener()->auto_clustering_activity().DebugString(),

    LogNotCompilable(node, uncompilable_reason);
    return false;
  }

  // _Arg nodes in a top-level function represent feeds and _Retval nodes in a
  // top-level function represent fetches.
  if (stack_depth == 1 &&
      (node.type_string() == "_Arg" || node.type_string() == "_Retval")) {
    absl::string_view uncompilable_reason = "top level _Arg or _Retval";
    MaybeMarkUncompilableNode(uncompilable_reason, *stack_trace,