return false;
    }
    return true;
  }
  bool HostCallback(Stream* stream,
                    absl::AnyInvocable<absl::Status() &&> callback) override {
    SP_Stream stream_handle = static_cast<CStream*>(stream)->Handle();
    HostCallbackContext* ctx = new HostCallbackContext{std::move(callback)};
    return stream_executor_->host_callback(&device_, stream_handle,

  ASSERT_FALSE(synchronize_all_called);
  ASSERT_TRUE(executor->SynchronizeAllActivity());
  ASSERT_TRUE(synchronize_all_called);
}

TEST_F(StreamExecutorTest, HostCallbackOk) {
  se_.host_callback = [](const SP_Device* const device, SP_Stream stream,
                         SE_StatusCallbackFn const callback_fn,
                         void* const callback_arg) -> TF_Bool {
    TF_Status* status = TF_NewStatus();

    // CHECK:  tf.XlaHostCompute
    // 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