// Given a Composite op that wraps a core.aten.avg_pool2d, and assuming that
// the padding part is extracted into a tfl.pad op prior to a
// tfl.average_pool_2d, this function finds the return type of the needed
// tfl.pad .
ShapedType GetPadOpType(mhlo::CompositeOp op);

// Given a Composite op that wraps a core.aten.avg_pool2d, finds the padding
// attribute to be passed to the a tfl.average_pool_2d that can fully replace

    if (!IsQI32Type(input_dequant.getType())) return failure();

    auto output_type =
        mlir::dyn_cast_or_null<ShapedType>(dequant_op.getOutput().getType());
    if (!output_type || !output_type.getElementType().isF32()) return failure();

    auto input_type = mlir::dyn_cast<ShapedType>(input_dequant.getType());
    // TODO(renjieliu): support UniformQuantizedPerAxisType.

// Convert mhlo.dot to mhlo.dot_general.
LogicalResult ConvertDotToDotGeneral(mhlo::DotOp op,
                                     PatternRewriter &rewriter) {
  auto lhs_type = mlir::cast<ShapedType>(op.getLhs().getType());
  auto rhs_type = mlir::cast<ShapedType>(op.getRhs().getType());
  if (!lhs_type.hasRank() || !rhs_type.hasRank()) {
    return rewriter.notifyMatchFailure(op, "unsupported unranked input type");
  }

      CPred<"$_op.getOperand(" # n #
      ").getType().cast<ShapedType>().getRank() == " # m>]>>;

// Returns true if the n-th result has unknown rank or has rank m.
class TF_ResultHasRank<int n, int m> :
  PredOpTrait<"result " # n # " is " # m # "-D",
    Or<[TF_ResultIsUnrankedPred<n>,
      CPred<"$_op.getResult(" # n #
      ").getType().cast<ShapedType>().getRank() == " # m>]>>;

    auto loc = op.getLoc();
    auto result_ty = mlir::cast<ShapedType>(op.getType());

    auto input = op.getImages();
    auto input_ty = mlir::cast<ShapedType>(input.getType());
    auto input_element_ty = input_ty.getElementType();
    auto out_size = op.getSize();
    auto out_size_ty = mlir::cast<ShapedType>(out_size.getType());
    auto out_size_element_ty = out_size_ty.getElementType();

          CPred<"$_op.getOperand(0).getType().cast<ShapedType>().getRank() <= 1">,
          CPred<"$_op.getOperand(0).getType().cast<ShapedType>().getRank() == 2 && !$_op.getOperand(0).getType().cast<ShapedType>().hasStaticShape()">,
          CPred<"$_op.getOperand(0).getType().cast<ShapedType>().getRank() == 2 && $_op.getOperand(0).getType().cast<ShapedType>().getShape()[1] <= 4">]>>]> {

          .isExactlyValue(0.0f));

  EXPECT_EQ(fused_lstm_func_.getFunctionType().getNumResults(), 1);
  auto output_types = fused_lstm_func_.getFunctionType().getResults();
  SmallVector<int64_t, 2> output_shape{1, mlir::ShapedType::kDynamic};
  EXPECT_EQ(mlir::cast<RankedTensorType>(output_types[0]).getShape().size(),
            output_shape.size());
  for (int i = 0; i < output_shape.size(); i++) {

      return failure();
    }

    // Remove identity reshape with both static result and input shape.
    auto result_type = mlir::cast<ShapedType>(op.getType());
    auto input_type = mlir::cast<ShapedType>(op.getInput().getType());

    // Constant folding
    // If the result type isn't static, tries to derive the result type from
    // the #2 operand.

// not.
const char kLegalizeTflVariables[] = "tfl._legalize_tfl_variables";

bool HasSupportedElementType(Operation* op) {
  return utils::IsSupportedVariableType(op);
}

bool IsSupportedElementType(ShapedType type) {
  return utils::IsSupportedVariableType(type);
}

#include "tensorflow/compiler/mlir/lite/transforms/generated_legalize_variables.inc"

// Pass which legalizes TF variables which are already passed as bounded

    const xla::ConvolutionDimensionNumbers &dnums, PatternRewriter &rewriter) {
  StringAttr conv_padding = op.getPaddingAttr();
  SmallVector<int64_t> padding_nums;
  ShapedType lhs_shape = mlir::cast<ShapedType>(op.getLhs().getType());
  ShapedType rhs_shape = mlir::cast<ShapedType>(op.getRhs().getType());

  // Handle only static shape cases.
  // TODO(b/260284866): Handle dynamic shape cases.
  if (!lhs_shape.hasStaticShape()) {