#include "tensorflow/c/tf_datatype.h"

#include "tensorflow/core/framework/types.h"

size_t TF_DataTypeSize(TF_DataType dt) {
  return static_cast<size_t>(
      tensorflow::DataTypeSize(static_cast<tensorflow::DataType>(dt)));

#include "tensorflow/core/framework/types.h"
#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) {

static Tensor MakeTensor(DataType dtype, const TensorShape& shape,
                         se::DeviceMemoryBase buffer, Allocator* allocator) {
  size_t expected_size = shape.num_elements() * DataTypeSize(dtype);
  auto* tensor_buffer = new XlaTensorBuffer(buffer.opaque(), expected_size,
                                            buffer.size(), allocator);
  Tensor t(dtype, shape, tensor_buffer);
  tensor_buffer->Unref();

  EXPECT_EQ(TF_HALF, static_cast<TF_DataType>(tensorflow::DT_HALF));
  EXPECT_EQ(TF_DataTypeSize(TF_DOUBLE),
            tensorflow::DataTypeSize(tensorflow::DT_DOUBLE));
  EXPECT_EQ(TF_DataTypeSize(TF_STRING),
            tensorflow::DataTypeSize(tensorflow::DT_STRING));
  // Test with invalid type; should always return 0 as documented
  EXPECT_EQ(TF_DataTypeSize(static_cast<TF_DataType>(0)), 0);
}