%0 = "mhlo.if"(%arg0) ({
    // CHECK: [[SEND0_TOKEN:%.*]] = "mhlo.send"([[ARG1]], [[INIT_TOKEN]])
    "tf.XlaSendToHost"(%arg1) {key = "send_key0"} : (tensor<f32>) -> ()

    // CHECK: [[SEND1_TOKEN:%.*]] = "mhlo.send"([[ARG1]], [[SEND0_TOKEN]])
    "tf.XlaSendToHost"(%arg1) {key = "send_key1"} : (tensor<f32>) -> ()

    // CHECK: mhlo.return [[ARG1]], [[SEND1_TOKEN]]
    "mhlo.return"(%arg1) : (tensor<f32>) -> ()
  },  {

  "tf.XlaSendToHost"(%0) {key = "send_key"} : (tensor<i32>) -> ()

  // CHECK: _xla_original_oc_node_name = [[NODE_NAME3:.*]], _xla_token_input_nodes = {{\[}}[[NODE_NAME2]]{{\]}}
  %1 = "tf._XlaHostComputeMlir"(%0) {recv_key = "host_compute_channel_recv1", send_key = "host_compute_channel_send1", host_mlir_module = ""} : (tensor<i32>) -> (tensor<f32>)
  func.return
}

      %control = tf_executor.island wraps "tf.XlaHostCompute"() {_xla_original_oc_node_name = "Add", _xla_token_input_nodes = ["_xla_token_arg_node"], ancestors = [], cost_estimate_ns = 1024 : i64, key = "host_compute_channel_1_retvals", send_key = "", recv_key = "", shapes = [], tpu_core = 0 : i64} : () -> ()
      tf_executor.fetch
    }
    func.return
  }
}

// CHECK: op: "XlaHostCompute"

// CHECK:       attr {

          %control_4 = tf_executor.island wraps "tf._XlaHostComputeMlir"(%outputs_2) {host_mlir_module = "", manual_sharding = true, recv_key = "host_compute_channel_0_retvals", send_key = "host_compute_channel_0_args"} : (tensor<1x2xf32>) -> ()

    flow ops when necessary.

    For example, given the program

    ```mlir
      func @send_to_host(%arg0: tensor<i32>) {
        "tf.XlaSendToHost"(%arg0) {key = "send_key"} : (tensor<i32>) -> ()
        return
      }
    ```

    This might be legalized like this:

    ```mlir
      func @send_to_host(%arg0: tensor<i32>) {

            {"Toutputs", absl::Span<const DataType>({DT_FLOAT})},
            {"ancestors", absl::Span<const string>({})},
            {"key", "host_compute_channel_F1_F1_O1"},
            {"send_key", ""},
            {"recv_key", ""},
            {"shape_inference_graph", shape_inference_graph},
            {"tpu_core", 0},
            {"cost_estimate_ns", 1000000},

    return
    }

    // CHECK:            "tf._XlaHostComputeMlir"(%[[A_OUTPUT]])
    // CHECK-SAME:       send_key = "host_compute_channel_0_args"
    // CHECK:            %[[C_OUTPUT:[0-9]*]] = "tf.C"
    // CHECK:            "tf._XlaHostComputeMlir"(%[[C_OUTPUT]])
    // CHECK-SAME:       send_key = "host_compute_channel_1_args"
    %1:2 = tf_device.replicate([%0, %arg0] as %ri_0: tensor<2xi32>) {n = 2 : i32} {

        %1 = "tf_device.cluster"() ( {
          %2 = "tf.Const"() {value = dense<1.000000e+00> : tensor<f32>} : () -> tensor<f32>
          %3 = "tf._XlaHostComputeMlir"(%2) {recv_key = "host_compute_channel_0_0_retvals", send_key = "host_compute_channel_0_0_args", tpu_core = 0 : i64} : (tensor<f32>) -> tensor<f32>
          %4 = "tf.AddV2"(%2, %3) : (tensor<f32>, tensor<f32>) -> tensor<f32>
          tf_device.return %4 : tensor<f32>

    // CHECK-SAME:_xla_original_oc_node_name = "hcb0", _xla_token_input_nodes = ["_xla_token_arg_node"] 
    "tf.XlaHostCompute"(%cst) <{ancestors = [], cost_estimate_ns = 1000000 : i64, key = "_host_callback", recv_key = "", send_key = "", shapes = [], tpu_core = 0 : i64}> {_xla_original_oc_node_name = "hcb0", _xla_token_input_nodes = ["_xla_token_arg_node"]} : (tensor<i32>) -> ()
    tf_device.return