// . So the TPU program should contain tf.Rank op and there should be no
// host-device communication.
// CHECK: tf._TPUCompileMlir
// CHECK-SAME: tf.Rank
// CHECK-NOT: tf._XlaHostComputeMlir
// CHECK-NOT: tf._XlaRecvAtHost
// CHECK-NOT: tf._XlaSendFromHost

    // CHECK:            %[[RECV_OUTPUT:[0-9]*]] = "tf._XlaRecvAtHost"(%[[PROGRAM_OUTPUT]])
    // CHECK:            _xla_has_host_transfer = true
    // CHECK:            %[[B_OUTPUT:[0-9]*]] = "tf.B"(%[[RECV_OUTPUT]])
    // CHECK:            %[[RECV_OUTPUT:[0-9]*]] = "tf._XlaRecvAtHost"(%[[PROGRAM_OUTPUT]])
    // CHECK:            _xla_has_host_transfer = true

  // CHECK-SAME: shape_inference_graph = @host_func
  // CHECK-SAME: shapes = [#tf_type.shape<*>, #tf_type.shape<3x?>]
  // CHECK-SAME: tpu_core = 0 : i64
  // CHECK: func @host_func
  // CHECK: %[[RECV_OUTPUT:[0-9]*]]:2 = "tf._XlaRecvAtHost"
  // CHECK-SAME: key = "host_compute_channel_send"
  // CHECK-SAME: -> (tensor<*xf32>, tensor<3x?xf64>)
  // CHECK: %[[ID_OUTPUT:[0-9]*]] = "tf.Identity"(%[[RECV_OUTPUT]]#0)
  // CHECK: "tf._XlaSendFromHost"

                                        &call_fbody));

    // Verify we have _XlaRecvAtHost and _XlaSendFromHost nodes.
    bool has_recv = false, has_send = false;
    for (Node *n : call_fbody->graph->nodes()) {
      if (n->type_string() == "_XlaRecvAtHost") {
        has_recv = true;
      } else if (n->type_string() == "_XlaSendFromHost") {
        has_send = true;
      }

    const std::vector<DataType>& recv_at_host_dtypes, Node* key_placeholder) {
  NodeDefBuilder recv_at_host_builder(
      absl::StrCat("outside_compilation_", oc_cluster_name, "_recv"),
      "_XlaRecvAtHost");
  NodeDef recv_at_host_def;
  recv_at_host_builder.Attr("Toutputs", recv_at_host_dtypes);
  // The correct device_ordinal will be inserted during replication in a
  // subsequent rewrite.

      %0 = "tf_device.parallel_execute"() ( {
        "tf_device.launch"() ( {
          %1 = "tf._XlaCompileMlirPlaceholderProgramKey"() : () -> tensor<3x!tf_type.string>
          %2 = "tf._XlaRecvAtHost"(%1) {device_ordinal = 0 : i64, key = "host_compute_channel_0_0_args"} : (tensor<3x!tf_type.string>) -> tensor<f32>
          %3 = "tf.Identity"(%2) : (tensor<f32>) -> tensor<f32>

          /*dynamic_key=*/compilation_key,
          builder.getStringAttr(communication_key),
          /*device_ordinal=*/builder.getI64IntegerAttr(0), device_type_attr),
      builder);
}

// Creates a tf._XlaRecvAtHost or tf._XlaRecvAtHostV2 op. If device ordinal is
// present, a tf._XlaRecvAtHostV2 op is created instead.
Operation* CreateRecvAtHostOp(OpBuilder& builder, Location loc,

  string name = absl::StrCat("outside_compilation_", cluster, "_",
                             new_func_name, "_", oc_cluster, "_recv");
  NodeBuilder node_builder(opts.WithName(name).GetNameForOp("_XlaRecvAtHost"),
                           "_XlaRecvAtHost", opts.op_registry());
  node_builder.Input(std::move(key_input));
  return opts.WithAttr("Toutputs", dtypes)
      .WithAttr("key", key)
      .WithAttr("device_ordinal", 0)

    // CHECK-NEXT: %cst_1 = "tf.Const"() <{value = dense<[1120, -1]> : tensor<2xi32>}> : () -> tensor<2xi32>
    // CHECK: %0 = "tf._XlaCompileMlirPlaceholderProgramKey"() : () -> tensor<3x!tf_type.string>
    // CHECK: %1:4 = "tf._XlaRecvAtHost"(%0) <{device_ordinal = 0 : i64, device_type = "TPU", key = "host_compute_channel_0_args"}> : (tensor<3x!tf_type.string>) -> (tensor<3360x8xi32>, tensor<3xi32>, tensor<i32>, tensor<2xi32>)

func @outside_compilation() -> tensor<f32> {
  %0 = "tf_device.parallel_execute"() ( {
    "tf_device.launch"() ( {
      %1 = "tf._XlaCompileMlirPlaceholderProgramKey"() : () -> tensor<3x!tf_type.string>
      %2 = "tf._XlaRecvAtHost"(%1) {device_ordinal = 0 : i64, key = "host_compute_channel_0_0_args"} : (tensor<3x!tf_type.string>) -> tensor<f32>
      %3 = "tf.Identity"(%2) : (tensor<f32>) -> tensor<f32>