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

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