*graph.add_node() = MakeNode("a", "Placeholder", {},
                               {{"dtype", DT_FLOAT}, {"shape", input_shape}});
  *graph.add_node() = MakeNode("b", "Placeholder", {},
                               {{"dtype", DT_FLOAT}, {"shape", input_shape}});
  *graph.add_node() = MakeNode("c", "Placeholder", {},
                               {{"dtype", DT_FLOAT}, {"shape", input_shape}});

}

Status DemangleTensor(absl::string_view str, TensorProto* proto) {
  return ParseTextProto(str, kTensorPrefix, proto);
}

string MangleDataType(const DataType& dtype) {
  return absl::StrCat(kDataTypePrefix, DataType_Name(dtype));
}

Status DemangleDataType(absl::string_view str, DataType* proto) {
  absl::string_view pbtxt;
  TF_RETURN_IF_ERROR(ConsumePrefix(str, kDataTypePrefix, &pbtxt));

  switch (dtype) {
#define CASE(type)                                       \
  case DataTypeToEnum<type>::value: {                    \
    type* typed_a = static_cast<type*>(a);               \
    type* typed_b = static_cast<type*>(b);               \
    for (int64_t i = 0; i < num_elements; ++i) {         \
      if (DataTypeIsFloating(dtype)) {                   \
        EXPECT_FLOAT_EQ(static_cast<float>(typed_a[i]),  \

#include "tsl/platform/statusor.h"

namespace tensorflow {

static size_t GetTensorSize(const TensorShape& shape, const DataType dtype) {
  return shape.num_elements() * DataTypeSize(dtype);
}

absl::StatusOr<Tensor> MakeTensorFromPjRtBuffer(
    const DataType dtype, const TensorShape& shape,
    std::unique_ptr<xla::PjRtBuffer> pjrt_buffer) {
  TF_ASSIGN_OR_RETURN(std::unique_ptr<xla::PjRtBuffer::ExternalReference> ref,

  Status status;
  if (reader_ != nullptr) {
    status = reader_->GetTensor(name, out_tensor);
  } else {
    tensorflow::DataType dtype;
    tensorflow::TensorShape shape;
    status = v2_reader_->LookupDtypeAndShape(name, &dtype, &shape);
    if (status.ok()) {
      out_tensor->reset(new Tensor(dtype, shape));
      status = v2_reader_->Lookup(name, out_tensor->get());
      if (!status.ok()) out_tensor->reset();
    }
  }

        xla_shape.dimensions(), &tensor_shape));
    arg_shapes.emplace_back(tensor_shape);

    DataType dtype;
    TF_RETURN_IF_ERROR(ConvertToDataType(input_type, &dtype));

    auto metadata_arg = metadata_proto.add_args();
    metadata_arg->set_kind(tpu::TPUCompileMetadataProto::Arg::PARAMETER);
    metadata_arg->set_dtype(dtype);
  }

  return absl::OkStatus();
}

absl::Status SetupReturnValues(mlir::ModuleOp module,

                                bool unknown_rank, DataType dtype) {
  TensorShapeProto shape;
  shape.set_unknown_rank(unknown_rank);
  for (int64_t dim : dims) {
    shape.add_dim()->set_size(dim);
  }
  mlir::MLIRContext context;
  mlir::Builder b(&context);
  auto status_or = ConvertToMlirTensorType(shape, dtype, &b);
  std::string buf;
  llvm::raw_string_ostream os(buf);
  status_or.value().print(os);

#include "llvm/Support/Debug.h"
#include "llvm/Support/raw_ostream.h"
#include "tensorflow/core/framework/types.pb.h"

// Returns whether the given dtype is a quantization type in TensorFlow.
static bool IsQuantizationType(tensorflow::DataType dtype) {
  switch (dtype) {
    case tensorflow::DT_QINT8:
    case tensorflow::DT_QUINT8:
    case tensorflow::DT_QINT16:
    case tensorflow::DT_QUINT16:
    case tensorflow::DT_QINT32:

  TF_DeleteExecutionContext(ctx);
  return wrap(func);
}

TF_AbstractTensor* TF_AddFunctionParameter(TF_ExecutionContext* func,
                                           TF_DataType dtype, TF_Shape shape,
                                           TF_Status* s) {
  DCHECK_GE(shape.num_dims, -1);
  TracingTensorHandle* t;
  TracingContext* tracing_ctx = dyn_cast<TracingContext>(unwrap(func));

func.func @main() -> (tensor<1x2xf16>, tensor<2xf16>) {
  %graph:2 = tf_executor.graph {
    %0:2 = tf_executor.island wraps "tf.Const"() {device = "", dtype = "tfdtype$DT_HALF", value = dense<1.0> : tensor<1x2xf16>} : () -> tensor<1x2xf16> loc("const1")
    %1:2 = tf_executor.island wraps "tf.Const"() {device = "", dtype = "tfdtype$DT_HALF", value = dense<[1.0, 2.0]> : tensor<2xf16>} : () -> tensor<2xf16> loc("const2")