tf_executor.fetch
    }
    func.return
  }
}

// CHECK: op: "XlaHostCompute"

// CHECK:       attr {
// CHECK:         key: "shape_inference_graph"
// CHECK:         value {

            {"key", "host_compute_channel_F1_F1_O1"},
            {"send_key", ""},
            {"recv_key", ""},
            {"shape_inference_graph", shape_inference_graph},
            {"tpu_core", 0},
            {"cost_estimate_ns", 1000000},
            {"shapes", absl::Span<const DataType>({})},
            {"_outside_compilation_subgraph", "O1"},
            {"_xla_token_input_nodes",

  EXPECT_NE(host_compute_1, nullptr);
  // Check XlaHostCompute nodes' "tpu_core" attr.
  int tpu_core;
  TF_CHECK_OK(GetNodeAttr(host_compute_0->attrs(), "tpu_core", &tpu_core));
  EXPECT_EQ(tpu_core, 1);
  TF_CHECK_OK(GetNodeAttr(host_compute_1->attrs(), "tpu_core", &tpu_core));
  EXPECT_EQ(tpu_core, 0);
  // Check XlaHostCompute nodes' "shapes" attr. "0" should not have shapes, and

  // CHECK-SAME: send_key = "host_compute_channel_send"
  // CHECK-SAME: shape_inference_graph = @host_func
  // CHECK-SAME: shapes = [#tf_type.shape<*>, #tf_type.shape<3x?>]
  // CHECK-SAME: tpu_core = 0 : i64
  // CHECK: func @host_func
  // CHECK: %[[RECV_OUTPUT:[0-9]*]]:2 = "tf._XlaRecvAtHost"
  // CHECK-SAME: key = "host_compute_channel_send"
  // CHECK-SAME: -> (tensor<*xf32>, tensor<3x?xf64>)

        /*key=*/rewriter.getStringAttr(""), op.getSendKeyAttr(),
        op.getRecvKeyAttr(),
        /*cost_estimate_ns=*/rewriter.getI64IntegerAttr(kDefaultCostEstimate),
        /*tpu_core=*/rewriter.getI64IntegerAttr(0));
  }
};

void UpdateArgAttributes(mlir::func::FuncOp func) {
  OpBuilder builder(func.getBody());
  for (int i = 0; i < func.getNumArguments(); ++i) {

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

    "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

                            metadata_op, compilation_op, loop_body_func,
                            "sc_forward", &forward_caller);
  if (failed(result)) return signalPassFailure();

  // Create tpu_core function
  builder.setInsertionPointAfter(forward_caller);
  Operation* core_tpu_caller = nullptr;
  result = ExtractOpsAsFunc(builder, module, core_tpu_ops, replication_attr,

  for (const auto& attr : call_node->attrs()) {
    host_compute_builder.Attr(attr.first, attr.second);
  }

  // Populate tpu_core assignment.
  const auto iter = host_compute_core.find(original_oc_name);
  if (iter != host_compute_core.end()) {
    int core = iter->second;
    host_compute_builder.Attr("tpu_core", core);
  }

  // Set input tokens and other outside compilation clusters that current

    // CHECK: "tf._XlaHostComputeMlir"
    // CHECK-SAME: -> tensor<*xf32>
    %0 = "tf._XlaHostComputeMlir"(%arg0) {recv_key = "host_compute_channel_recv", send_key = "host_compute_channel_send", tpu_core = 0, host_mlir_module = ""} : (tensor<2xf32>) -> tensor<*xf32>
    func.return %0 : tensor<*xf32>
  }

  // CHECK-LABEL: func @xla_host_compute_mlir_shape_inferred