#include "tensorflow/core/framework/cpp_shape_inference.pb.h"
#include "tensorflow/core/framework/full_type.pb.h"
#include "tensorflow/core/framework/kernel_def.pb.h"
#include "tensorflow/core/framework/log_memory.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/partial_tensor_shape.h"
#include "tensorflow/core/framework/shape_inference.h"

  // scheduled on that same threadpool, causing a deadlock in cases where the
  // caller of event_mgr->ThenExecute() blocks on the completion of the callback
  // (as in the case of ConstOp kernel creation on GPU, which involves copying a
  // CPU tensor to GPU).
  // Setting a larger thread pool does not help with the Swift caller, as we use
  // a different TFE context for each thread of execution (for running graph

#include "absl/strings/str_cat.h"
#include "absl/time/time.h"
#include "absl/types/span.h"
#include "tensorflow/c/experimental/next_pluggable_device/tensor_pjrt_buffer_util.h"
#include "tensorflow/c/kernels.h"
#include "tensorflow/c/kernels_experimental.h"
#include "tensorflow/c/tf_buffer.h"
#include "tensorflow/c/tf_status.h"
#include "tensorflow/c/tf_status_internal.h"
#include "tensorflow/c/tf_tensor.h"

  ASSERT_EQ(kernel_list.kernel_size(), 1);
  EXPECT_EQ(kernel_list.kernel(0).op(), "TestOpWithSingleKernel");
  EXPECT_EQ(kernel_list.kernel(0).device_type(), "CPU");
  TF_DeleteBuffer(kernel_list_buf);
  TF_DeleteStatus(status);
}

TEST(TestKernel, TestGetRegisteredKernelsForOpNoKernels) {
  TF_Status* status = TF_NewStatus();
  TF_Buffer* kernel_list_buf = TF_GetRegisteredKernelsForOp("Unknown", status);

    .SetIsDistributedCommunication()
    .Doc(R"doc(Test communication op throwing Unavailable error.)doc");

// Kernel that throws an Unavailable error.
class TestUnavailableErrorOp : public tensorflow::OpKernel {
 public:
  explicit TestUnavailableErrorOp(tensorflow::OpKernelConstruction* ctx)
      : tensorflow::OpKernel(ctx) {}
  void Compute(tensorflow::OpKernelContext* ctx) override {