return mlir::TF::ShapeAttr::get(type.getContext(), tensor_type.getShape());
  }

  // If type is not a RankedTensor or UnrankedTensor, it must be a scalar.
  // Empty TensorShape indicates a scalar.
  return mlir::TF::ShapeAttr::get(type.getContext(), ArrayRef<int64_t>());
}

absl::StatusOr<TensorSpecProto> ConvertTypeToTensorSpecProto(
    const mlir::Type& type) {
  DataType dtype;

                                     const char* attr_name,
                                     const TF_DataType type,
                                     TF_Status* status) {
  tensorflow::DataType dtype = static_cast<tensorflow::DataType>(type);
  tensorflow::AddTypeConstraint(kernel_builder, attr_name, dtype, status);
}

void TF_KernelBuilder_HostMemory(TF_KernelBuilder* kernel_builder,

        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,

//     inputs: {
//       key  : "a"
//       value: {
//         name : "serving_default_a:0"
//         dtype: DT_FLOAT
//         tensor_shape: {
//         }
//       }
//     }
//     inputs: {
//       key  : "b"
//       value: {
//         name : "serving_default_b:0"
//         dtype: DT_FLOAT
//         tensor_shape: {
//         }
//       }
//     }
//     outputs: {

  }

  tensorflow::DataType dtype;
  auto status = tensorflow::ConvertToDataType(operand.getType(), &dtype);
  if (!status.ok()) {
    op->emitRemark() << "skipping legalization due to " << status.ToString();
    return tensorflow::XlaExpression::Invalid();
  }
  return tensorflow::XlaExpression::XlaOp(xla_op, dtype);
}

}  // namespace mhlo

    }
    if (TFE_TensorHandleDataType(component.get()) != dtype) {
      TF_SetStatus(status, TF_INTERNAL,
                   "Components of a ParallelTensor must all have "
                   "the same dtype");
      return nullptr;
    }
  }
  return std::unique_ptr<ParallelTensor>(
      new ParallelTensor(parallel_device, std::move(components), shape, dtype));
}

    AllocatorAttributes attr;
    if (cc_input.dtype() == ::tensorflow::DT_VARIANT) {
      attr.set_on_host(true);
    }
    TF_RETURN_IF_ERROR(cc_ctx->allocate_temp(cc_input.dtype(), cc_input.shape(),
                                             cc_out, attr));

    switch (cc_input.dtype()) {
      case ::tensorflow::DT_INVALID: {
        *cc_out = Tensor(::tensorflow::DT_INVALID);

#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:

    }

    se::Stream* const src_compute_stream = src_xla_context->stream();
    TF_RET_CHECK(src_compute_stream != nullptr);
    TF_RET_CHECK(input->dtype() == output->dtype())
        << "input type: " << DataTypeString(input->dtype()) << " output type "
        << DataTypeString(output->dtype());
    TF_RET_CHECK(input->shape() == output->shape());
    TF_RET_CHECK(DMAHelper::CanUseDMA(input));
    auto* const src_compute_stream_impl =

  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));