// %1 = "tfl.transpose"(%0, %cst_0) : (tensor<56x8x56x7xf32>, tensor<4xi32>) -> tensor<8x56x56x7xf32>
// %2 = "tfl.reshape"(%1, %cst_1) : (tensor<8x56x56x7xf32>, tensor<7xi32>) -> tensor<1x1x8x56x56x7x1xf32>
def ConvertTransposeToDecreaseRank : Pat<
  (TFL_TransposeOp:$output_transpose $input, (Arith_ConstantOp:$permutation $_)),
  (TFL_ReshapeOp
    (TFL_TransposeOp
      (TFL_ReshapeOp $input, (Arith_ConstantOp (GetSqueezedShape $input))),

  }
}

}  // namespace

void TransposeOp::getCanonicalizationPatterns(RewritePatternSet& results,
                                              MLIRContext* context) {
  results.add<ConvertTransposeToDecreaseRank>(context);
}

OpFoldResult TransposeOp::fold(FoldAdaptor adaptor) {
  auto operands = adaptor.getOperands();

  auto input_tensor = operands[0].dyn_cast_or_null<DenseElementsAttr>();