public:
  // Constructs a scalar, DT_FLOAT TensorSpec
  TensorSpec();

  TensorSpec(PartialTensorShape shape, DataType dtype);

  explicit TensorSpec(const TensorSpecProto& proto);

  const PartialTensorShape& shape() const;
  DataType dtype() const;

 private:
  PartialTensorShape shape_;
  DataType dtype_;
};

}  // namespace tensorflow

#include "tensorflow/core/protobuf/saved_object_graph.pb.h"

namespace tensorflow {
namespace {

class SavedVariableLoadingTest
    : public ::testing::TestWithParam<
          std::tuple<DataType, std::vector<int64_t>>> {
 public:
  SavedVariableLoadingTest() {
    SessionOptions options;
    options.config.mutable_device_count()->insert({"CPU", 3});
    std::vector<std::unique_ptr<Device>> devices;

        val error = result.errors.singleOrNull()
        assertNotNull(error)
        val reason = error.errorReason
        assertIs<ErrorReason.AssignmentTypeMismatch>(reason)
        assertIs<DataType.StringDataType>(reason.expected)
        assertIs<DataType.IntDataType>(reason.actual)
    }

    data class ValReassignment(val localVal: LocalValue) : ErrorReason
    data class ExternalReassignment(val external: ObjectOrigin.External) : ErrorReason
    data class AssignmentTypeMismatch(val expected: DataType, val actual: DataType) : ErrorReason

    // TODO: these two are never reported for now, instead it is UnresolvedFunctionCallSignature
    data object UnusedConfigureLambda : ErrorReason
    data object MissingConfigureLambda : ErrorReason

                                           DataType dtype, TensorShape shape,
                                           const char* raw_device_name,
                                           ImmediateTensorHandlePtr* handle);

// Executes an AssignVariableOp using `ctx`, assigning the variable associated
// with `variable_handle` with `value`. `dtype` must be the datatype of the

// Demangle a string mangled with MangleTensor.
Status DemangleTensor(absl::string_view str, TensorProto* proto);

// Return a DataType mangled as a string.
string MangleDataType(const DataType& dtype);
// Demangle a string mangled with MangleDataType.
Status DemangleDataType(absl::string_view str, DataType* proto);

}  // namespace mangling_util
}  // namespace tensorflow

#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) {

  TFE_TensorHandle* input_eager =
      TestTensorHandleWithDims<T, datatype>(eager_ctx, data, dims, num_dims);
  *tensor =
      unwrap(TF_CreateAbstractTensorFromEagerTensor(input_eager, status.get()));
  return absl::OkStatus();
}

// Return a scalar tensor handle with given value.
template <class T, TF_DataType datatype>
Status TestScalarTensorHandle(AbstractContext* ctx, const T value,

TensorSpec::TensorSpec(PartialTensorShape shape, DataType dtype)
    : shape_(std::move(shape)), dtype_(dtype) {}

TensorSpec::TensorSpec(const TensorSpecProto& proto)
    : shape_(proto.shape()), dtype_(proto.dtype()) {}

const PartialTensorShape& TensorSpec::shape() const { return shape_; }

DataType TensorSpec::dtype() const { return dtype_; }

    return mlir::cast<mlir::StringAttr>(index_path[0]).getValue();
  }
  return "";
}

absl::StatusOr<std::pair<tensorflow::DataType, tensorflow::PartialTensorShape>>
ProcessTensorSpec(mlir::TensorType type) {
  tensorflow::DataType dtype;
  TF_RETURN_IF_ERROR(
      ConvertScalarTypeToDataType(type.getElementType(), &dtype));

  if (!type.hasRank())