// Returns true if 'op' has type that is supported by native TFLite
// variables.
bool IsSupportedVariableType(Operation* op);

// Returns true if 'type' is supported by native tflite variables.
bool IsSupportedVariableType(ShapedType type);

}  // namespace utils
}  // namespace TFL
}  // namespace mlir

  }

  // Here we only add the shapes for the leading values with ShapedType,
  // assuming values with non-ShapedType are put at the end of the result.
  if (!ignore_unregistered_attrs && inst->getNumResults() > 0) {
    auto values = inst->getResults();
    auto begin = values.begin();
    auto end = values.begin();
    while (end != values.end() && mlir::isa<mlir::ShapedType>((*end).getType()))
      end++;
    if (begin != end) {

  CPred<"$0.getType().isa<RankedTensorType>()">,
  CPred<"!$0.getType().cast<ShapedType>().isDynamicDim( "
  "  $0.getType().cast<RankedTensorType>().getRank() - 1)">]>>;

def IsNotComplexType : Constraint<And<[
  CPred<"$0.getType().isa<RankedTensorType>()">,
  CPred<"!$0.getType().cast<ShapedType>().getElementType().isa<ComplexType>()">
]>>;

    // we have separate ops for them. If only one of them is used then the other
    // one will be garbage collected later.
    if (!mlir::isa<ShapedType>(operand.getType())) return failure();
    auto operand_type = mlir::cast<ShapedType>(operand.getType());
    if (operand_type.getElementType().isInteger(1)) {
      // TF does not support min or max on boolean (int1) arguments.

bool IsLargeFloatType(Type type) {
  type = getElementTypeOrSelf(type);
  return isa<FloatType>(type) && type.getIntOrFloatBitWidth() > 16;
}

Type ToBfloat16Type(Type type) {
  if (auto shaped = mlir::dyn_cast<ShapedType>(type)) {
    const Type elem = shaped.getElementType();
    if (IsLargeFloatType(elem)) {
      return shaped.clone(BFloat16Type::get(type.getContext()));
    }
  } else if (IsLargeFloatType(type)) {

      }
    } else if (function_name.contains("batch_matmul")) {
      // For BatchMatMul, the input must be ranked to determine the batch
      // dimensions.
      ShapedType shaped_type =
          mlir::dyn_cast<ShapedType>(call_op->getOperand(0).getType());
      if (!shaped_type || !shaped_type.hasRank()) {
        return absl::InternalError("The input of BatchMatMul must have rank.");
      }

  std::string device_type_;
  bool prefer_tf2xla_;
  bool use_tf2xla_hlo_importer_;
};

bool ShouldRefineTypeTo(Type original_ty, Type updated_ty) {
  auto updated = mlir::dyn_cast<ShapedType>(updated_ty);
  auto original = mlir::dyn_cast<ShapedType>(original_ty);

  // Both types must be shaped types.
  if (!original || !updated) return false;

  // Element types must match.

    auto transpose_op = *transpose_ops.begin();
    auto result_type =
        mlir::dyn_cast_or_null<ShapedType>(transpose_op.getResult().getType());
    auto is_valid_move =
        llvm::all_of(op->getOperands(), [result_type](Value operand) -> bool {
          auto operand_type =
              mlir::dyn_cast_or_null<ShapedType>(operand.getType());
          return result_type && operand_type && result_type.hasRank() &&

namespace tensorflow {

using llvm::ArrayRef;
using llvm::SmallVector;
using mlir::Builder;
using mlir::DenseStringElementsAttr;
using mlir::ElementsAttr;
using mlir::RankedTensorType;
using mlir::ShapedType;
using mlir::Type;
using tensorflow::errors::InvalidArgument;

static TensorProto ConvertToProto(const Tensor& input_tensor,
                                  bool use_tensor_content = true) {

      if (!cast_op || cast_op.getResult().use_empty()) continue;

      // Get types
      Type old_result_type = op.getResult().getType();
      ShapedType new_result_type =
          mlir::dyn_cast<ShapedType>(cast_op.getType());

      // Proceeds only if the casting is to float16
      if (!new_result_type.getElementType().isF16()) continue;

      // Cast values