const NameAttrList& func_name_attrs, Node* node,
                 const std::map<string, int>& host_compute_core)
      : cluster_name(cluster_name),
        func_name_attrs(func_name_attrs),
        node(node),
        host_compute_core(host_compute_core) {}
  // XLA cluster name. It might be different from `func_name`.
  const string cluster_name;

  std::map<string, int> host_compute_core;
  std::vector<string> shape_inference_graphs;
  bool has_outside_compilation;
  NameAttrList name_attrs;
  name_attrs.set_name("cluster");
  *name_attrs.mutable_attr() = attrs;
  TF_CHECK_OK(ExtractOutsideCompilationTest(
      "_xla", "_oc", "cluster", name_attrs, "cluster_rewritten", "host_graph",
      host_compute_core, &fld, &shape_inference_graphs,

func.func @convert_tpu_replicate() {
  tf_executor.graph {

      auto host_compute_node_or = ReplaceOutsideCompilationCallNode(
          graph_out.get(), n, host_compute_core, *cluster_deps);
      TF_RETURN_IF_ERROR(host_compute_node_or.status());
      Node* host_compute_node = host_compute_node_or.value();
      host_compute_nodes[host_compute_node->name()] = host_compute_node;
    }
    // For XlaHostCompute nodes with dependencies, add control edges between

//   because XLA cluster name might be different from `func_name`.
// func_name_attrs: they will be used to instantiate the XLA computation func.
// new_func_name: new function name for rewritten XLA computation func.
// host_compute_core: mapping from outside compilation cluster name to XLA
//   device assignment.
// fld: FunctionLibraryDefinition object.
// host_graph: Graph object to store host side graph for all outside

    %control_4 = tf_executor.island(%control_2) wraps "tf.TPUReplicateMetadata"() {_replication_info = "cluster", _xla_compile_device_type = "TPU", allow_soft_placement = false, computation_shape = [], device = "", device_assignment = [], host_compute_core = [], num_cores_per_replica = 1 : i64, num_replicas = 1 : i64, padding_map = [], step_marker_location = "STEP_MARK_AT_ENTRY", topology = "", use_spmd_for_xla_partitioning = true, use_tpu = true} : () -> ()

"tf_device.cluster_func"(%arg0) <{func = @empty_func}> {_dynamic_arg_index = [], _replication_info = "cluster", _xla_compile_device_type = "TPU", allow_soft_placement = false, computation_shape = [], device = "", device_assignment = [0, 0, 0, 0, 0, 0, 0, 1], host_compute_core = [], input_sharding_configuration = ["{devices=[2,1]0,1}"], num_cores_per_replica = 2 : i64, output_sharding_configuration = [""], padding_map = [], step_marker_location = "STEP_MARK_AT_ENTRY", topology = "\0A\04\02\02\01\02\10\01\18\08\22 ...

    "tf_device.cluster"() ({
      "tf.OpA"(%arg0) : (tensor<f32>) -> ()
      tf_device.return

    "tf.TPUReplicateMetadata"() {_tpu_replicate = "cluster_Fn", allow_soft_placement = true, computation_shape = [], device = "", device_assignment = [], host_compute_core = [], num_cores_per_replica = 1 : i64, num_replicas = 1 : i64, padding_map = [], step_marker_location = "STEP_MARK_AT_ENTRY", topology = "", use_spmd_for_xla_partitioning = true, use_tpu = true} : () -> ()

    %cst_1 = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
    "tf.TPUReplicateMetadata"() {_tpu_replicate = "cluster", allow_soft_placement = false, computation_shape = [], device = "", device_assignment = [], host_compute_core = [], num_cores_per_replica = 1 : i64, num_replicas = 1 : i64, padding_map = [], step_marker_location = "STEP_MARK_AT_ENTRY", topology = "", use_spmd_for_xla_partitioning = true, use_tpu = true} : () -> ()