mutex_lock ml(*variable->mu());
  OP_REQUIRES(
      context,
      !variable->is_initialized || variable->tensor()->dtype() == dtype_,
      errors::InvalidArgument(
          "Trying to assign variable with wrong dtype. Expected ",
          DataTypeString(variable->tensor()->dtype()), " got ",
          DataTypeString(dtype_)));
  variable->is_initialized = true;
  *variable->tensor() = value;
}

// Functor that compares the arg values of two different signatures. Returns
// true when the args are not equal.
struct SignatureNotEqual {
  bool operator()(const Tensor& arg, const Tensor& other) {
    return arg.dtype() != other.dtype() || arg.shape() != other.shape() ||
           arg.tensor_data() != other.tensor_data();
  }
  bool operator()(const TensorTypeAndShape& arg,
                  const TensorTypeAndShape& other) {

  *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();
    }
  }

  def.set_device(DEVICE_CPU);

  AttrValue typeAttr;
  SetAttrValue(input_tensor->dtype(), &typeAttr);

  AttrValue outTypeAttr;
  SetAttrValue(out_type, &outTypeAttr);

  (*def.mutable_attr())["T"] = typeAttr;
  (*def.mutable_attr())["type"] = outTypeAttr;

  def.add_input(
      strings::StrCat("input1: ", DataTypeString(input_tensor->dtype())));

  std::unique_ptr<OpKernel> kernel =

        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,

                .values();
        for (const StructuredValue& arg : args) {
          PartialTensorShape shape = arg.tensor_spec_value().shape();
          DataType dtype = arg.tensor_spec_value().dtype();
          TaggedValue tensor_spec(shape, dtype);
          input_signature.tuple().emplace_back(tensor_spec);
        }

        // Setup output signature.
        TensorSpecProto output_tensor_spec_proto =