// 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()) {

    const TensorType new_input_tensor_type = GetTransposedTensorType(
        mlir::cast<TensorType>(input.getType()), kNchwToNhwcPermutation);

    auto input_transpose_op = rewriter.create<mlir::stablehlo::TransposeOp>(
        op.getLoc(), /*resultType0=*/new_input_tensor_type, /*operand=*/input,
        rewriter.getDenseI64ArrayAttr(kNchwToNhwcPermutation));

    // Transpose the filter tensor: [o, i, 0, 1] => [0, 1, i, o]

class RewriteQuantizedTransposeOp
    : public OpRewritePattern<stablehlo::TransposeOp> {
 public:
  using OpRewritePattern<stablehlo::TransposeOp>::OpRewritePattern;

  LogicalResult match(stablehlo::TransposeOp op) const override {
    return success(IsOpFullyQuantized(op));
  }

  void rewrite(stablehlo::TransposeOp op,
               PatternRewriter& rewriter) const override {

  // has one result operand.
  return isa<TF::IdentityOp, TF::CastOp, TF::ReshapeOp, TF::XlaShardingOp,
             TF::GatherOp, TF::GatherV2Op, TF::XlaGatherOp, TF::ExpandDimsOp,
             TF::SqueezeOp, TF::TransposeOp>(op);
}

bool IsOpWithQuantizableTrait(Operation* op) {
  // Supported quantizable ops.
  return isa<TF::XlaConvV2Op, TF::XlaDotV2Op, TF::MatMulOp, TF::Conv2DOp,

  // CHECK: %[[TRANSPOSE_OP:[0-9]*]] = "tf.Transpose"(%arg0, %[[PERM]])
  // CHECK: %[[PADDING:.*]] = "tf.Const"() <{value = dense<{{\[\[}}0, 0], [0, 0], [1, 1], [1, 1], [0, 0]]> : tensor<5x2xi32>}>
  // CHECK: %[[PAD_OP:[0-9]*]] = "tf.Pad"(%arg0, %[[PADDING]])
  // CHECK: %[[DUP_TRANSPOSE_OP:[0-9]*]] = "tf.Transpose"(%[[PAD_OP]], %[[PERM]])
  // CHECK: return %[[TRANSPOSE_OP]], %[[DUP_TRANSPOSE_OP]]

      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 =

    permutation_array[i] = dim;
    ++dim;
  }

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

  auto transposed_updates = rewriter.create<mhlo::TransposeOp>(
      scatter_op->getLoc(), permutation_and_shape.shape, updates,
      permutation_and_shape.permutation);

  updates = transposed_updates;
  updates_type = permutation_and_shape.shape;
  return success();
}

def ConvertTrivialTransposeOpToReshapeOp : Pat<
  (TFL_TransposeOp:$transpose_op $input, (Arith_ConstantOp:$permutation $p1)),
  (TFL_ReshapeOp $input, (Arith_ConstantOp (GetShape $transpose_op))),
  [(IsTransposeTrivial $input, $permutation),
   (AnyStaticShapeTensor $input),
   (AnyStaticShapeTensor $transpose_op)]>;

// Pattern to fuse redundant tanspose op
def FoldDoubleTranspose : Pat<

  // Converted MHLO ops should be marked illegal here.
  // TODO: b/304003568 - Add TF_TransposeOp folding logic to tflite.
  target.addIllegalOp<mhlo::DotGeneralOp, mhlo::DotOp, mhlo::TransposeOp>();
  if (failed(applyPartialConversion(getOperation(), target,
                                    std::move(patterns)))) {
    getOperation().emitError("mhlo to TFLite legalization failed.");
    signalPassFailure();