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

        %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

    manual_inputs.push_back(in_manual);
  }

  // Create the _XlaHostComputeMlirOp
  auto host_compute = builder.create<mlir::TF::_XlaHostComputeMlirOp>(
      loc, shard_output_types, manual_inputs,
      /*send_key=*/builder.getStringAttr(args_communication_key),
      /*recv_key=*/builder.getStringAttr(retvals_communication_key),
      /*host_mlir_module=*/builder.getStringAttr(serialized_func_module),

    StrArrayAttr:$ancestors,
    TF_ShapeAttrArray:$shapes,
    OptionalAttr<SymbolRefAttr>:$shape_inference_graph,
    StrAttr:$key,
    DefaultValuedStrAttr<StrAttr, "">:$send_key,
    DefaultValuedStrAttr<StrAttr, "">:$recv_key,
    DefaultValuedOptionalAttr<I64Attr, "1000000">:$cost_estimate_ns,
    DefaultValuedOptionalAttr<I64Attr, "0">:$tpu_core
  );

  let results = (outs

    %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>
  }
}

// -----

// The following xla.OpSharding is used:

A pseudo-op to represent host-side computation in an XLA program.
  }];

  let arguments = (ins
    Arg<Variadic<TF_Tensor>, [{A list of tensors that will be sent to the host.}]>:$inputs,

    StrAttr:$send_key,
    StrAttr:$recv_key,
    DefaultValuedOptionalAttr<StrAttr, "\"\"">:$host_mlir_module,
    DefaultValuedOptionalAttr<BoolAttr, "false">:$manual_sharding
  );

  let results = (outs