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,

      OpT op, typename OpT::Adaptor adaptor,
      ConversionPatternRewriter &rewriter) const override {
    Type dtype = op.getElementDtype();
    if (!(dtype.isF16() || dtype.isF32() || dtype.isF64() ||
          dtype.isInteger(1) || dtype.isInteger(8) || dtype.isInteger(16) ||
          dtype.isInteger(32) || dtype.isInteger(64))) {
      const char *error_info =
          "requires element_dtype to be 1-bit/8-bit/16-bit/32-bit/64-bit "

    for (auto type : inst.getResultTypes())
      if (!mlir::isa<mlir::TensorType, mlir::tf_executor::ControlType,
                     mlir::tf_executor::TokenType>(type))
        return errors::InvalidArgument(
            "Values must be of tensor type, TensorFlow control type, or "
            "TensorFlow token type. Found ",
            mlir::debugString(type));

string AttrView::DefaultValue() const {
  const AttrValue &attr_value = attr_.default_value();
  switch (attr_value.value_case()) {
    case AttrValue::VALUE_NOT_SET:
      return "";
    case AttrValue::kType:
      return DataType_Name(attr_value.type());
    case AttrValue::kS:
      return "\"" + attr_value.s() + "\"";
    case AttrValue::kI:
      return std::to_string(attr_value.i());
    case AttrValue::kF:

    }
    if (ftype.isF64()) {
      return tflite::TensorType_COMPLEX128;
    }
    return Status(absl::StatusCode::kInvalidArgument, "Unsupported type");
  } else if (auto itype = mlir::dyn_cast<mlir::IntegerType>(type)) {
    switch (itype.getWidth()) {
      case 1:
        return tflite::TensorType_BOOL;
      case 4:
        if (itype.isUnsigned()) {
          return Status(absl::StatusCode::kInvalidArgument,

  else
    return UnrankedTensorType::get(type.getValue());
}
}  // namespace

// Returns whether type can be further refined.
bool CanBeRefined(Type type) {
  auto shape_type = mlir::dyn_cast<ShapedType>(type);
  if (!shape_type) return false;

  // Returns whether type with subtypes can be further refined.

    rewriter.replaceOpWithNewOp<mhlo::ReturnOp>(op, adaptor.getOperands());
    return success();
  }
};

// Returns a new tensor type from the given type with element type updated to
// the given type.
TensorType UpdateElementTypeTo(Type ty, Type element_ty) {
  auto ranked_ty = mlir::dyn_cast<RankedTensorType>(ty);
  if (!ranked_ty) {
    return UnrankedTensorType::get(element_ty);
  }