if (!s.ok()) {
    ::tensorflow::Set_TF_Status_from_Status(status, s);
    return nullptr;
  }
  TF_Tensor* tf_tensor = TF_TensorFromTensor(*tensor, &s);
  if (!s.ok()) {
    ::tensorflow::Set_TF_Status_from_Status(status, s);
    return nullptr;
  }
  return tf_tensor;
}

TF_Tensor* TF_ForwardInputOrAllocateOutput(
    TF_OpKernelContext* context, const int* candidate_input_indices,

}

/// Parse a type registered to this dialect.
Type FallbackDialect::parseType(DialectAsmParser &parser) const {
  StringRef keyword;
  if (parser.parseKeyword(&keyword)) return Type();

  if (keyword == "tf_tensor") return TFTensorType::get(getContext());
  if (keyword == "tf_allocator") return TFAllocatorType::get(getContext());

  parser.emitError(parser.getNameLoc(), "unknown type: ") << keyword;
  return Type();
}

    theta_inputs[i] = inputs[i];
  }

  AbstractTensorHandle* theta =
      theta_inputs[input_index];  // parameter we are grad checking

  // Convert from AbstractTensor to TF_Tensor.
  TF_Tensor* theta_tensor;
  TF_RETURN_IF_ERROR(GetValue(theta, &theta_tensor));

  // Get number of elements and fill data.
  int num_elems = TF_TensorElementCount(theta_tensor);
  vector<float> theta_data(num_elems);

// if the tensor does not have a `PjRtCApiBuffer`.
TF_CAPI_EXPORT extern PJRT_Buffer* TF_GetPjRtCBuffer(TF_Tensor* c_tensor,
                                                     TF_Status* status);

// Creates a `PjRtCApiBuffer` with the `PJRT_Buffer*` passed in and set to the
// tensor.
TF_CAPI_EXPORT extern void TF_CreatePjRtBuffer(TF_Tensor* c_tensor,
                                               PJRT_Buffer* c_buffer,

    name = "tf_tensor_hdrs",
    hdrs = [
        "tf_tensor.h",
        "tf_tensor_helper.h",
    ],
    visibility = ["//visibility:public"],
    deps = [
        ":c_api_macros_hdrs",
        ":tf_datatype_hdrs",
        ":tf_status_headers",
        "//tensorflow/core/platform:status",
    ],
)

cc_library(
    name = "tf_tensor",
    srcs = ["tf_tensor.cc"],
    hdrs = [
        "tf_tensor.h",

// RUN: tf-tfrt-opt -tf-executor-to-tfrt-pipeline %s | FileCheck %s --dump-input=fail

// CHECK-LABEL: func @__inference_pruned_131
// CHECK-SAME: ([[in_chain:%.*]]: !tfrt.chain) -> (!tfrt.chain, !tfrt_fallback.tf_tensor)
// CHECK-NEXT: [[o_chain:%.*]], [[o:%.*]] = tfrt_fallback_async.executeop.seq([[in_chain]]) key(0) cost({{.*}}) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.VarHandleOp"()

  // CHECK:       [[pred:%.*]] = tfrt_fallback_async.predicate
  // CHECK-NEXT:  tfrt.while [[pred]] @"[[while_func_prefix:.*]]/tfrt_body_1"
  // CHECK-SAME:  (!tfrt.chain, !tfrt_fallback.tf_tensor) -> (!tfrt.chain, !tfrt_fallback.tf_tensor)
  // CHECK-NOT:   func.call
  func.func @while_test_remove_unused_results(%arg0: tensor<i32>, %arg1: tensor<i32>) -> (tensor<i32>, tensor<i32>) {

    ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_);
    TF_Tensor* expected_out0 = expected_csession->output_tensor(0);
    TF_Tensor* expected_out1 = expected_csession->output_tensor(1);

    CompareTensors(out0, expected_out0);
    CompareTensors(out1, expected_out1);
  }

  void CompareTensors(TF_Tensor* a, TF_Tensor* b) {
    float* a_data = static_cast<float*>(TF_TensorData(a));

#include "tensorflow/c/tf_tensor.h"

// Required for IS_MOBILE_PLATFORM definition
#include "tsl/platform/platform.h"  // IWYU pragma: keep

#if !defined(IS_MOBILE_PLATFORM) && !defined(IS_SLIM_BUILD)
#include "tensorflow/core/common_runtime/next_pluggable_device/c/tf_rendezvous_c_api.h"
#endif

#ifdef __cplusplus
extern "C" {
#endif

typedef struct TF_Tensor TF_Tensor;

// passed to or returned from C functions *by pointer*. Otherwise, changes to
// its internal structure will break the C API's binary interface.
typedef struct TF_Tensor {
  tensorflow::AbstractTensorInterface* tensor;
} TF_Tensor;

class TF_ManagedBuffer : public tensorflow::TensorBuffer {
 public:
  TF_ManagedBuffer(void* data, size_t len,
                   void (*deallocator)(void* data, size_t len, void* arg),