RankedTensorType pad_type = pad_op.getType().cast<RankedTensorType>();

    auto transpose_op = pad_input.getDefiningOp<stablehlo::TransposeOp>();
    if (!transpose_op || !transpose_op->hasOneUse()) return failure();
    Value transpose_input = transpose_op.getOperand();

    ArrayRef<int64_t> transpose_perm = transpose_op.getPermutation();
    SmallVector<int64_t> new_padding_low =

    ),
    $bmm_rhs, ConstBoolAttrTrue, $transpose_rhs,
    $asymmetric_quantize_inputs_bmm
  ),
  (TFL_BatchMatMulOp
    (TFL_BatchMatMulOp $fc_rhs, $fc_lhs, ConstBoolAttrFalse, ConstBoolAttrFalse,
                       $asymmetric_quantize_inputs_fc,
                       (returnType (TransposeFCLastTwoDims $fc_output))
    ),
    $bmm_rhs, ConstBoolAttrFalse, $transpose_rhs,
    $asymmetric_quantize_inputs_bmm
  ),

  SmallVector<int64_t> permutation_;
};

class FoldTransposedConstantOp
    : public OpRewritePattern<mlir::stablehlo::TransposeOp> {
 public:
  using OpRewritePattern<mlir::stablehlo::TransposeOp>::OpRewritePattern;

  LogicalResult match(mlir::stablehlo::TransposeOp op) const override {
    Value operand = op.getOperand();
    auto const_op =

    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]

  // 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]]

    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();
}

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

TAC_REGISTER_GPU_OP(SquaredDifferenceOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(StridedSliceOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(TanhOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(TransposeOp, CreateBasicOpNoCost);
TAC_REGISTER_GPU_OP(TransposeConvOp, CreateBasicOpNoCost);

TAC_REGISTER_GPU_OP(ConcatenationOp, CreateConcatOp);
TAC_REGISTER_GPU_OP(ReshapeOp, CreateConcatOp);

                mlir::stablehlo::PadOp, mlir::stablehlo::ReduceWindowOp,
                mlir::stablehlo::ReshapeOp, mlir::stablehlo::SelectOp,
                mlir::stablehlo::SliceOp, mlir::stablehlo::TransposeOp>(op)) {
    scale_spec->has_same_scale_requirement = true;
  }
  if (llvm::isa<mlir::stablehlo::DynamicSliceOp, mlir::stablehlo::GatherOp,
                mlir::stablehlo::PadOp, mlir::stablehlo::SliceOp>(op)) {