// host-device communication.
// CHECK: tf._TPUCompileMlir
// CHECK-SAME: tf.Rank
// CHECK-NOT: tf._XlaHostComputeMlir
// CHECK-NOT: tf._XlaRecvAtHost
// CHECK-NOT: tf._XlaSendFromHost

  // CHECK-NEXT: "tf._XlaSendFromHost"(%arg0, %0, [[key]]) <{device_ordinal = 0 : i64, key = "host_compute_channel_0_retvals"}> {_xla_has_host_transfer = true, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<*xi32>, tensor<*xi32>, tensor<3x!tf_type.string>) ->...

  // 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"
  // CHECK-SAME: [[ID_OUTPUT]]
  // CHECK-SAME: key = "host_compute_channel_recv"

                                        &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;
      }
    }
    EXPECT_TRUE(has_recv);

    // CHECK-NOT:        "tf._TPUDeviceOrdinalPlaceholder"
    // CHECK:            %[[B_OUTPUT:[0-9]*]] = "tf.B"()
    // CHECK:            "tf._XlaSendFromHost"(%[[B_OUTPUT]], %[[PROGRAM_OUTPUT]])
    // CHECK-SAME:       device_ordinal = 0
    // CHECK-SAME:       device_type = "TPU"
    // CHECK-SAME:       key = "host_compute_channel_0_retvals"

  for (Node* n : g->nodes()) {
    if (n->type_string() == "_XlaSendFromHost") {
      send_from_host = n;
      break;
    }
  }
  if (!send_from_host) {
    return errors::Internal("Shape inference graph ",
                            shape_inference_graph_name,
                            " does not have _XlaSendFromHost node.");
  }

  // See if the SendFromHost node exists in `host_graph`.

          %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>
          "tf._XlaSendFromHost"(%3, %1) {device_ordinal = 0 : i64, key = "host_compute_channel_0_0_retvals"} : (tensor<f32>, tensor<3x!tf_type.string>) -> ()
          tf_device.return
        }) {device = "/job:worker/replica:0/task:0/device:CPU:0"} : () -> ()

}

Operation* ApplyXlaHostTransferAttr(Operation* op, OpBuilder& builder) {
  op->setAttr("_xla_has_host_transfer", builder.getBoolAttr(true));
  return op;
}

// Creates a tf._XlaSendFromHost or tf._XlaSendFromHostV2 op. If device ordinal
// is present, a tf._XlaSendFromHostV2 op is created instead.
Operation* CreateSendFromHostOp(OpBuilder& builder, Location loc,

  string name = absl::StrCat("outside_compilation_", cluster, "_",
                             new_func_name, "_", oc_cluster, "_send");
  NodeBuilder node_builder(opts.WithName(name).GetNameForOp("_XlaSendFromHost"),
                           "_XlaSendFromHost", opts.op_registry());
  node_builder.Input(inputs);
  node_builder.Input(std::move(key_input));
  std::vector<DataType> dtypes;
  for (const auto& node : inputs) {

    // CHECK: %5 = "tf.Shape"(%4) {_xla_outside_compilation = "0"} : (tensor<1120x8xi32>) -> tensor<2xi32>
    // CHECK: "tf._XlaSendFromHost"(%3#1, %3#2, %4, %5, %0) <{device_ordinal = 0 : i64, device_type = "TPU", key = "host_compute_channel_0_retvals"}> : (tensor<1x1120x8xi32>, tensor<1x1120x8xi32>, tensor<1120x8xi32>, tensor<2xi32>, tensor<3x!tf_type.string>) -> ()