namespace {

// Helper method that returns an op from 'transpose_ops' that match criteria
// for an 'operand' and 'permutation'
TransposeOp ReuseExistingTranspose(const OpOperand* operand,
                                   const SmallVector<int64_t, 4>& permutation,
                                   Operation* op, ConstOp permutation_op,
                                   SmallVector<TransposeOp, 2>* transpose_ops) {

        transpose_op.getLoc(), transpose_op.getType().getShape(), new_qtype);
    rewriter.setInsertionPointAfter(transpose_op);
    auto new_q_op = rewriter.create<quantfork::QuantizeCastOp>(
        transpose_op.getLoc(), new_tensor_type, q_op.getArg());
    auto new_dq_op = rewriter.create<quantfork::DequantizeCastOp>(
        new_q_op.getLoc(), transpose_op.getResult().getType(),
        new_q_op.getResult());

using ::mlir::stablehlo::TransposeOp;

// Returns `success()` if `op` is a `TransposeOp` with permutation attribute
// equivalent to `permuation`.
LogicalResult IsTransposeOpWithPermuation(absl::Nullable<Operation*> op,
                                          const ArrayRef<int64_t> permutation) {
  auto transpose_op = dyn_cast_or_null<TransposeOp>(op);
  return success(transpose_op != nullptr && transpose_op.getPermutation() ==

  EXPECT_EQ(fused_lstm_func_.getFunctionType().getNumResults(), 1);

  auto transpose_op = fused_lstm_func_.getBody().front().begin();
  transpose_op++;
  EXPECT_EQ(mlir::cast<RankedTensorType>(transpose_op->getOperand(0).getType())
                .getDimSize(0),
            3);
  EXPECT_EQ(mlir::cast<RankedTensorType>(transpose_op->getOperand(0).getType())
                .getDimSize(1),
            12);

  // Transpose.
  auto transpose_op = GetTransposeOpForConv2DFilter(&builder, reshape_op);

  // Reshape Back.
  SmallVector<int64_t, 4> final_shape = {
      height / block_size, width / block_size,
      channel * block_size * block_size, out_channel};
  auto final_reshape_op =
      GetReshapeOpForConv2DFilter(final_shape, transpose_op, &builder);
  // Update filter of Conv2D.

    if (any_blargs) {
      return failure();
    }

    TFL::TransposeOp tpose_arg = nullptr;
    Operation *cst_arg = nullptr;

    for (Value opr : op->getOperands()) {
      auto *defining_op = opr.getDefiningOp();
      if (llvm::isa<TFL::TransposeOp>(defining_op) && !tpose_arg) {
        tpose_arg = llvm::dyn_cast<TFL::TransposeOp>(defining_op);
        continue;
      }

  // Transpose and constant-fold the weight if needed.
  if (transpose_b.getValue()) {
    Value perm = Create1DConstValue<int32_t>(builder, loc, {1, 0});
    auto transpose_op = builder.create<TF::TransposeOp>(loc, weight, perm);
    weight = ConstantFoldOpIfPossible(transpose_op).front();
  }

  xla::DotDimensionNumbers dnums;
  dnums.add_rhs_contracting_dimensions(0);
  if (transpose_a.getValue()) {

      return success();
    }
    return failure();
  }
};

// Fold the constant quantized Transpose ops.
struct FoldTransposeOp : public OpRewritePattern<TransposeOp> {
  explicit FoldTransposeOp(MLIRContext* context)
      : OpRewritePattern<TransposeOp>(context, 1) {}

  // Computes the permutation of a constant `input_tensor` according to `perm`.
  // The function recursively traverses the dimensions of the output tensor in

      auto permutation_and_shape = GetPermutationAndTransposedShape(
          permutation_array, operand_type, rewriter);

      Location loc = scatter_op.getLoc();
      auto transposed_operand = rewriter.create<mhlo::TransposeOp>(
          loc, permutation_and_shape.shape, operands[0],
          permutation_and_shape.permutation);

      Value new_indices = indices;
      int64_t index_depth =

def FoldTrivialReshapeIntoTranspose : Pat<
  (TF_ReshapeOp:$output
      (TF_TransposeOp:$transpose_out1 $input, (Arith_ConstantOp:$permutation1 $p1)), $_),
  (TF_TransposeOp:$transpose_op $input,
      (Arith_ConstantOp
        (RemapPermutation $permutation1,
            (GetPermutationFromTrivialReshape $transpose_out1, $output)))),
  [(IsReshapeEquivalentToTranspose $transpose_out1, $output),