%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"} : () -> ()
        tf_device.return
      },  {

      %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"} : () -> ()
    tf_device.return
  },  {

    // CHECK-SAME:       device_type = "TPU"
    // CHECK-SAME:       key = "host_compute_channel_0_retvals"
    // CHECK:          "tf_device.cluster"
    // CHECK:            %[[A_OUTPUT:[0-9]*]] = "tf.A"
    // CHECK:            %[[HOST_OUTPUT:[0-9]*]] = "tf._XlaHostComputeMlir"()
    // CHECK-SAME:       recv_key = "host_compute_channel_0_retvals"
    // CHECK-SAME:       send_key = "host_compute_channel_0_args"

    return
    }

          %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>) -> ()

  // 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-NEXT: return [[exec_result]] : tensor<*xi32>

        // expected-remark@above {{ID: 1}}
        // expected-remark@above {{Successors: {2}}}
        "tf._XlaSendFromHostV2"(%const, %arg0, %arg1) {_xla_has_host_transfer = true, key = "host_compute_channel_0_retvals"} : (tensor<f32>, tensor<!tf_type.string>, tensor<i64>) -> ()

    %2336:4 = "tf.XlaHostCompute"(%2335, %cst_2, %cst_68, %cst_1) <{ancestors = [], cost_estimate_ns = 1000000 : i64, key = "", recv_key = "host_compute_channel_0_retvals", send_key = "host_compute_channel_0_args", shape_inference_graph = @host_func_0, shapes = [#tf_type.shape<1x1120x8>, #tf_type.shape<1x1120x8>, #tf_type.shape<1120x8>, #tf_type.shape<2>], tpu_core = 0 : i64}> {_xla_original_oc_node_name...

  }
  func.func private @_func() -> (tensor<?xi32> {mhlo.sharding = "\08\01\1A\01\01\22\01\00"}) {
    %0 = "tf.A"() : () -> tensor<?xi32>
    "tf._XlaHostComputeMlir"(%0) {host_mlir_module = "", recv_key = "host_compute_channel_0_retvals", send_key = "host_compute_channel_0_args"} : (tensor<?xi32>) -> ()
    %1 = "tf.C"() : () -> tensor<?xi32>
    return %1 : tensor<?xi32>
  }
}

// -----